def project(point_3d, focal, principal_point, name=None): r"""Projects a 3d point onto the 2d camera plane. Projects a 3d point \\((x, y, z)\\) to a 2d point \\((x', y')\\) onto the image plane with $$ \begin{matrix} x' = \frac{f_x}{z}x + c_x, & y' = \frac{f_y}{z}y + c_y, \end{matrix} $$ where \\((f_x, f_y)\\) is the focal length and \\((c_x, c_y)\\) the principal point. Note: In the following, A1 to An are optional batch dimensions that must be broadcast compatible. Args: point_3d: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a 3d point to project. focal: A tensor of shape `[A1, ..., An, 2]`, where the last dimension represents a camera focal length. principal_point: A tensor of shape `[A1, ..., An, 2]`, where the last dimension represents a camera principal point. name: A name for this op that defaults to "perspective_project". Returns: A tensor of shape `[A1, ..., An, 2]`, where the last dimension represents a 2d point. Raises: ValueError: If the shape of `point_3d`, `focal`, or `principal_point` is not supported. """ with tf.compat.v1.name_scope(name, "perspective_project", [point_3d, focal, principal_point]): point_3d = tf.convert_to_tensor(value=point_3d) focal = tf.convert_to_tensor(value=focal) principal_point = tf.convert_to_tensor(value=principal_point) shape.check_static( tensor=point_3d, tensor_name="point_3d", has_dim_equals=(-1, 3)) shape.check_static( tensor=focal, tensor_name="focal", has_dim_equals=(-1, 2)) shape.check_static( tensor=principal_point, tensor_name="principal_point", has_dim_equals=(-1, 2)) shape.compare_batch_dimensions( tensors=(point_3d, focal, principal_point), tensor_names=("point_3d", "focal", "principal_point"), last_axes=-2, broadcast_compatible=True) point_2d, depth = tf.split(point_3d, (2, 1), axis=-1) point_2d *= safe_ops.safe_signed_div(focal, depth) point_2d += principal_point return point_2d
def between_two_vectors_3d(vector1, vector2, name=None): """Computes quaternion over the shortest arc between two vectors. Result quaternion describes shortest geodesic rotation from vector1 to vector2. Note: In the following, A1 to An are optional batch dimensions. Args: vector1: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents the first vector. vector2: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents the second vector. name: A name for this op that defaults to "quaternion_between_two_vectors_3d". Returns: A tensor of shape `[A1, ..., An, 4]`, where the last dimension represents a normalized quaternion. Raises: ValueError: If the shape of `vector1` or `vector2` is not supported. """ with tf.compat.v1.name_scope(name, "quaternion_between_two_vectors_3d", [vector1, vector2]): vector1 = tf.convert_to_tensor(value=vector1) vector2 = tf.convert_to_tensor(value=vector2) shape.check_static( tensor=vector1, tensor_name="vector1", has_dim_equals=(-1, 3)) shape.check_static( tensor=vector2, tensor_name="vector2", has_dim_equals=(-1, 3)) shape.compare_batch_dimensions( tensors=(vector1, vector2), last_axes=-2, broadcast_compatible=True) # Make sure that we are dealing with unit vectors. vector1 = tf.nn.l2_normalize(vector1, axis=-1) vector2 = tf.nn.l2_normalize(vector2, axis=-1) cos_theta = vector.dot(vector1, vector2) real_part = 1.0 + cos_theta axis = vector.cross(vector1, vector2) # Compute arbitrary antiparallel axes to rotate around in case of opposite # vectors. x, y, z = tf.split(vector1, (1, 1, 1), axis=-1) x_bigger_z = tf.abs(x) > tf.abs(z) x_bigger_z = tf.concat([x_bigger_z] * 3, axis=-1) antiparallel_axis = tf.compat.v1.where( x_bigger_z, tf.concat((-y, x, tf.zeros_like(z)), axis=-1), tf.concat((tf.zeros_like(x), -z, y), axis=-1)) # Compute rotation between two vectors. is_antiparallel = real_part < 1e-6 is_antiparallel = tf.concat([is_antiparallel] * 4, axis=-1) rot = tf.compat.v1.where( is_antiparallel, tf.concat((antiparallel_axis, tf.zeros_like(real_part)), axis=-1), tf.concat((axis, real_part), axis=-1)) return tf.nn.l2_normalize(rot, axis=-1)
def stratified_geomspace_1d(near: TensorLike, far: TensorLike, num_samples: int, name="stratified_geomspace_1d") -> tf.Tensor: """Stratified sampling based on evenly-spaced bins on a geometric progression. Args: near: A tensor of shape `[A1, ... An]` containing the starting points of the sampling interval. far: A tensor of shape `[A1, ... An]` containing the ending points of the sampling interval. num_samples: The number M of points to be sampled. name: A name for this op that defaults to "stratified". Returns: A tensor of shape `[A1, ..., An, M]` indicating the M points on the ray """ with tf.name_scope(name): near = tf.convert_to_tensor(near) far = tf.convert_to_tensor(far) shape.compare_batch_dimensions(tensors=(tf.expand_dims(near, axis=-1), tf.expand_dims(far, axis=-1)), tensor_names=("near", "far"), last_axes=-1, broadcast_compatible=True) bin_borders = geomspace_1d(near, far, num_samples + 1) bin_below = bin_borders[..., :-1] bin_above = bin_borders[..., 1:] target_shape = tf.concat([tf.shape(near), [num_samples]], axis=-1) random_point_in_bin = tf.random.uniform(target_shape) z_values = bin_below + (bin_above - bin_below) * random_point_in_bin return z_values
def uniform_1d(near: TensorLike, far: TensorLike, num_samples: int, name="uniform_1d") -> tf.Tensor: """Sample uniformly numbers on an interval, with the numbers being sorted. Args: near: A tensor of shape `[A1, ... An]` containing the starting points of the sampling interval. far: A tensor of shape `[A1, ... An]` containing the ending points of the sampling interval. num_samples: The number M of points to be sampled. name: A name for this op that defaults to "uniform_1d". Returns: A tensor of shape `[A1, ..., An, M]` indicating the M points on the ray """ with tf.name_scope(name): near = tf.convert_to_tensor(near) far = tf.convert_to_tensor(far) shape.compare_batch_dimensions(tensors=(tf.expand_dims(near, axis=-1), tf.expand_dims(far, axis=-1)), tensor_names=("near", "far"), last_axes=-1, broadcast_compatible=True) target_shape = tf.concat([tf.shape(near), [num_samples]], axis=-1) random_samples = tf.random.uniform(target_shape, minval=tf.expand_dims(near, -1), maxval=tf.expand_dims(far, -1)) return tf.sort(random_samples, axis=-1)
def logspace_1d(near: TensorLike, far: TensorLike, num_samples: int, base: float = 10.0, name="logspace_1d") -> tf.Tensor: """Sample evenly spaced numbers from an interval on a log scale. Args: near: A tensor of shape `[A1, ... An]` containing the starting points of the sampling interval. far: A tensor of shape `[A1, ... An]` containing the ending points of the sampling interval. num_samples: The number M of points to be sampled. base: The logarithmic base. name: A name for this op that defaults to "logspace_1d". Returns: A tensor of shape `[A1, ..., An, M]` indicating the M points on the ray """ with tf.name_scope(name): near = tf.convert_to_tensor(near) far = tf.convert_to_tensor(far) shape.compare_batch_dimensions(tensors=(tf.expand_dims(near, axis=-1), tf.expand_dims(far, axis=-1)), tensor_names=("near", "far"), last_axes=-1, broadcast_compatible=True) linspace = tf.linspace(near, far, num_samples, axis=-1) return tf.math.pow(base, linspace)
def geomspace_1d(near: TensorLike, far: TensorLike, num_samples: int, name="geomspace_1d") -> tf.Tensor: """Sample evenly spaced numbers on a geometric progression (log scale). Args: near: A tensor of shape `[A1, ... An]` containing the starting points of the sampling interval. far: A tensor of shape `[A1, ... An]` containing the ending points of the sampling interval. num_samples: The number M of points to be sampled. name: A name for this op that defaults to "geomspace_1d". Returns: A tensor of shape `[A1, ..., An, M]` indicating the M points on the ray """ with tf.name_scope(name): near = tf.convert_to_tensor(near) far = tf.convert_to_tensor(far) shape.compare_batch_dimensions(tensors=(tf.expand_dims(near, axis=-1), tf.expand_dims(far, axis=-1)), tensor_names=("near", "far"), last_axes=-1, broadcast_compatible=True) return logspace_1d(_log10(near), _log10(far), num_samples)
def camera_rays_from_extrinsics(rays, rotation_matrix, translation_vector): """Transform the rays from a camera located at (0, 0, 0) to ray origins and directions for a camera with given extrinsics. Args: rays: A tensor of shape `[A1, ..., An, N, 3]` where N is the number of rays. rotation_matrix: A tensor of shape `[A1, ..., An, 3, 3]`. translation_vector: A tensor of shape `[A1, ..., An, 3, 1]`. Returns: A tensor of shape `[A1, ..., An, N, 3]` representing the ray origin and a tensor of shape `[A1, ..., An, N, 3]` representing the ray direction. """ shape.check_static(tensor=rays, tensor_name="pixels", has_rank_greater_than=1) shape.compare_batch_dimensions(tensors=(rays, rotation_matrix, translation_vector), tensor_names=("points_on_rays", "rotation_matrix", "translation_vector"), last_axes=-3, broadcast_compatible=False) rays_org = _move_in_front_of_camera(tf.zeros_like(rays), rotation_matrix, translation_vector) rays_dir_ = _move_in_front_of_camera(rays, rotation_matrix, 0 * translation_vector) rays_dir = rays_dir_ / tf.norm(rays_dir_, axis=-1, keepdims=True) return rays_org, rays_dir
def from_axis_angle_translation( axis: type_alias.TensorLike, angle: type_alias.TensorLike, translation_vector: type_alias.TensorLike, name: str = "dual_quat_from_axis_angle_trans" ) -> type_alias.TensorLike: """Converts an axis-angle rotation and translation to a dual quaternion. Note: In the following, A1 to An are optional batch dimensions. Args: axis: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized axis. angle: A tensor of shape `[A1, ..., An, 1]`, where the last dimension represents an angle. translation_vector: A `[A1, ..., An, 3]`-tensor, where the last dimension represents a translation vector. name: A name for this op that defaults to "dual_quat_from_axis_angle_trans". Returns: A `[A1, ..., An, 8]`-tensor, where the last dimension represents a normalized dual quaternion. Raises: ValueError: If the shape of `axis`, `angle`, or `translation_vector` is not supported. """ with tf.name_scope(name): axis = tf.convert_to_tensor(value=axis) angle = tf.convert_to_tensor(value=angle) translation_vector = tf.convert_to_tensor(value=translation_vector) shape.check_static(tensor=axis, tensor_name="axis", has_dim_equals=(-1, 3)) shape.check_static(tensor=angle, tensor_name="angle", has_dim_equals=(-1, 1)) shape.check_static(tensor=translation_vector, tensor_name="translation_vector", has_dim_equals=(-1, 3)) shape.compare_batch_dimensions(tensors=(axis, angle, translation_vector), last_axes=-2, broadcast_compatible=True) scalar_shape = tf.concat((tf.shape(translation_vector)[:-1], (1, )), axis=-1) dtype = translation_vector.dtype quaternion_rotation = quaternion.from_axis_angle(axis, angle) quaternion_translation = tf.concat( (translation_vector, tf.zeros(scalar_shape, dtype)), axis=-1) dual_quaternion_dual_part = 0.5 * quaternion.multiply( quaternion_translation, quaternion_rotation) return tf.concat((quaternion_rotation, dual_quaternion_dual_part), axis=-1)
def from_axis_angle(axis, angle, name=None): """Convert an axis-angle representation to a rotation matrix. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. Args: axis: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized axis. angle: A tensor of shape `[A1, ..., An, 1]`, where the last dimension represents a normalized axis. name: A name for this op that defaults to "rotation_matrix_3d_from_axis_angle". Returns: A tensor of shape `[A1, ..., An, 3, 3]`, where the last two dimensions represents a 3d rotation matrix. Raises: ValueError: If the shape of `axis` or `angle` is not supported. """ with tf.compat.v1.name_scope(name, "rotation_matrix_3d_from_axis_angle", [axis, angle]): axis = tf.convert_to_tensor(value=axis) angle = tf.convert_to_tensor(value=angle) shape.check_static(tensor=axis, tensor_name="axis", has_dim_equals=(-1, 3)) shape.check_static( tensor=angle, tensor_name="angle", has_dim_equals=(-1, 1)) shape.compare_batch_dimensions( tensors=(axis, angle), tensor_names=("axis", "angle"), last_axes=-2, broadcast_compatible=True) axis = asserts.assert_normalized(axis) sin_axis = tf.sin(angle) * axis cos_angle = tf.cos(angle) cos1_axis = (1.0 - cos_angle) * axis _, axis_y, axis_z = tf.unstack(axis, axis=-1) cos1_axis_x, cos1_axis_y, _ = tf.unstack(cos1_axis, axis=-1) sin_axis_x, sin_axis_y, sin_axis_z = tf.unstack(sin_axis, axis=-1) tmp = cos1_axis_x * axis_y m01 = tmp - sin_axis_z m10 = tmp + sin_axis_z tmp = cos1_axis_x * axis_z m02 = tmp + sin_axis_y m20 = tmp - sin_axis_y tmp = cos1_axis_y * axis_z m12 = tmp - sin_axis_x m21 = tmp + sin_axis_x diag = cos1_axis * axis + cos_angle diag_x, diag_y, diag_z = tf.unstack(diag, axis=-1) matrix = tf.stack((diag_x, m01, m02, m10, diag_y, m12, m20, m21, diag_z), axis=-1) # pyformat: disable output_shape = tf.concat((tf.shape(input=axis)[:-1], (3, 3)), axis=-1) return tf.reshape(matrix, shape=output_shape)
def _points_from_z_values(ray_org: TensorLike, ray_dir: TensorLike, z_values: TensorLike) -> tf.Tensor: """Sample points on rays given the z values (distances along the rays). Args: ray_org: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents the 3D position of the ray origin. ray_dir: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents the 3D direction of the ray. z_values: A tensor of shape `[A1, ..., An, M]` containing the 1D position of M points along the ray. Returns: A tensor of shape `[A1, ..., An, M, 3]` """ shape.check_static(tensor=ray_dir, tensor_name="ray_dir", has_dim_equals=(-1, 3)) shape.check_static(tensor=ray_org, tensor_name="ray_org", has_dim_equals=(-1, 3)) shape.compare_batch_dimensions(tensors=(ray_org, ray_dir, z_values), tensor_names=("ray_org", "ray_dir", "z_values"), last_axes=-2, broadcast_compatible=False) points3d = (tf.expand_dims(ray_dir, axis=-2) * tf.expand_dims(z_values, axis=-1)) points3d = tf.expand_dims(ray_org, -2) + points3d return points3d
def dot(vector1, vector2, axis=-1, keepdims=True, name=None): """Computes the dot product between two tensors along an axis. Note: In the following, A1 to An are optional batch dimensions, which should be broadcast compatible. Args: vector1: Tensor of rank R and shape `[A1, ..., Ai, ..., An]`, where the dimension i = axis represents a vector. vector2: Tensor of rank R and shape `[A1, ..., Ai, ..., An]`, where the dimension i = axis represents a vector. axis: The dimension along which to compute the dot product. keepdims: If True, retains reduced dimensions with length 1. name: A name for this op which defaults to "vector_dot". Returns: A tensor of shape `[A1, ..., Ai = 1, ..., An]`, where the dimension i = axis represents the result of the dot product. """ with tf.compat.v1.name_scope(name, "vector_dot", [vector1, vector2]): vector1 = tf.convert_to_tensor(value=vector1) vector2 = tf.convert_to_tensor(value=vector2) shape.compare_batch_dimensions(tensors=(vector1, vector2), last_axes=-1, broadcast_compatible=True) shape.compare_dimensions(tensors=(vector1, vector2), axes=axis, tensor_names=("vector1", "vector2")) return tf.reduce_sum(input_tensor=vector1 * vector2, axis=axis, keepdims=keepdims)
def regular_inverse_1d(near: TensorLike, far: TensorLike, num_samples: int, name="regular_inverse_1d") -> tf.Tensor: """Sample inverse evenly spaced numbers on an interval. Args: near: A tensor of shape `[A1, ... An]` containing the starting points of the sampling interval. far: A tensor of shape `[A1, ... An]` containing the ending points of the sampling interval. num_samples: The number M of points to be sampled. name: A name for this op that defaults to "regular_inverse_1d". Returns: A tensor of shape `[A1, ..., An, M]` indicating the M points on the ray """ with tf.name_scope(name): near = tf.convert_to_tensor(near) far = tf.convert_to_tensor(far) shape.compare_batch_dimensions(tensors=(tf.expand_dims(near, axis=-1), tf.expand_dims(far, axis=-1)), tensor_names=("near", "far"), last_axes=-1, broadcast_compatible=True) return 1. / tf.linspace(1. / near, 1. / far, num_samples, axis=-1)
def is_normalized(axis, angle, atol=1e-3, name=None): """Determines if the axis-angle is normalized or not. Note: In the following, A1 to An are optional batch dimensions. Args: axis: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized axis. angle: A tensor of shape `[A1, ..., An, 1]` where the last dimension represents an angle. atol: The absolute tolerance parameter. name: A name for this op that defaults to "axis_angle_is_normalized". Returns: A tensor of shape `[A1, ..., An, 1]`, where False indicates that the axis is not normalized. """ with tf.compat.v1.name_scope(name, "axis_angle_is_normalized", [axis, angle]): axis = tf.convert_to_tensor(value=axis) angle = tf.convert_to_tensor(value=angle) shape.check_static(tensor=axis, tensor_name="axis", has_dim_equals=(-1, 3)) shape.check_static( tensor=angle, tensor_name="angle", has_dim_equals=(-1, 1)) shape.compare_batch_dimensions( tensors=(axis, angle), tensor_names=("axis", "angle"), last_axes=-2, broadcast_compatible=True) norms = tf.norm(tensor=axis, axis=-1, keepdims=True) return tf.abs(norms - 1.) < atol
def evaluate(point_set_a: type_alias.TensorLike, point_set_b: type_alias.TensorLike, name: str = "hausdorff_distance_evaluate") -> tf.Tensor: """Computes the Hausdorff distance from point_set_a to point_set_b. Note: Hausdorff distance from point_set_a to point_set_b is defined as the maximum of all distances from a point in point_set_a to the closest point in point_set_b. It is an asymmetric metric. Note: This function returns the exact Hausdorff distance and not an approximation. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. Args: point_set_a: A tensor of shape `[A1, ..., An, N, D]`, where the last axis represents points in a D dimensional space. point_set_b: A tensor of shape `[A1, ..., An, M, D]`, where the last axis represents points in a D dimensional space. name: A name for this op. Defaults to "hausdorff_distance_evaluate". Returns: A tensor of shape `[A1, ..., An]` storing the hausdorff distance from from point_set_a to point_set_b. Raises: ValueError: if the shape of `point_set_a`, `point_set_b` is not supported. """ with tf.name_scope(name): point_set_a = tf.convert_to_tensor(value=point_set_a) point_set_b = tf.convert_to_tensor(value=point_set_b) shape.compare_batch_dimensions( tensors=(point_set_a, point_set_b), tensor_names=("point_set_a", "point_set_b"), last_axes=-3, broadcast_compatible=True) # Verify that the last axis of the tensors has the same dimension. dimension = point_set_a.shape.as_list()[-1] shape.check_static( tensor=point_set_b, tensor_name="point_set_b", has_dim_equals=(-1, dimension)) # Create N x M matrix where the entry i,j corresponds to ai - bj (vector of # dimension D). difference = ( tf.expand_dims(point_set_a, axis=-2) - tf.expand_dims(point_set_b, axis=-3)) # Calculate the square distances between each two points: |ai - bj|^2. square_distances = tf.einsum("...i,...i->...", difference, difference) minimum_square_distance_a_to_b = tf.reduce_min( input_tensor=square_distances, axis=-1) return tf.sqrt( tf.reduce_max(input_tensor=minimum_square_distance_a_to_b, axis=-1))
def model_to_eye(point_model_space, camera_position, look_at_point, up_vector, name=None): """Transforms points from model to eye coordinates. Note: In the following, A1 to An are optional batch dimensions which must be broadcast compatible. Args: point_model_space: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents the 3D points in model space. camera_position: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents the 3D position of the camera. look_at_point: A tensor of shape `[A1, ..., An, 3]`, with the last dimension storing the position where the camera is looking at. up_vector: A tensor of shape `[A1, ..., An, 3]`, where the last dimension defines the up vector of the camera. name: A name for this op. Defaults to 'model_to_eye'. Raises: ValueError: if the all the inputs are not of the same shape, or if any input of of an unsupported shape. Returns: A tensor of shape `[A1, ..., An, 3]`, containing `point_model_space` in eye coordinates. """ with tf.compat.v1.name_scope( name, "model_to_eye", [point_model_space, camera_position, look_at_point, up_vector]): point_model_space = tf.convert_to_tensor(value=point_model_space) camera_position = tf.convert_to_tensor(value=camera_position) look_at_point = tf.convert_to_tensor(value=look_at_point) up_vector = tf.convert_to_tensor(value=up_vector) shape.check_static(tensor=point_model_space, tensor_name="point_model_space", has_dim_equals=(-1, 3)) shape.compare_batch_dimensions(tensors=(point_model_space, camera_position), last_axes=-2, tensor_names=("point_model_space", "camera_position"), broadcast_compatible=True) model_to_eye_matrix = look_at.right_handed(camera_position, look_at_point, up_vector) batch_shape = tf.shape(input=point_model_space)[:-1] one = tf.ones(shape=tf.concat((batch_shape, (1, )), axis=-1), dtype=point_model_space.dtype) point_model_space = tf.concat((point_model_space, one), axis=-1) point_model_space = tf.expand_dims(point_model_space, axis=-1) res = tf.squeeze(tf.matmul(model_to_eye_matrix, point_model_space), axis=-1) return res[..., :-1]
def vector_weights(vector1: type_alias.TensorLike, vector2: type_alias.TensorLike, percent: Union[type_alias.Float, type_alias.TensorLike], eps: Optional[type_alias.Float] = None, name: str = "vector_weights") -> Tuple[tf.Tensor, tf.Tensor]: """Spherical linear interpolation (slerp) between two unnormalized vectors. This function applies geometric slerp to unnormalized vectors by first normalizing them to return the interpolation weights. It reduces to lerp when input vectors are exactly anti-parallel. Note: In the following, A1 to An are optional batch dimensions. Args: vector1: A tensor of shape `[A1, ... , An, M]`, which stores a normalized vector in its last dimension. vector2: A tensor of shape `[A1, ... , An, M]`, which stores a normalized vector in its last dimension. percent: A `float` or tensor with shape broadcastable to the shape of input vectors. eps: A small float for operation safety. If left None, its value is automatically selected using dtype of input vectors. name: A name for this op. Defaults to "vector_weights". Raises: ValueError: if the shape of `vector1`, `vector2`, or `percent` is not supported. Returns: Two tensors of shape `[A1, ... , An, 1]`, representing interpolation weights for each input vector. """ with tf.name_scope(name): vector1 = tf.convert_to_tensor(value=vector1) vector2 = tf.convert_to_tensor(value=vector2) percent = tf.convert_to_tensor(value=percent, dtype=vector1.dtype) if percent.shape.ndims == 0: percent = tf.expand_dims(percent, axis=0) shape.compare_dimensions( tensors=(vector1, vector2), axes=-1, tensor_names=("vector1", "vector2")) shape.compare_batch_dimensions( tensors=(vector1, vector2, percent), last_axes=(-2, -2, -1), broadcast_compatible=True, tensor_names=("vector1", "vector2", "percent")) normalized1 = tf.nn.l2_normalize(vector1, axis=-1) normalized2 = tf.nn.l2_normalize(vector2, axis=-1) dot_product = _safe_dot(normalized1, normalized2, eps) theta = tf.acos(dot_product) scale1 = safe_ops.safe_sinpx_div_sinx(theta, 1.0 - percent, eps) scale2 = safe_ops.safe_sinpx_div_sinx(theta, percent, eps) return scale1, scale2
def rotate_zonal_harmonics( zonal_coeffs: TensorLike, theta: TensorLike, phi: TensorLike, name: str = "spherical_harmonics_rotate_zonal_harmonics" ) -> TensorLike: """Rotates zonal harmonics. Note: In the following, A1 to An are optional batch dimensions. Args: zonal_coeffs: A tensor of shape `[C]` storing zonal harmonics coefficients. theta: A tensor of shape `[A1, ..., An, 1]` storing polar angles. phi: A tensor of shape `[A1, ..., An, 1]` storing azimuthal angles. name: A name for this op. Defaults to "spherical_harmonics_rotate_zonal_harmonics". Returns: A tensor of shape `[A1, ..., An, C*C]` storing coefficients of the rotated harmonics. Raises: ValueError: If the shape of `zonal_coeffs`, `theta` or `phi` is not supported. """ with tf.name_scope(name): zonal_coeffs = tf.convert_to_tensor(value=zonal_coeffs) theta = tf.convert_to_tensor(value=theta) phi = tf.convert_to_tensor(value=phi) shape.check_static(tensor=zonal_coeffs, tensor_name="zonal_coeffs", has_rank=1) shape.check_static(tensor=phi, tensor_name="phi", has_dim_equals=(-1, 1)) shape.check_static(tensor=theta, tensor_name="theta", has_dim_equals=(-1, 1)) shape.compare_batch_dimensions(tensors=(theta, phi), last_axes=-2, tensor_names=("theta", "phi"), broadcast_compatible=False) tiled_zonal_coeffs = tile_zonal_coefficients(zonal_coeffs) max_band = zonal_coeffs.shape.as_list()[-1] l, m = generate_l_m_permutations(max_band - 1) broadcast_shape = theta.shape.as_list()[:-1] + l.shape.as_list() l_broadcasted = tf.broadcast_to(l, broadcast_shape) m_broadcasted = tf.broadcast_to(m, broadcast_shape) n_star = tf.sqrt(4.0 * np.pi / (2.0 * tf.cast(l, dtype=theta.dtype) + 1.0)) return n_star * tiled_zonal_coeffs * evaluate_spherical_harmonics( l_broadcasted, m_broadcasted, theta, phi)
def compute_radiance( rgba_values: type_alias.TensorLike, distances: type_alias.TensorLike, name: str = "ray_radiance") -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]: """Renders the rgba values for points along a ray, as described in ["NeRF Representing Scenes as Neural Radiance Fields for View Synthesis"](https://github.com/bmild/nerf). Note: In the following, A1 to An are optional batch dimensions. Args: rgba_values: A tensor of shape `[A1, ..., An, N, 4]`, where N are the samples on the ray. distances: A tensor of shape `[A1, ..., An, N]` containing the distances between the samples, where N are the samples on the ray. name: A name for this op. Defaults to "ray_radiance". Returns: A tensor of shape `[A1, ..., An, 3]` for the estimated rgb values, a tensor of shape `[A1, ..., An, 1]` for the estimated density values, and a tensor of shape `[A1, ..., An, N]` for the sample weights. """ with tf.name_scope(name): rgba_values = tf.convert_to_tensor(value=rgba_values) distances = tf.convert_to_tensor(value=distances) distances = tf.expand_dims(distances, -1) shape.check_static(tensor=rgba_values, tensor_name="rgba_values", has_dim_equals=(-1, 4)) shape.check_static(tensor=rgba_values, tensor_name="rgba_values", has_rank_greater_than=1) shape.check_static(tensor=distances, tensor_name="distances", has_rank_greater_than=1) shape.compare_batch_dimensions(tensors=(rgba_values, distances), tensor_names=("ray_values", "dists"), last_axes=-3, broadcast_compatible=True) shape.compare_dimensions(tensors=(rgba_values, distances), tensor_names=("ray_values", "dists"), axes=-2) rgb, density = tf.split(rgba_values, [3, 1], axis=-1) alpha = 1. - tf.exp(-density * distances) alpha = tf.squeeze(alpha, -1) ray_sample_weights = alpha * tf.math.cumprod( 1. - alpha + 1e-10, -1, exclusive=True) ray_rgb = tf.reduce_sum( input_tensor=tf.expand_dims(ray_sample_weights, -1) * rgb, axis=-2) ray_alpha = tf.expand_dims(tf.reduce_sum( input_tensor=ray_sample_weights, axis=-1), axis=-1) return ray_rgb, ray_alpha, ray_sample_weights
def rotate(point, axis, angle, name=None): r"""Rotates a 3d point using an axis-angle by applying the Rodrigues' formula. Rotates a vector $$\mathbf{v} \in {\mathbb{R}^3}$$ into a vector $$\mathbf{v}' \in {\mathbb{R}^3}$$ using the Rodrigues' rotation formula: $$\mathbf{v}'=\mathbf{v}\cos(\theta)+(\mathbf{a}\times\mathbf{v})\sin(\theta) +\mathbf{a}(\mathbf{a}\cdot\mathbf{v})(1-\cos(\theta)).$$ Note: In the following, A1 to An are optional batch dimensions. Args: point: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a 3d point to rotate. axis: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized axis. angle: A tensor of shape `[A1, ..., An, 1]`, where the last dimension represents an angle. name: A name for this op that defaults to "axis_angle_rotate". Returns: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a 3d point. Raises: ValueError: If `point`, `axis`, or `angle` are of different shape or if their respective shape is not supported. """ with tf.compat.v1.name_scope(name, "axis_angle_rotate", [point, axis, angle]): point = tf.convert_to_tensor(value=point) axis = tf.convert_to_tensor(value=axis) angle = tf.convert_to_tensor(value=angle) shape.check_static(tensor=point, tensor_name="point", has_dim_equals=(-1, 3)) shape.check_static(tensor=axis, tensor_name="axis", has_dim_equals=(-1, 3)) shape.check_static(tensor=angle, tensor_name="angle", has_dim_equals=(-1, 1)) shape.compare_batch_dimensions(tensors=(point, axis, angle), tensor_names=("point", "axis", "angle"), last_axes=-2, broadcast_compatible=True) axis = asserts.assert_normalized(axis) cos_angle = tf.cos(angle) axis_dot_point = vector.dot(axis, point) return point * cos_angle + vector.cross(axis, point) * tf.sin( angle) + axis * axis_dot_point * (1.0 - cos_angle)
def random_rays(focal: tf.Tensor, principal_point: tf.Tensor, height: int, width: int, n_rays: int, margin: int = 0, name: str = "random_rays") -> Tuple[tf.Tensor, tf.Tensor]: """Sample rays at random pixel location from a perspective camera. Args: focal: A tensor of shape `[A1, ..., An, 2]` where the last dimension contains the fx and fy focal length values. principal_point: A tensor of shape `[A1, ..., An, 2]` where the last dimension contains the cx and cy principal point values. height: The height of the image plane in pixels width: The width of the image plane in pixels. n_rays: The number M of rays to sample. margin: The margin around the borders of the image. name: A name for this op that defaults to "random_rays". Returns: A tensor of shape `[A1, ..., An, M, 3]` with the ray directions and a tensor of shape `[A1, ..., An, M, 2]` with the pixel x, y locations. """ with tf.name_scope(name): focal = tf.convert_to_tensor(value=focal) principal_point = tf.convert_to_tensor(value=principal_point) shape.check_static(tensor=focal, tensor_name="focal", has_dim_equals=(-1, 2)) shape.check_static(tensor=principal_point, tensor_name="principal_point", has_dim_equals=(-1, 2)) shape.compare_batch_dimensions(tensors=(focal, principal_point), tensor_names=("focal", "principal_point"), last_axes=-2, broadcast_compatible=True) batch_dims = tf.shape(focal)[:-1] target_shape = tf.concat([batch_dims, [n_rays]], axis=0) random_x = tf.random.uniform(target_shape, minval=margin, maxval=width - margin, dtype=tf.int32) random_y = tf.random.uniform(target_shape, minval=margin, maxval=height - margin, dtype=tf.int32) pixels = tf.cast(tf.stack((random_x, random_y), axis=-1), tf.float32) rays = ray(pixels, tf.expand_dims(focal, -2), tf.expand_dims(principal_point, -2)) return rays, tf.cast(pixels, tf.int32)
def linear_coefficients(matte: type_alias.TensorLike, pseudo_inverse: type_alias.TensorLike, name: str = "matting_linear_coefficients" ) -> Tuple[tf.Tensor, tf.Tensor]: """Computes the matting linear coefficients. Computes the matting linear coefficients (a, b) based on the `pseudo_inverse` generated by the `build_matrices` function which implements the approach proposed by Levin et al. in "A Closed Form Solution to Natural Image Matting". Args: matte: A tensor of shape `[B, H, W, 1]`. pseudo_inverse: A tensor of shape `[B, H - pad, W - pad, C + 1, size^2]` containing the pseudo-inverse matrices computed by the `build_matrices` function, where `pad` is equal to `size - 1` and `size` is the patch size used to compute this tensor. name: A name for this op. Defaults to "matting_linear_coefficients". Returns: A tuple contraining two Tensors for the linear coefficients (a, b) of shape `[B, H, W, C]` and `[B, H, W, 1]`. Raises: ValueError: If the last dimension of `matte` is not 1. If `matte` is not of rank 4. If `pseudo_inverse` is not of rank 5. If `B` is different between `matte` and `pseudo_inverse`. """ with tf.name_scope(name): matte = tf.convert_to_tensor(value=matte) pseudo_inverse = tf.convert_to_tensor(value=pseudo_inverse) pixels = tf.compat.dimension_value(pseudo_inverse.shape[-1]) shape.check_static(matte, has_rank=4, has_dim_equals=(-1, 1)) shape.check_static(pseudo_inverse, has_rank=5) shape.compare_batch_dimensions( tensors=(matte, pseudo_inverse), last_axes=0, broadcast_compatible=False) size = np.sqrt(pixels) # Computes the linear coefficients. patches = tf.expand_dims(_image_patches(matte, size), axis=-1) coeffs = tf.squeeze(tf.matmul(pseudo_inverse, patches), axis=-1) # Averages the linear coefficients over patches. height = tf.shape(input=coeffs)[1] width = tf.shape(input=coeffs)[2] ones = tf.ones(shape=_shape((1,), height, width, 1), dtype=matte.dtype) height = tf.shape(input=matte)[1] + size - 1 width = tf.shape(input=matte)[2] + size - 1 coeffs = tf.image.resize_with_crop_or_pad(coeffs, height, width) ones = tf.image.resize_with_crop_or_pad(ones, height, width) coeffs = _image_average(coeffs, size) / _image_average(ones, size) return tf.split(coeffs, (-1, 1), axis=-1)
def evaluate(ground_truth: type_alias.TensorLike, prediction: type_alias.TensorLike, precision_function: Callable[..., Any] = precision_module.evaluate, recall_function: Callable[..., Any] = recall_module.evaluate, name: str = "fscore_evaluate") -> tf.Tensor: """Computes the fscore metric for the given ground truth and predicted labels. The fscore is calculated as 2 * (precision * recall) / (precision + recall) where the precision and recall are evaluated by the given function parameters. The precision and recall functions default to their definition for boolean labels (see https://en.wikipedia.org/wiki/Precision_and_recall for more details). Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. Args: ground_truth: A tensor of shape `[A1, ..., An, N]`, where the last axis represents the ground truth values. prediction: A tensor of shape `[A1, ..., An, N]`, where the last axis represents the predicted values. precision_function: The function to use for evaluating the precision. Defaults to the precision evaluation for binary ground-truth and predictions. recall_function: The function to use for evaluating the recall. Defaults to the recall evaluation for binary ground-truth and prediction. name: A name for this op. Defaults to "fscore_evaluate". Returns: A tensor of shape `[A1, ..., An]` that stores the fscore metric for the given ground truth labels and predictions. Raises: ValueError: if the shape of `ground_truth`, `prediction` is not supported. """ with tf.name_scope(name): ground_truth = tf.convert_to_tensor(value=ground_truth) prediction = tf.convert_to_tensor(value=prediction) shape.compare_batch_dimensions(tensors=(ground_truth, prediction), tensor_names=("ground_truth", "prediction"), last_axes=-1, broadcast_compatible=True) recall = recall_function(ground_truth, prediction) precision = precision_function(ground_truth, prediction) return safe_ops.safe_signed_div(2 * precision * recall, precision + recall)
def rotate(point: type_alias.TensorLike, matrix: type_alias.TensorLike, name: str = "rotation_matrix_3d_rotate") -> tf.Tensor: """Rotate a point using a rotation matrix 3d. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. Args: point: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a 3d point. matrix: A tensor of shape `[A1, ..., An, 3,3]`, where the last dimension represents a 3d rotation matrix. name: A name for this op that defaults to "rotation_matrix_3d_rotate". Returns: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a 3d point. Raises: ValueError: If the shape of `point` or `rotation_matrix_3d` is not supported. """ with tf.name_scope(name): point = tf.convert_to_tensor(value=point) matrix = tf.convert_to_tensor(value=matrix) shape.check_static(tensor=point, tensor_name="point", has_dim_equals=(-1, 3)) shape.check_static(tensor=matrix, tensor_name="matrix", has_rank_greater_than=1, has_dim_equals=((-2, 3), (-1, 3))) shape.compare_batch_dimensions(tensors=(point, matrix), tensor_names=("point", "matrix"), last_axes=(-2, -3), broadcast_compatible=True) matrix = assert_rotation_matrix_normalized(matrix) point = tf.expand_dims(point, axis=-1) common_batch_shape = shape.get_broadcasted_shape( point.shape[:-2], matrix.shape[:-2]) def dim_value(dim): return 1 if dim is None else tf.compat.dimension_value(dim) common_batch_shape = [dim_value(dim) for dim in common_batch_shape] point = tf.broadcast_to(point, common_batch_shape + [3, 1]) matrix = tf.broadcast_to(matrix, common_batch_shape + [3, 3]) rotated_point = tf.matmul(matrix, point) return tf.squeeze(rotated_point, axis=-1)
def check_valid_graph_convolution_input(data: type_alias.TensorLike, neighbors: tf.sparse.SparseTensor, sizes: type_alias.TensorLike): """Checks that the inputs are valid for graph convolution ops. Note: In the following, A1 to An are optional batch dimensions. Args: data: A `float` tensor with shape `[A1, ..., An, V1, V2]`. neighbors: A SparseTensor with the same type as `data` and with shape `[A1, ..., An, V1, V1]`. sizes: An `int` tensor of shape `[A1, ..., An]`. Optional, can be `None`. Raises: TypeError: if the input types are invalid. ValueError: if the input dimensions are invalid. """ if not data.dtype.is_floating: raise TypeError("'data' must have a float type.") if neighbors.dtype != data.dtype: raise TypeError("'neighbors' and 'data' must have the same type.") if sizes is not None and not sizes.dtype.is_integer: raise TypeError("'sizes' must have an integer type.") if not isinstance(neighbors, tf.sparse.SparseTensor): raise ValueError("'neighbors' must be a SparseTensor.") data_ndims = data.shape.ndims shape.check_static(tensor=data, tensor_name="data", has_rank_greater_than=1) shape.check_static(tensor=neighbors, tensor_name="neighbors", has_rank=data_ndims) if not _is_dynamic_shape(tensors=(data, neighbors)): shape.compare_dimensions(tensors=(data, neighbors, neighbors), tensor_names=("data", "neighbors", "neighbors"), axes=(-2, -2, -1)) if sizes is None: shape.compare_batch_dimensions(tensors=(data, neighbors), tensor_names=("data", "neighbors"), last_axes=-3, broadcast_compatible=False) else: shape.check_static(tensor=sizes, tensor_name="sizes", has_rank=data_ndims - 2) shape.compare_batch_dimensions(tensors=(data, neighbors, sizes), tensor_names=("data", "neighbors", "sizes"), last_axes=(-3, -3, -1), broadcast_compatible=False)
def check_valid_graph_unpooling_input(data: type_alias.TensorLike, pool_map: tf.sparse.SparseTensor, sizes: type_alias.TensorLike): """Checks that the inputs are valid for graph unpooling. Note: In the following, A1 to A3 are optional batch dimensions. Args: data: A `float` tensor with shape `[A1, ..., A3, V1, C]`. pool_map: A `SparseTensor` with the same type as `data` and with shape `[A1, ..., A3, V1, V2]`. sizes: An `int` tensor of shape `[A1, ..., A3, 2]`. Can be `None`. Raises: TypeError: if the input types are invalid. ValueError: if the input dimensions are invalid. """ if not data.dtype.is_floating: raise TypeError("'data' must have a float type.") if pool_map.dtype != data.dtype: raise TypeError("'pool_map' and 'data' must have the same type.") if sizes is not None and not sizes.dtype.is_integer: raise TypeError("'sizes' must have an integer type.") if not isinstance(pool_map, tf.sparse.SparseTensor): raise ValueError("'pool_map' must be a SparseTensor.") data_ndims = data.shape.ndims shape.check_static(tensor=data, tensor_name="data", has_rank_greater_than=1) shape.check_static(tensor=data, tensor_name="data", has_rank_less_than=6) shape.check_static(tensor=pool_map, tensor_name="pool_map", has_rank=data_ndims) if not _is_dynamic_shape(tensors=(data, pool_map)): shape.compare_dimensions(tensors=(data, pool_map), tensor_names=("data", "pool_map"), axes=(-2, -2)) if sizes is None: shape.compare_batch_dimensions(tensors=(data, pool_map), tensor_names=("data", "pool_map"), last_axes=-3, broadcast_compatible=False) else: shape.check_static(tensor=sizes, tensor_name="sizes", has_rank=data_ndims - 1) shape.compare_batch_dimensions(tensors=(data, pool_map, sizes), tensor_names=("data", "pool_map", "sizes"), last_axes=(-3, -3, -2), broadcast_compatible=False)
def evaluate(ground_truth_labels: type_alias.TensorLike, predicted_labels: type_alias.TensorLike, grid_size: int = 1, name: str = "intersection_over_union_evaluate") -> tf.Tensor: """Computes the Intersection-Over-Union metric for the given ground truth and predicted labels. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible, and G1 to Gm are the grid dimensions. Args: ground_truth_labels: A tensor of shape `[A1, ..., An, G1, ..., Gm]`, where the last m axes represent a grid of ground truth attributes. Each attribute can either be 0 or 1. predicted_labels: A tensor of shape `[A1, ..., An, G1, ..., Gm]`, where the last m axes represent a grid of predicted attributes. Each attribute can either be 0 or 1. grid_size: The number of grid dimensions. Defaults to 1. name: A name for this op. Defaults to "intersection_over_union_evaluate". Returns: A tensor of shape `[A1, ..., An]` that stores the intersection-over-union metric of the given ground truth labels and predictions. Raises: ValueError: if the shape of `ground_truth_labels`, `predicted_labels` is not supported. """ with tf.name_scope(name): ground_truth_labels = tf.convert_to_tensor(value=ground_truth_labels) predicted_labels = tf.convert_to_tensor(value=predicted_labels) shape.compare_batch_dimensions(tensors=(ground_truth_labels, predicted_labels), tensor_names=("ground_truth_labels", "predicted_labels"), last_axes=-grid_size, broadcast_compatible=True) ground_truth_labels = asserts.assert_binary(ground_truth_labels) predicted_labels = asserts.assert_binary(predicted_labels) sum_ground_truth = tf.math.reduce_sum(input_tensor=ground_truth_labels, axis=list(range(-grid_size, 0))) sum_predictions = tf.math.reduce_sum(input_tensor=predicted_labels, axis=list(range(-grid_size, 0))) intersection = tf.math.reduce_sum(input_tensor=ground_truth_labels * predicted_labels, axis=list(range(-grid_size, 0))) union = sum_ground_truth + sum_predictions - intersection return tf.where(tf.math.equal(union, 0), tf.ones_like(union), intersection / union)
def inverse_transform_sampling_1d( bins: TensorLike, pdf: TensorLike, num_samples: int, name="inverse_transform_sampling_1d") -> tf.Tensor: """Sampling 1D points from a distribution using the inverse transform. The target distrubution is defined by its probability density function and the spatial 1D location of its bins. The new random samples correspond to the centers of the bins. Args: bins: A tensor of shape `[A1, ..., An, M]` containing 1D location of M bins. For example, a tensor [a, b, c, d] corresponds to the bin structure |--a--|-b-|--c--|d|. pdf: A tensor of shape `[A1, ..., An, M]` containing the probability distribution in M bins. num_samples: The number N of new samples. name: A name for this op that defaults to "inverse_transform_sampling_1d". Returns: A tensor of shape `[A1, ..., An, N]` indicating the new N random points. """ with tf.name_scope(name): bins = tf.convert_to_tensor(value=bins) pdf = tf.convert_to_tensor(value=pdf) shape.check_static(tensor=bins, tensor_name="bins", has_rank_greater_than=0) shape.check_static(tensor=pdf, tensor_name="pdf", has_rank_greater_than=0) shape.compare_batch_dimensions(tensors=(bins, pdf), tensor_names=("bins", "pdf"), last_axes=-2, broadcast_compatible=True) shape.compare_dimensions(tensors=(bins, pdf), tensor_names=("bins", "pdf"), axes=-1) # Do not consider the last bin, as the cdf contains has +1 dimension. pdf = _normalize_pdf(pdf[..., :-1]) cdf = _get_cdf(pdf) batch_shape = tf.shape(pdf)[:-1] # TODO(krematas): Use dynamic values batch_dims = tf.get_static_value(tf.rank(pdf) - 1) target_shape = tf.concat([batch_shape, [num_samples]], axis=-1) uniform_samples = tf.random.uniform(target_shape) bin_indices = tf.searchsorted(cdf, uniform_samples, side="right") bin_indices = tf.maximum(0, bin_indices - 1) z_values = tf.gather(bins, bin_indices, axis=-1, batch_dims=batch_dims) return z_values
def gather_faces(vertices: type_alias.TensorLike, indices: type_alias.TensorLike, name: str = "normals_gather_faces") -> type_alias.TensorLike: """Gather corresponding vertices for each face. Note: In the following, A1 to An are optional batch dimensions. Args: vertices: A tensor of shape `[A1, ..., An, V, D]`, where `V` is the number of vertices and `D` the dimensionality of each vertex. The rank of this tensor should be at least 2. indices: A tensor of shape `[A1, ..., An, F, M]`, where `F` is the number of faces, and `M` is the number of vertices per face. The rank of this tensor should be at least 2. name: A name for this op. Defaults to "normals_gather_faces". Returns: A tensor of shape `[A1, ..., An, F, M, D]` containing the vertices of each face. Raises: ValueError: If the shape of `vertices` or `indices` is not supported. """ with tf.name_scope(name): vertices = tf.convert_to_tensor(value=vertices) indices = tf.convert_to_tensor(value=indices) shape.check_static(tensor=vertices, tensor_name="vertices", has_rank_greater_than=1) shape.check_static(tensor=indices, tensor_name="indices", has_rank_greater_than=1) shape.compare_batch_dimensions(tensors=(vertices, indices), last_axes=(-3, -3), broadcast_compatible=False) if hasattr(tf, "batch_gather"): expanded_vertices = tf.expand_dims(vertices, axis=-3) broadcasted_shape = tf.concat( [tf.shape(input=indices)[:-1], tf.shape(input=vertices)[-2:]], axis=-1) broadcasted_vertices = tf.broadcast_to(expanded_vertices, broadcasted_shape) return tf.gather(broadcasted_vertices, indices, batch_dims=-1) else: return tf.gather(vertices, indices, axis=-2, batch_dims=indices.shape.ndims - 2)
def rotate(point, matrix, name=None): """Rotates a 2d point using a 2d rotation matrix. Note: In the following, A1 to An are optional batch dimensions, which must be identical. Args: point: A tensor of shape `[A1, ..., An, 2]`, where the last dimension represents a 2d point. matrix: A tensor of shape `[A1, ..., An, 2, 2]`, where the last two dimensions represent a 2d rotation matrix. name: A name for this op that defaults to "rotation_matrix_2d_rotate". Returns: A tensor of shape `[A1, ..., An, 2]`, where the last dimension represents a 2d point. Raises: ValueError: If the shape of `point` or `matrix` is not supported. """ with tf.compat.v1.name_scope(name, "rotation_matrix_2d_rotate", [point, matrix]): point = tf.convert_to_tensor(value=point) matrix = tf.convert_to_tensor(value=matrix) shape.check_static(tensor=point, tensor_name="point", has_dim_equals=(-1, 2)) shape.check_static(tensor=matrix, tensor_name="matrix", has_rank_greater_than=1, has_dim_equals=((-2, 2), (-1, 2))) shape.compare_batch_dimensions(tensors=(point, matrix), tensor_names=("point", "matrix"), last_axes=(-2, -3), broadcast_compatible=True) point = tf.expand_dims(point, axis=-1) common_batch_shape = shape.get_broadcasted_shape( point.shape[:-2], matrix.shape[:-2]) def dim_value(dim): return 1 if dim is None else tf.compat.v1.dimension_value(dim) common_batch_shape = [dim_value(dim) for dim in common_batch_shape] point = tf.broadcast_to(point, common_batch_shape + [2, 1]) matrix = tf.broadcast_to(matrix, common_batch_shape + [2, 2]) rotated_point = tf.matmul(matrix, point) return tf.squeeze(rotated_point, axis=-1)
def normal(v0: type_alias.TensorLike, v1: type_alias.TensorLike, v2: type_alias.TensorLike, clockwise: bool = False, normalize: bool = True, name: str = "triangle_normal") -> tf.Tensor: """Computes face normals (triangles). Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. Args: v0: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents the first vertex of a triangle. v1: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents the second vertex of a triangle. v2: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents the third vertex of a triangle. clockwise: Winding order to determine front-facing triangles. normalize: A `bool` indicating whether output normals should be normalized by the function. name: A name for this op. Defaults to "triangle_normal". Returns: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized vector. Raises: ValueError: If the shape of `v0`, `v1`, or `v2` is not supported. """ with tf.name_scope(name): v0 = tf.convert_to_tensor(value=v0) v1 = tf.convert_to_tensor(value=v1) v2 = tf.convert_to_tensor(value=v2) shape.check_static(tensor=v0, tensor_name="v0", has_dim_equals=(-1, 3)) shape.check_static(tensor=v1, tensor_name="v1", has_dim_equals=(-1, 3)) shape.check_static(tensor=v2, tensor_name="v2", has_dim_equals=(-1, 3)) shape.compare_batch_dimensions(tensors=(v0, v1, v2), last_axes=-2, broadcast_compatible=True) normal_vector = vector.cross(v1 - v0, v2 - v0, axis=-1) normal_vector = asserts.assert_nonzero_norm(normal_vector) if not clockwise: normal_vector *= -1.0 if normalize: return tf.nn.l2_normalize(normal_vector, axis=-1) return normal_vector