def tf_scatter(self, v):
"""
Scatter the given direction tensor v.
Parameters
----------
v : TF tensor, shape(?, 3)
Direction vectors of the photons which are being scattered.
Returns
-------
The scattered direction tensor of shape(?, 3).
"""
# sample cos(theta)
cosTs = 2 * self._uni_pdf.sample(tf.shape(v)[0])**(1 / 19) - 1
cosT2s = tf.sqrt((cosTs + 1) / 2)
sinT2s = tf.sqrt((1 - cosTs) / 2)
ns = tf.transpose(
self.tf_sample_normal_vectors(v) * tf.expand_dims(sinT2s, axis=-1))
# ignore the fact that n could be parallel to v, what's the probability
# of that happening?
q = tfq.Quaternion(tf.transpose([cosT2s, ns[0], ns[1], ns[2]]))
return tfq.rotate_vector_by_quaternion(q, v)
def test_rotate_vector_by_quaternion(self):
v = [3., 4., 5.]
qs_np, qs_tf = get_quaternions()
with self.test_session():
for q in qs_tf[QS_FINITENONZERO]:
rotated = q * tfq.vector3d_to_quaternion(v) * q.inverse()
rotated = tfq.quaternion_to_vector3d(rotated)
self.assertAllClose(tfq.rotate_vector_by_quaternion(q, v),
rotated)
def augment(points, segment_ids, training_batch_size):
"""Rotate a full scene around the z axis and scale on each axis."""
scene_ids = segment_ids // 2
bs = training_batch_size * 3
rand_angle = tf.random_uniform([bs], maxval=(2 * pi))
rot = tf.stack([[1.] * bs, [0.] * bs, [0.] * bs, rand_angle], axis=1)
rotations = tfq.Quaternion(tf.gather(rot, scene_ids))
points = tfq.rotate_vector_by_quaternion(rotations, points, 2, 2)
return points
def _input_fn(self, obj_ids, translations, rotations, train_batch_size,
num_objs, do_augmentation):
"""The input function 's part that is shared.
This function creates the scene point clouds from scene descriptions.
Returns: Two tf.Tensors, the first contains all points of the
objects in the batch with shape (N, 3) and the second contains the
corresponding segment ids, the shape is (N,).
"""
batch_size = tf.shape(obj_ids)[0]
# flatten all inputs
obj_ids = tf.reshape(obj_ids, (-1, ))
translations = tf.reshape(translations, (-1, 3))
rotations = tf.reshape(rotations, (-1, 4))
clouds_num_points = (self.cloud_slice_indices[1:] -
self.cloud_slice_indices[:-1])
# vector with the number of points of each cloud
num_points = tf.gather(clouds_num_points, obj_ids)
# vector with a range where each number i is num_points[i] repeated
segment_ids = self._repeat(tf.range(tf.shape(num_points)[0]),
tf.to_int32(num_points), batch_size,
num_objs)
segment_ids = tf.to_int32(segment_ids)
# repeat translations[i] and rotations[i] num_points[i] times
translations = tf.gather(translations, segment_ids)
rotations = tf.gather(rotations, segment_ids)
rotations = tfq.Quaternion(rotations)
obj_ids = tf.gather(tf.to_float(obj_ids), segment_ids)
# indices of points consist of the start index plus range(num_points)
start = tf.gather(self.cloud_slice_indices, tf.to_int32(obj_ids))
ranges = tf.cond(
tf.equal(batch_size, 1),
lambda: tf.concat([tf.range(num_points[i]) for i in range(2)],
axis=0), lambda: tf.
concat([tf.range(num_points[i]) for i in range(num_objs)], axis=0))
point_ids = tf.to_int32(start + ranges)
points = tf.gather(self.clouds_tensor, point_ids)
# Rotate objects. Note that the quaternions are relative to the object
# clouds' origins, so no centering using the mean is required.
points = tfq.rotate_vector_by_quaternion(rotations, points, 2, 2)
points = tf.squeeze(points) + translations
# if we're training, randomly rotate around the z_axis
if do_augmentation:
points = augment.pointcloud(points, segment_ids, batch_size,
train_batch_size)
return points, tf.to_float(segment_ids)
def z_rotate():
global central_angle
central_angle = tfq.rotate_vector_by_quaternion(z_rotation, central_angle)
redraw()