def test_inverse_random(self):
"""Checks that inverting rotated points results in no transformation."""
random_euler_angle = test_helpers.generate_random_test_euler_angles()
tensor_tile = random_euler_angle.shape[:-1]
random_matrix = rotation_matrix_3d.from_euler(random_euler_angle)
random_point = np.random.normal(size=tensor_tile + (3,))
rotated_random_points = rotation_matrix_3d.rotate(random_point,
random_matrix)
predicted_invert_random_matrix = rotation_matrix_3d.inverse(random_matrix)
predicted_invert_rotated_random_points = rotation_matrix_3d.rotate(
rotated_random_points, predicted_invert_random_matrix)
self.assertAllClose(
random_point, predicted_invert_rotated_random_points, rtol=1e-6)
def test_rotate_vs_rotate_quaternion_random(self):
"""Tests that the rotate provide the same results as quaternion.rotate."""
random_euler_angle = test_helpers.generate_random_test_euler_angles()
tensor_tile = random_euler_angle.shape[:-1]
random_matrix = rotation_matrix_3d.from_euler(random_euler_angle)
random_quaternion = quaternion.from_rotation_matrix(random_matrix)
random_point = np.random.normal(size=tensor_tile + (3, ))
ground_truth = quaternion.rotate(random_point, random_quaternion)
prediction = rotation_matrix_3d.rotate(random_point, random_matrix)
self.assertAllClose(ground_truth, prediction, rtol=1e-6)
def test_from_euler_random(self):
"""Tests that quaternions can be constructed from Euler angles."""
random_euler_angles = test_helpers.generate_random_test_euler_angles()
tensor_shape = random_euler_angles.shape[:-1]
random_matrix = rotation_matrix_3d.from_euler(random_euler_angles)
random_quaternion = quaternion.from_euler(random_euler_angles)
random_point = np.random.normal(size=tensor_shape + (3,))
rotated_with_matrix = rotation_matrix_3d.rotate(random_point, random_matrix)
rotated_with_quaternion = quaternion.rotate(random_point, random_quaternion)
self.assertAllClose(rotated_with_matrix, rotated_with_quaternion)
def test_rotate_jacobian_random(self):
"""Test the Jacobian of the rotate function."""
x_matrix_init = test_helpers.generate_random_test_rotation_matrix_3d()
x_matrix = tf.convert_to_tensor(value=x_matrix_init)
tensor_shape = x_matrix.shape[:-1]
x_point_init = np.random.uniform(size=tensor_shape)
x_point = tf.convert_to_tensor(value=x_point_init)
y = rotation_matrix_3d.rotate(x_point, x_matrix)
self.assert_jacobian_is_correct(x_matrix, x_matrix_init, y)
self.assert_jacobian_is_correct(x_point, x_point_init, y)
def test_from_euler_random(self):
"""Tests that from_euler allows to perform the expect rotation of points."""
random_euler_angles = test_helpers.generate_random_test_euler_angles()
tensor_shape = random_euler_angles.shape[:-1]
random_point = np.random.normal(size=tensor_shape + (3,))
random_matrix = rotation_matrix_3d.from_euler(random_euler_angles)
random_axis, random_angle = axis_angle.from_euler(random_euler_angles)
rotated_with_matrix = rotation_matrix_3d.rotate(random_point, random_matrix)
rotated_with_axis_angle = axis_angle.rotate(random_point, random_axis,
random_angle)
self.assertAllClose(rotated_with_matrix, rotated_with_axis_angle)
def test_rotate_random(self):
"""Tests the rotation using a quaternion vs a rotation matrix."""
random_quaternion = test_helpers.generate_random_test_quaternions()
tensor_shape = random_quaternion.shape[:-1]
random_point = np.random.normal(size=tensor_shape + (3,))
rotated_point_quaternion = quaternion.rotate(random_point,
random_quaternion)
matrix = rotation_matrix_3d.from_quaternion(random_quaternion)
rotated_point_matrix = rotation_matrix_3d.rotate(random_point, matrix)
self.assertAllClose(
rotated_point_matrix, rotated_point_quaternion, rtol=1e-3)
def generate_ground_image(height,
width,
focal,
principal_point,
camera_rotation_matrix,
camera_translation_vector,
ground_color=(0.43, 0.43, 0.8)):
"""Generate an image depicting only the ground."""
batch_size = camera_rotation_matrix.shape[0]
background_image = np.ones((batch_size, height, width, 1, 1),
dtype=np.float32)
background_image[:, -1, ...] = 0 # Zero the bottom line for proper sampling
# The projection of the ground depends on the top right corner (approximation)
plane_point_np = np.tile(np.array([[3.077984, 2.905388, 0.]],
dtype=np.float32), (batch_size, 1))
plane_point_rotated = rotation_matrix_3d.rotate(plane_point_np,
camera_rotation_matrix)
plane_point_translated = plane_point_rotated + camera_translation_vector
plane_point2d = \
perspective.project(plane_point_translated, focal, principal_point)
_, y = tf.split(plane_point2d, [1, 1], axis=-1)
sfactor = height/y
helper_matrix1 = np.tile(np.array([[[1, 0, 0],
[0, 0, 0],
[0, 0, 0]]]), (batch_size, 1, 1))
helper_matrix2 = np.tile(np.array([[[0, 0, 0],
[0, 1, 0],
[0, 0, 1]]]), (batch_size, 1, 1))
transformation_matrix = tf.multiply(tf.expand_dims(sfactor, -1),
helper_matrix1) + helper_matrix2
plane_points = grid.generate((0., 0., 0.),
(float(height), float(width), 0.),
(height, width, 1))
plane_points = tf.reshape(plane_points, [-1, 3])
transf_plane_points = tf.matmul(transformation_matrix,
plane_points,
transpose_b=True)
interpolated_points = \
trilinear.interpolate(background_image,
tf.linalg.matrix_transpose(transf_plane_points))
ground_alpha = (1- tf.reshape(interpolated_points,
[batch_size, height, width, 1]))
ground_image = tf.ones((batch_size, height, width, 3))*ground_color
return ground_image, ground_alpha
def generate_ground_image(height,
width,
focal,
principal_point,
camera_rotation_matrix,
camera_translation_vector,
ground_color=(0.43, 0.43, 0.8)):
"""Generate an image depicting only the ground."""
background_image = np.ones((height, width, 1, 1), dtype=np.float32)
background_image[-1,
...] = 0 # Set the bottom line to 0 for proper sampling
# The projection of the ground depends on the top right corner (approximation)
plane_point_np = np.array([[3.077984, 2.905388, 0.]], dtype=np.float32)
plane_point2d = \
perspective.project(rotation_matrix_3d.rotate(plane_point_np,
camera_rotation_matrix) +
camera_translation_vector.T, focal, principal_point)
_, y = tf.split(plane_point2d, [1, 1], axis=-1)
sfactor = y / 256.
transformation_matrix = (1/sfactor)*np.array([[1, 0, 0],
[0, 0, 0],
[0, 0, 0]]) + \
np.array([[0, 0, 0],
[0, 1, 0],
[0, 0, 1]])
plane_points = grid.generate(
(0., 0., 0.), (float(height), float(width), 0.), (height, width, 1))
plane_points = tf.reshape(plane_points, [-1, 3])
transf_plane_points = tf.matmul(transformation_matrix,
plane_points,
transpose_b=True)
interpolated_points = \
trilinear.interpolate(background_image, tf.transpose(transf_plane_points))
ground_alpha = (1 - tf.reshape(interpolated_points, [256, 256, 1]))
ground_image = tf.ones((256, 256, 3)) * ground_color
return ground_image, ground_alpha
def __getitem__(self, idx: int):
# Fetch data for this batch
batch_slice = slice(idx * self.batch_size,
min((idx + 1) * self.batch_size, self.len))
src = self.data[batch_slice]
batch_len = len(src)
# Compute true transformation
euler_angles = np.random.uniform(low=-np.pi,
high=np.pi,
size=(batch_len, 1,
3)).astype(np.float32)
R_true = rotation_matrix_3d.from_euler(euler_angles)
t_true = np.random.uniform(low=-0.5, high=0.5,
size=(batch_len, 1, 3)).astype(np.float32)
# Output batch
tgt = rotation_matrix_3d.rotate(src, R_true) + t_true
src = tf.expand_dims(src, axis=1)
tgt = tf.expand_dims(tgt, axis=1)
x = tf.concat([src, tgt], 1)
y_true = tf.concat([tf.squeeze(R_true), t_true], 1)
return x, y_true
def blend(points,
skinning_weights,
bone_rotations,
bone_translations,
name=None):
"""Transforms the points using Linear Blend Skinning.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible and allow transforming full 3D shapes at once.
In the following, B1 to Bm are optional batch dimensions, which allow
transforming multiple poses at once.
Args:
points: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents a 3d point.
skinning_weights: A tensor of shape `[A1, ..., An, W]`, where the last
dimension represents the skinning weights of each bone.
bone_rotations: A tensor of shape `[B1, ..., Bm, W, 3, 3]`, which represents
the 3d rotations applied to each bone.
bone_translations: A tensor of shape `[B1, ..., Bm, W, 3]`, which represents
the 3d translation vectors applied to each bone.
name: A name for this op that defaults to "linear_blend_skinning_blend".
Returns:
A tensor of shape `[B1, ..., Bm, A1, ..., An, 3]`, where the last dimension
represents a 3d point.
Raises:
ValueError: If the shape of the input tensors are not supported.
"""
with tf.compat.v1.name_scope(
name, "linear_blend_skinning_blend",
[points, skinning_weights, bone_rotations, bone_translations]):
points = tf.convert_to_tensor(value=points)
skinning_weights = tf.convert_to_tensor(value=skinning_weights)
bone_rotations = tf.convert_to_tensor(value=bone_rotations)
bone_translations = tf.convert_to_tensor(value=bone_translations)
shape.check_static(
tensor=points, tensor_name="points", has_dim_equals=(-1, 3))
shape.check_static(
tensor=bone_rotations,
tensor_name="bone_rotations",
has_rank_greater_than=2,
has_dim_equals=((-2, 3), (-1, 3)))
shape.check_static(
tensor=bone_translations,
tensor_name="bone_translations",
has_rank_greater_than=1,
has_dim_equals=(-1, 3))
shape.compare_dimensions(
tensors=(skinning_weights, bone_rotations),
tensor_names=("skinning_weights", "bone_rotations"),
axes=(-1, -3))
shape.compare_dimensions(
tensors=(skinning_weights, bone_translations),
tensor_names=("skinning_weights", "bone_translations"),
axes=(-1, -2))
shape.compare_batch_dimensions(
tensors=(points, skinning_weights),
tensor_names=("points", "skinning_weights"),
last_axes=(-2, -2),
broadcast_compatible=True)
shape.compare_batch_dimensions(
tensors=(bone_rotations, bone_translations),
tensor_names=("bone_rotations", "bone_translations"),
last_axes=(-3, -2),
broadcast_compatible=True)
num_bones = skinning_weights.shape[-1]
def dim_value(dim):
return 1 if dim is None else tf.compat.v1.dimension_value(dim)
# TODO(b/148362025): factorize this block out
points_batch_shape = shape.get_broadcasted_shape(
points.shape[:-1], skinning_weights.shape[:-1])
points_batch_shape = [dim_value(dim) for dim in points_batch_shape]
points = tf.broadcast_to(points, points_batch_shape + [3])
skinning_weights = tf.broadcast_to(skinning_weights,
points_batch_shape + [num_bones])
bones_batch_shape = shape.get_broadcasted_shape(
bone_rotations.shape[:-3], bone_translations.shape[:-2])
bones_batch_shape = [dim_value(dim) for dim in bones_batch_shape]
bone_rotations = tf.broadcast_to(bone_rotations,
bones_batch_shape + [num_bones, 3, 3])
bone_translations = tf.broadcast_to(bone_translations,
bones_batch_shape + [num_bones, 3])
points_batch_dims = points.shape.ndims - 1
bones_batch_dims = bone_rotations.shape.ndims - 3
points = tf.reshape(points,
[1] * bones_batch_dims + points_batch_shape + [1, 3])
skinning_weights = tf.reshape(skinning_weights, [1] * bones_batch_dims +
points_batch_shape + [num_bones, 1])
bone_rotations = tf.reshape(
bone_rotations,
bones_batch_shape + [1] * points_batch_dims + [num_bones, 3, 3])
bone_translations = tf.reshape(
bone_translations,
bones_batch_shape + [1] * points_batch_dims + [num_bones, 3])
transformed_points = rotation_matrix_3d.rotate(
points, bone_rotations) + bone_translations
weighted_points = tf.multiply(skinning_weights, transformed_points)
blended_points = tf.reduce_sum(input_tensor=weighted_points, axis=-2)
return blended_points
def func(angles, point):
matrix = rotation_matrix_3d.from_euler(angles)
return rotation_matrix_3d.rotate(point, matrix)
def func(axis, angle, point):
matrix = rotation_matrix_3d.from_axis_angle(axis, angle)
return rotation_matrix_3d.rotate(point, matrix)
def rotate(points):
"""Randomly rotates a point cloud around the Y axis (UP)."""
axis = tf.constant([[0, 1, 0]], dtype=tf.float32) # [1, 3]
angle = tf.random.uniform((1, 1), minval=0., maxval=2 * np.pi) # [1, 1]
matrix = rotation_matrix_3d.from_axis_angle(axis, angle) # [3, 3]
return rotation_matrix_3d.rotate(points, matrix) # [N, 3]
