def test_assert_rotation_matrix_normalized_preset(self, dtype): """Checks that assert_normalized function works as expected.""" angles = test_helpers.generate_preset_test_euler_angles().astype(dtype) matrix = rotation_matrix_3d.from_euler(angles) matrix_rescaled = matrix * 1.01 matrix_normalized = rotation_matrix_3d.assert_rotation_matrix_normalized( matrix) self.evaluate(matrix_normalized) with self.assertRaises(tf.errors.InvalidArgumentError): # pylint: disable=g-error-prone-assert-raises self.evaluate(rotation_matrix_3d.assert_rotation_matrix_normalized( matrix_rescaled))
def test_assert_rotation_matrix_normalized_passthrough(self): """Checks that the assert is a passthrough when the flag is False.""" angles = test_helpers.generate_preset_test_euler_angles() matrix_input = rotation_matrix_3d.from_euler(angles) matrix_output = rotation_matrix_3d.assert_rotation_matrix_normalized( matrix_input) self.assertTrue(matrix_input is matrix_output) # pylint: disable=g-generic-assert
def from_rotation_matrix(rotation_matrix: type_alias.TensorLike, name: str = "axis_angle_from_rotation_matrix" ) -> Tuple[tf.Tensor, tf.Tensor]: """Converts a rotation matrix to an axis-angle representation. Note: In the current version the returned axis-angle representation is not unique for a given rotation matrix. Since a direct conversion would not really be faster, we first transform the rotation matrix to a quaternion, and finally perform the conversion from that quaternion to the corresponding axis-angle representation. Note: In the following, A1 to An are optional batch dimensions. Args: rotation_matrix: A tensor of shape `[A1, ..., An, 3, 3]`, where the last two dimensions represent a rotation matrix. name: A name for this op that defaults to "axis_angle_from_rotation_matrix". Returns: A tuple of two tensors, respectively of shape `[A1, ..., An, 3]` and `[A1, ..., An, 1]`, where the first tensor represents the axis, and the second represents the angle. The resulting axis is a normalized vector. Raises: ValueError: If the shape of `rotation_matrix` is not supported. """ with tf.name_scope(name): rotation_matrix = tf.convert_to_tensor(value=rotation_matrix) shape.check_static( tensor=rotation_matrix, tensor_name="rotation_matrix", has_rank_greater_than=1, has_dim_equals=((-2, 3), (-1, 3))) rotation_matrix = rotation_matrix_3d.assert_rotation_matrix_normalized( rotation_matrix) quaternion = quaternion_lib.from_rotation_matrix(rotation_matrix) return from_quaternion(quaternion)
def from_rotation_matrix(rotation_matrix, name=None): """Converts a rotation matrix representation to a quaternion. Warning: This function is not smooth everywhere. Note: In the following, A1 to An are optional batch dimensions. Args: rotation_matrix: A tensor of shape `[A1, ..., An, 3, 3]`, where the last two dimensions represent a rotation matrix. name: A name for this op that defaults to "quaternion_from_rotation_matrix". Returns: A tensor of shape `[A1, ..., An, 4]`, where the last dimension represents a normalized quaternion. Raises: ValueError: If the shape of `rotation_matrix` is not supported. """ with tf.compat.v1.name_scope(name, "quaternion_from_rotation_matrix", [rotation_matrix]): rotation_matrix = tf.convert_to_tensor(value=rotation_matrix) shape.check_static( tensor=rotation_matrix, tensor_name="rotation_matrix", has_rank_greater_than=1, has_dim_equals=((-1, 3), (-2, 3))) rotation_matrix = rotation_matrix_3d.assert_rotation_matrix_normalized( rotation_matrix) trace = tf.linalg.trace(rotation_matrix) eps_addition = asserts.select_eps_for_addition(rotation_matrix.dtype) rows = tf.unstack(rotation_matrix, axis=-2) entries = [tf.unstack(row, axis=-1) for row in rows] def tr_positive(): sq = tf.sqrt(trace + 1.0) * 2. # sq = 4 * qw. qw = 0.25 * sq qx = safe_ops.safe_unsigned_div(entries[2][1] - entries[1][2], sq) qy = safe_ops.safe_unsigned_div(entries[0][2] - entries[2][0], sq) qz = safe_ops.safe_unsigned_div(entries[1][0] - entries[0][1], sq) return tf.stack((qx, qy, qz, qw), axis=-1) def cond_1(): sq = tf.sqrt(1.0 + entries[0][0] - entries[1][1] - entries[2][2] + eps_addition) * 2. # sq = 4 * qx. qw = safe_ops.safe_unsigned_div(entries[2][1] - entries[1][2], sq) qx = 0.25 * sq qy = safe_ops.safe_unsigned_div(entries[0][1] + entries[1][0], sq) qz = safe_ops.safe_unsigned_div(entries[0][2] + entries[2][0], sq) return tf.stack((qx, qy, qz, qw), axis=-1) def cond_2(): sq = tf.sqrt(1.0 + entries[1][1] - entries[0][0] - entries[2][2] + eps_addition) * 2. # sq = 4 * qy. qw = safe_ops.safe_unsigned_div(entries[0][2] - entries[2][0], sq) qx = safe_ops.safe_unsigned_div(entries[0][1] + entries[1][0], sq) qy = 0.25 * sq qz = safe_ops.safe_unsigned_div(entries[1][2] + entries[2][1], sq) return tf.stack((qx, qy, qz, qw), axis=-1) def cond_3(): sq = tf.sqrt(1.0 + entries[2][2] - entries[0][0] - entries[1][1] + eps_addition) * 2. # sq = 4 * qz. qw = safe_ops.safe_unsigned_div(entries[1][0] - entries[0][1], sq) qx = safe_ops.safe_unsigned_div(entries[0][2] + entries[2][0], sq) qy = safe_ops.safe_unsigned_div(entries[1][2] + entries[2][1], sq) qz = 0.25 * sq return tf.stack((qx, qy, qz, qw), axis=-1) def cond_idx(cond): cond = tf.expand_dims(cond, -1) cond = tf.tile(cond, [1] * (rotation_matrix.shape.ndims - 2) + [4]) return cond where_2 = tf.compat.v1.where( cond_idx(entries[1][1] > entries[2][2]), cond_2(), cond_3()) where_1 = tf.compat.v1.where( cond_idx((entries[0][0] > entries[1][1]) & (entries[0][0] > entries[2][2])), cond_1(), where_2) quat = tf.compat.v1.where(cond_idx(trace > 0), tr_positive(), where_1) return quat
def from_rotation_matrix(rotation_matrix, name=None): """Converts rotation matrices to Euler angles. The rotation matrices are assumed to have been constructed by rotation around the $$x$$, then $$y$$, and finally the $$z$$ axis. Note: There is an infinite number of solutions to this problem. There are Gimbal locks when abs(rotation_matrix(2,0)) == 1, which are not handled. Note: In the following, A1 to An are optional batch dimensions. Args: rotation_matrix: A tensor of shape `[A1, ..., An, 3, 3]`, where the last two dimensions represent a rotation matrix. name: A name for this op that defaults to "euler_from_rotation_matrix". Returns: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents the three Euler angles. Raises: ValueError: If the shape of `rotation_matrix` is not supported. """ def general_case(rotation_matrix, r20, eps_addition): """Handles the general case.""" theta_y = -tf.asin(r20) sign_cos_theta_y = safe_ops.nonzero_sign(tf.cos(theta_y)) r00 = rotation_matrix[..., 0, 0] r10 = rotation_matrix[..., 1, 0] r21 = rotation_matrix[..., 2, 1] r22 = rotation_matrix[..., 2, 2] r00 = safe_ops.nonzero_sign(r00) * eps_addition + r00 r22 = safe_ops.nonzero_sign(r22) * eps_addition + r22 # cos_theta_y evaluates to 0 on Gimbal locks, in which case the output of # this function will not be used. theta_z = tf.atan2(r10 * sign_cos_theta_y, r00 * sign_cos_theta_y) theta_x = tf.atan2(r21 * sign_cos_theta_y, r22 * sign_cos_theta_y) angles = tf.stack((theta_x, theta_y, theta_z), axis=-1) return angles def gimbal_lock(rotation_matrix, r20, eps_addition): """Handles Gimbal locks.""" r01 = rotation_matrix[..., 0, 1] r02 = rotation_matrix[..., 0, 2] sign_r20 = safe_ops.nonzero_sign(r20) r02 = safe_ops.nonzero_sign(r02) * eps_addition + r02 theta_x = tf.atan2(-sign_r20 * r01, -sign_r20 * r02) theta_y = -sign_r20 * tf.constant(math.pi / 2.0, dtype=r20.dtype) theta_z = tf.zeros_like(theta_x) angles = tf.stack((theta_x, theta_y, theta_z), axis=-1) return angles with tf.compat.v1.name_scope(name, "euler_from_rotation_matrix", [rotation_matrix]): rotation_matrix = tf.convert_to_tensor(value=rotation_matrix) shape.check_static(tensor=rotation_matrix, tensor_name="rotation_matrix", has_rank_greater_than=1, has_dim_equals=((-1, 3), (-2, 3))) rotation_matrix = rotation_matrix_3d.assert_rotation_matrix_normalized( rotation_matrix) r20 = rotation_matrix[..., 2, 0] eps_addition = asserts.select_eps_for_addition(rotation_matrix.dtype) general_solution = general_case(rotation_matrix, r20, eps_addition) gimbal_solution = gimbal_lock(rotation_matrix, r20, eps_addition) is_gimbal = tf.equal(tf.abs(r20), 1) gimbal_mask = tf.stack((is_gimbal, is_gimbal, is_gimbal), axis=-1) return tf.compat.v1.where(gimbal_mask, gimbal_solution, general_solution)
def from_rotation_matrix(rotation_matrix, name=None): """Converts a rotation matrix representation to a quaternion. Warning: This function is not smooth everywhere. Note: In the following, A1 to An are optional batch dimensions. Args: rotation_matrix: A tensor of shape `[A1, ..., An, 3, 3]`, where the last two dimensions represent a rotation matrix. name: A name for this op that defaults to "quaternion_from_rotation_matrix". Returns: A tensor of shape `[A1, ..., An, 4]`, where the last dimension represents a normalized quaternion. Raises: ValueError: If the shape of `rotation_matrix` is not supported. """ with tf.compat.v1.name_scope(name, "quaternion_from_rotation_matrix", [rotation_matrix]): rotation_matrix = tf.convert_to_tensor(value=rotation_matrix) shape.check_static( tensor=rotation_matrix, tensor_name="rotation_matrix", has_rank_greater_than=1, has_dim_equals=((-1, 3), (-2, 3))) rotation_matrix = rotation_matrix_3d.assert_rotation_matrix_normalized( rotation_matrix) eps_addition = asserts.select_eps_for_addition(rotation_matrix.dtype) def tr_positive(matrix, trace): sq = tf.sqrt(trace + 1.0) * 2. # sq = 4 * qw. qw = 0.25 * sq qx = safe_ops.safe_unsigned_div(matrix[2, 1] - matrix[1, 2], sq) qy = safe_ops.safe_unsigned_div(matrix[0, 2] - matrix[2, 0], sq) qz = safe_ops.safe_unsigned_div(matrix[1, 0] - matrix[0, 1], sq) return tf.convert_to_tensor([qx, qy, qz, qw], dtype=tf.float32) def cond_1(matrix): sq = tf.sqrt(1.0 + matrix[0, 0] - matrix[1, 1] - matrix[2, 2] + eps_addition) * 2. # sq = 4 * qx. qw = safe_ops.safe_unsigned_div(matrix[2, 1] - matrix[1, 2], sq) qx = 0.25 * sq qy = safe_ops.safe_unsigned_div(matrix[0, 1] + matrix[1, 0], sq) qz = safe_ops.safe_unsigned_div(matrix[0, 2] + matrix[2, 0], sq) return tf.convert_to_tensor([qx, qy, qz, qw], dtype=tf.float32) def cond_2(matrix): sq = tf.sqrt(1.0 + matrix[1, 1] - matrix[0, 0] - matrix[2, 2] + eps_addition) * 2. # sq = 4 * qy. qw = safe_ops.safe_unsigned_div(matrix[0, 2] - matrix[2, 0], sq) qx = safe_ops.safe_unsigned_div(matrix[0, 1] + matrix[1, 0], sq) qy = 0.25 * sq qz = safe_ops.safe_unsigned_div(matrix[1, 2] + matrix[2, 1], sq) return tf.convert_to_tensor([qx, qy, qz, qw], dtype=tf.float32) def cond_3(matrix): sq = tf.sqrt(1.0 + matrix[2, 2] - matrix[0, 0] - matrix[1, 1] + eps_addition) * 2. # sq = 4 * qz. qw = safe_ops.safe_unsigned_div(matrix[1, 0] - matrix[0, 1], sq) qx = safe_ops.safe_unsigned_div(matrix[0, 2] + matrix[2, 0], sq) qy = safe_ops.safe_unsigned_div(matrix[1, 2] + matrix[2, 1], sq) qz = 0.25 * sq return tf.convert_to_tensor([qx, qy, qz, qw], dtype=tf.float32) def quat_from_matrix(matrix): trace = tf.linalg.trace(matrix) where_2 = tf.cond(matrix[1, 1] > matrix[2, 2], lambda: cond_2(matrix), lambda: cond_3(matrix)) where_1 = tf.cond( (matrix[0, 0] > matrix[1, 1]) & (matrix[0, 0] > matrix[2, 2]), lambda: cond_1(matrix), lambda: where_2) return tf.cond(trace > 0, lambda: tr_positive(matrix, trace), lambda: where_1) return tf.map_fn(quat_from_matrix, rotation_matrix)