def test_normal_random(self, clockwise): """Tests the triangle normal computation in each axis.""" tensor_size = np.random.randint(3) tensor_shape = np.random.randint(1, 10, size=(tensor_size)).tolist() zeros = np.zeros(shape=tensor_shape + [1]) for i in range(3): v0 = np.random.random(size=tensor_shape + [3]) v1 = np.random.random(size=tensor_shape + [3]) v2 = np.random.random(size=tensor_shape + [3]) v0[..., i] = 0. v1[..., i] = 0. v2[..., i] = 0. n = np.zeros_like(v0) n[..., i] = 1. normal = triangle.normal(v0, v1, v2, clockwise) with self.subTest(name="n"): self.assertAllClose(tf.abs(normal), n) with self.subTest(name="v1-v0"): self.assertAllClose(vector.dot(normal, (v1 - v0)), zeros) with self.subTest(name="v2-v0"): self.assertAllClose(vector.dot(normal, (v2 - v0)), zeros)
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 _safe_dot(vector1, vector2, eps): """Calculates dot product while ensuring it is in the range [-1, 1].""" dot_product = vector.dot(vector1, vector2) # Safely shrink to make sure machine precision does not cause the dot # product to be outside the [-1.0, 1.0] range. return safe_ops.safe_shrink( vector=dot_product, minval=-1.0, maxval=1.0, open_bounds=False, eps=eps)
def test_dot_jacobian_preset(self, u_init, v_init): """Tests the Jacobian of the dot product.""" u_tensor = tf.convert_to_tensor(value=u_init) v_tensor = tf.convert_to_tensor(value=v_init) y = vector.dot(u_tensor, v_tensor) self.assert_jacobian_is_correct(u_tensor, u_init, y) self.assert_jacobian_is_correct(v_tensor, v_init, y)
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 test_dot_random(self): """Tests the dot product function.""" tensor_size = np.random.randint(2, 4) tensor_shape = np.random.randint(1, 10, size=tensor_size).tolist() axis = np.random.randint(tensor_size) u = np.random.random(size=tensor_shape) v = np.random.random(size=tensor_shape) dot = tf.linalg.tensor_diag_part(tf.tensordot(u, v, axes=[[axis], [axis]])) dot = tf.expand_dims(dot, axis=axis) self.assertAllClose(vector.dot(u, v, axis=axis), dot)
def test_dot_jacobian_random(self): """Tests the Jacobian of the dot product.""" tensor_size = np.random.randint(3) tensor_shape = np.random.randint(1, 10, size=(tensor_size)).tolist() u_init = np.random.random(size=tensor_shape + [3]) v_init = np.random.random(size=tensor_shape + [3]) u_tensor = tf.convert_to_tensor(value=u_init) v_tensor = tf.convert_to_tensor(value=v_init) y = vector.dot(u_tensor, v_tensor) self.assert_jacobian_is_correct(u_tensor, u_init, y) self.assert_jacobian_is_correct(v_tensor, v_init, y)
def inverse(dual_quaternion: type_alias.TensorLike, name: str = "dual_quaternion_inverse") -> tf.Tensor: """Computes the inverse of a dual quaternion. Note: In the following, A1 to An are optional batch dimensions. Args: dual_quaternion: A TensorLike of shape `[A1, ..., An, 8]`, where the last dimension represents a dual quaternion. name: A name for this op that defaults to "dual_quaternion_inverse". Returns: A tensor of shape `[A1, ..., An, 8]`, where the last dimension represents a dual quaternion. Raises: ValueError: If the shape of `dual quaternion` is not supported. """ with tf.name_scope(name): dual_quaternion = tf.convert_to_tensor(value=dual_quaternion) shape.check_static(tensor=dual_quaternion, tensor_name="dual_quaternion", has_dim_equals=(-1, 8)) quaternion_real, quaternion_dual = tf.split(dual_quaternion, (4, 4), axis=-1) quaternion_real_norm_squared = tf.norm(tensor=quaternion_real, axis=-1, keepdims=True)**2 quaternion_real_conj = quaternion.conjugate(quaternion_real) quaternion_output_real = safe_ops.safe_signed_div( quaternion_real_conj, quaternion_real_norm_squared) normalized_dual = safe_ops.safe_signed_div( quaternion.conjugate(quaternion_dual), quaternion_real_norm_squared) normalized_dot_product = safe_ops.safe_signed_div( vector.dot(quaternion_real, quaternion_dual, keepdims=True), quaternion_real_norm_squared**2) quaternion_output_dual = ( normalized_dual - 2 * quaternion_real_conj * normalized_dot_product) return tf.concat((quaternion_output_real, quaternion_output_dual), axis=-1)
def relative_angle(quaternion1: type_alias.TensorLike, quaternion2: type_alias.TensorLike, name: str = "quaternion_relative_angle") -> tf.Tensor: r"""Computes the unsigned relative rotation angle between 2 unit quaternions. Given two normalized quanternions $$\mathbf{q}_1$$ and $$\mathbf{q}_2$$, the relative angle is computed as $$\theta = 2\arccos(\mathbf{q}_1^T\mathbf{q}_2)$$. Note: In the following, A1 to An are optional batch dimensions. Args: quaternion1: A tensor of shape `[A1, ..., An, 4]`, where the last dimension represents a normalized quaternion. quaternion2: A tensor of shape `[A1, ..., An, 4]`, where the last dimension represents a normalized quaternion. name: A name for this op that defaults to "quaternion_relative_angle". Returns: A tensor of shape `[A1, ..., An, 1]` where the last dimension represents rotation angles in the range [0.0, pi]. Raises: ValueError: If the shape of `quaternion1` or `quaternion2` is not supported. """ with tf.name_scope(name): quaternion1 = tf.convert_to_tensor(value=quaternion1) quaternion2 = tf.convert_to_tensor(value=quaternion2) shape.check_static(tensor=quaternion1, tensor_name="quaternion1", has_dim_equals=(-1, 4)) shape.check_static(tensor=quaternion2, tensor_name="quaternion2", has_dim_equals=(-1, 4)) quaternion1 = asserts.assert_normalized(quaternion1) quaternion2 = asserts.assert_normalized(quaternion2) dot_product = vector.dot(quaternion1, quaternion2, keepdims=False) # Ensure dot product is in range [-1. 1]. eps_dot_prod = 4.0 * asserts.select_eps_for_addition(dot_product.dtype) dot_product = safe_ops.safe_shrink(dot_product, -1.0, 1.0, False, eps=eps_dot_prod) return 2.0 * tf.acos(tf.abs(dot_product))
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 we deal with unit vectors. vector1 = tf.nn.l2_normalize(vector1, axis=-1) vector2 = tf.nn.l2_normalize(vector2, axis=-1) axis = vector.cross(vector1, vector2) cos_theta = vector.dot(vector1, vector2) rot = tf.concat((axis, 1. + cos_theta), axis=-1) return tf.nn.l2_normalize(rot, axis=-1)
def distance_to_ray(point: type_alias.TensorLike, origin: type_alias.TensorLike, direction: type_alias.TensorLike, keepdims: bool = True, name: str = "point_distance_to_ray" ) -> tf.Tensor: """Computes the distance from a M-d point to a M-d ray. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. Args: point: A tensor of shape `[A1, ..., An, M]`. origin: A tensor of shape `[A1, ..., An, M]`. direction: A tensor of shape `[A1, ..., An, M]`. The last dimension must be normalized. keepdims: A `bool`, whether to keep the last dimension with length 1 or to remove it. name: A name for this op. Defaults to "point_distance_to_ray". Returns: A tensor of shape `[A1, ..., An, 1]` containing the distance from each point to the corresponding ray. Raises: ValueError: If the shape of `point`, `origin`, or 'direction' is not supported. """ with tf.name_scope(name): point = tf.convert_to_tensor(value=point) origin = tf.convert_to_tensor(value=origin) direction = tf.convert_to_tensor(value=direction) shape.compare_dimensions((point, origin, direction), -1, ("point", "origin", "direction")) shape.compare_batch_dimensions( tensors=(point, origin, direction), last_axes=-2, broadcast_compatible=True) direction = asserts.assert_normalized(direction) vec = point - origin dot = vector.dot(vec, direction) vec -= dot * direction return tf.norm(tensor=vec, axis=-1, keepdims=keepdims)
def d_q(q1, q2): """Distance between 2 quaternions The quaternion distance takes values between [0, pi] Parameters ---------- q1: tf.tensor/np.ndarray 1st quaternion q2: tf.tensor/np.ndarray 2nd quaternion Returns ------- : distnace between these 2 quaternions """ q1 = tf.cast(tf.convert_to_tensor(value=q1), dtype=tf.float64) q2 = tf.cast(tf.convert_to_tensor(value=q2), dtype=tf.float64) shape.check_static(tensor=q1, tensor_name="quaternion1", has_dim_equals=(-1, 4)) shape.check_static(tensor=q2, tensor_name="quaternion2", has_dim_equals=(-1, 4)) q1 = quaternion.normalize(q1) q2 = quaternion.normalize(q2) dot_product = vector.dot(q1, q2, keepdims=False) # Ensure dot product is in range [-1. 1]. eps_dot_prod = 1.8 * asserts.select_eps_for_addition(dot_product.dtype) dot_product = safe_ops.safe_shrink(dot_product, -1, 1, open_bounds=False, eps=eps_dot_prod) return 2.0 * tf.acos(tf.abs(dot_product))
def integration_product(harmonics1, harmonics2, keepdims=True, name="spherical_harmonics_convolution"): """Computes the integral of harmonics1.harmonics2 over the sphere. Note: In the following, A1 to An are optional batch dimensions. Args: harmonics1: A tensor of shape `[A1, ..., An, C]`, where the last dimension represents spherical harmonics coefficients. harmonics2: A tensor of shape `[A1, ..., An, C]`, where the last dimension represents spherical harmonics coefficients. keepdims: If True, retains reduced dimensions with length 1. name: A name for this op. Defaults to "spherical_harmonics_convolution". Returns: A tensor of shape `[A1, ..., An]` containing scalar values resulting from integrating the product of the spherical harmonics `harmonics1` and `harmonics2`. Raises: ValueError: if the last dimension of `harmonics1` is different from the last dimension of `harmonics2`. """ with tf.name_scope(name): harmonics1 = tf.convert_to_tensor(value=harmonics1) harmonics2 = tf.convert_to_tensor(value=harmonics2) shape.compare_dimensions( tensors=(harmonics1, harmonics2), axes=-1, tensor_names=("harmonics1", "harmonics2")) shape.compare_batch_dimensions( tensors=(harmonics1, harmonics2), last_axes=-2, tensor_names=("harmonics1", "harmonics2"), broadcast_compatible=True) return vector.dot(harmonics1, harmonics2, keepdims=keepdims)
def project_to_ray(point: type_alias.TensorLike, origin: type_alias.TensorLike, direction: type_alias.TensorLike, name: str = "point_project_to_ray" ) -> tf.Tensor: """Computes the projection of a M-d point on a M-d ray. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. Args: point: A tensor of shape `[A1, ..., An, M]`. origin: A tensor of shape `[A1, ..., An, M]`. direction: A tensor of shape `[A1, ..., An, M]`. The last dimension must be normalized. name: A name for this op. Defaults to "point_project_to_ray". Returns: A tensor of shape `[A1, ..., An, M]` containing the projected point. Raises: ValueError: If the shape of `point`, `origin`, or 'direction' is not supported. """ with tf.name_scope(name): point = tf.convert_to_tensor(value=point) origin = tf.convert_to_tensor(value=origin) direction = tf.convert_to_tensor(value=direction) shape.compare_dimensions((point, origin, direction), -1, ("point", "origin", "direction")) shape.compare_batch_dimensions( tensors=(point, origin, direction), last_axes=-2, broadcast_compatible=True) direction = asserts.assert_normalized(direction) vec = point - origin dot = vector.dot(vec, direction) return origin + dot * direction
def reconstruct(image, coeff_mul, coeff_add, name=None): """Reconstruct the matte from the image using the linear coefficients. Reconstruct the matte from the image using the linear coefficients (a, b) returned by the linear_coefficients function. Args: image: A tensor of shape `[B, H, W, C]` . coeff_mul: A tensor of shape `[B, H, W, C]` representing the multiplicative part of the linear coefficients. coeff_add: A tensor of shape `[B, H, W, 1]` representing the additive part of the linear coefficients. name: A name for this op. Defaults to "matting_reconstruct". Returns: A tensor of shape `[B, H, W, 1]` containing the mattes. Raises: ValueError: If `image`, `coeff_mul`, or `coeff_add` are not of rank 4. If the last dimension of `coeff_add` is not 1. If the batch dimensions of `image`, `coeff_mul`, and `coeff_add` do not match. """ with tf.compat.v1.name_scope(name, "matting_reconstruct", [image, coeff_mul, coeff_add]): image = tf.convert_to_tensor(value=image) coeff_mul = tf.convert_to_tensor(value=coeff_mul) coeff_add = tf.convert_to_tensor(value=coeff_add) shape.check_static(image, has_rank=4) shape.check_static(coeff_mul, has_rank=4) shape.check_static(coeff_add, has_rank=4, has_dim_equals=(-1, 1)) shape.compare_batch_dimensions(tensors=(image, coeff_mul), last_axes=-1, broadcast_compatible=False) shape.compare_batch_dimensions(tensors=(image, coeff_add), last_axes=-2, broadcast_compatible=False) return tfg_vector.dot(coeff_mul, image) + coeff_add
def interpolate(points, weights, indices, normalize=True, allow_negative_weights=False, name=None): """Weighted interpolation for M-D point sets. Given an M-D point set, this function can be used to generate a new point set that is formed by interpolating a subset of points in the set. Note: In the following, A1 to An, and B1 to Bk are optional batch dimensions. Args: points: A tensor with shape `[B1, ..., Bk, M] and rank R > 1, where M is the dimensionality of the points. weights: A tensor with shape `[A1, ..., An, P]`, where P is the number of points to interpolate for each output point. indices: A tensor of dtype tf.int32 and shape `[A1, ..., An, P, R-1]`, which contains the point indices to be used for each output point. The R-1 dimensional axis gives the slice index of a single point in `points`. The first n+1 dimensions of weights and indices must match, or be broadcast compatible. normalize: A `bool` describing whether or not to normalize the weights on the last axis. allow_negative_weights: A `bool` describing whether or not negative weights are allowed. name: A name for this op. Defaults to "weighted_interpolate". Returns: A tensor of shape `[A1, ..., An, M]` storing the interpolated M-D points. The first n dimensions will be the same as weights and indices. """ with tf.compat.v1.name_scope(name, "weighted_interpolate", [points, weights, indices]): points = tf.convert_to_tensor(value=points) weights = tf.convert_to_tensor(value=weights) indices = tf.convert_to_tensor(value=indices) shape.check_static(tensor=points, tensor_name="points", has_rank_greater_than=1) shape.check_static(tensor=indices, tensor_name="indices", has_rank_greater_than=1, has_dim_equals=(-1, points.shape.ndims - 1)) shape.compare_dimensions(tensors=(weights, indices), axes=(-1, -2), tensor_names=("weights", "indices")) shape.compare_batch_dimensions(tensors=(weights, indices), last_axes=(-2, -3), tensor_names=("weights", "indices"), broadcast_compatible=True) if not allow_negative_weights: weights = asserts.assert_all_above(weights, 0.0, open_bound=False) if normalize: sums = tf.reduce_sum(input_tensor=weights, axis=-1, keepdims=True) sums = asserts.assert_nonzero_norm(sums) weights = safe_ops.safe_signed_div(weights, sums) point_lists = tf.gather_nd(points, indices) return vector.dot(point_lists, tf.expand_dims(weights, axis=-1), axis=-2, keepdims=False)
def brdf(direction_incoming_light: type_alias.TensorLike, direction_outgoing_light: type_alias.TensorLike, surface_normal: type_alias.TensorLike, albedo: type_alias.TensorLike, name: str = "lambertian_brdf") -> tf.Tensor: """Evaluates the brdf of a Lambertian surface. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. Note: The gradient of this function is not smooth when the dot product of the normal with any light is 0.0. Args: direction_incoming_light: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized incoming light vector. direction_outgoing_light: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized outgoing light vector. surface_normal: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized surface normal. albedo: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents albedo with values in [0,1]. name: A name for this op. Defaults to "lambertian_brdf". Returns: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents the amount of reflected light in any outgoing direction. Raises: ValueError: if the shape of `direction_incoming_light`, `direction_outgoing_light`, `surface_normal`, `shininess` or `albedo` is not supported. InvalidArgumentError: if at least one element of `albedo` is outside of [0,1]. """ with tf.name_scope(name): direction_incoming_light = tf.convert_to_tensor( value=direction_incoming_light) direction_outgoing_light = tf.convert_to_tensor( value=direction_outgoing_light) surface_normal = tf.convert_to_tensor(value=surface_normal) albedo = tf.convert_to_tensor(value=albedo) shape.check_static(tensor=direction_incoming_light, tensor_name="direction_incoming_light", has_dim_equals=(-1, 3)) shape.check_static(tensor=direction_outgoing_light, tensor_name="direction_outgoing_light", has_dim_equals=(-1, 3)) shape.check_static(tensor=surface_normal, tensor_name="surface_normal", has_dim_equals=(-1, 3)) shape.check_static(tensor=albedo, tensor_name="albedo", has_dim_equals=(-1, 3)) shape.compare_batch_dimensions( tensors=(direction_incoming_light, direction_outgoing_light, surface_normal, albedo), tensor_names=("direction_incoming_light", "direction_outgoing_light", "surface_normal", "albedo"), last_axes=-2, broadcast_compatible=True) direction_incoming_light = asserts.assert_normalized( direction_incoming_light) direction_outgoing_light = asserts.assert_normalized( direction_outgoing_light) surface_normal = asserts.assert_normalized(surface_normal) albedo = asserts.assert_all_in_range(albedo, 0.0, 1.0, open_bounds=False) # Checks whether the incoming or outgoing light point behind the surface. dot_incoming_light_surface_normal = vector.dot( -direction_incoming_light, surface_normal) dot_outgoing_light_surface_normal = vector.dot( direction_outgoing_light, surface_normal) min_dot = tf.minimum(dot_incoming_light_surface_normal, dot_outgoing_light_surface_normal) common_shape = shape.get_broadcasted_shape(min_dot.shape, albedo.shape) d_val = lambda dim: 1 if dim is None else tf.compat.dimension_value(dim ) common_shape = [d_val(dim) for dim in common_shape] condition = tf.broadcast_to(tf.greater_equal(min_dot, 0.0), common_shape) albedo = tf.broadcast_to(albedo, common_shape) return tf.where(condition, albedo / math.pi, tf.zeros_like(albedo))
def energy(vertices_rest_pose, vertices_deformed_pose, quaternions, edges, vertex_weight=None, edge_weight=None, conformal_energy=True, aggregate_loss=True, name=None): """Estimates an As Conformal As Possible (ACAP) fitting energy. For a given mesh in rest pose, this function evaluates a variant of the ACAP [1] fitting energy for a batch of deformed meshes. The vertex weights and edge weights are defined on the rest pose. The method implemented here is similar to [2], but with an added free variable capturing a scale factor per vertex. [1]: Yusuke Yoshiyasu, Wan-Chun Ma, Eiichi Yoshida, and Fumio Kanehiro. "As-Conformal-As-Possible Surface Registration." Computer Graphics Forum. Vol. 33. No. 5. 2014.</br> [2]: Olga Sorkine, and Marc Alexa. "As-rigid-as-possible surface modeling". Symposium on Geometry Processing. Vol. 4. 2007. Note: In the description of the arguments, V corresponds to the number of vertices in the mesh, and E to the number of edges in this mesh. Note: In the following, A1 to An are optional batch dimensions. Args: vertices_rest_pose: A tensor of shape `[V, 3]` containing the position of all the vertices of the mesh in rest pose. vertices_deformed_pose: A tensor of shape `[A1, ..., An, V, 3]` containing the position of all the vertices of the mesh in deformed pose. quaternions: A tensor of shape `[A1, ..., An, V, 4]` defining a rigid transformation to apply to each vertex of the rest pose. See Section 2 from [1] for further details. edges: A tensor of shape `[E, 2]` defining indices of vertices that are connected by an edge. vertex_weight: An optional tensor of shape `[V]` defining the weight associated with each vertex. Defaults to a tensor of ones. edge_weight: A tensor of shape `[E]` defining the weight of edges. Common choices for these weights include uniform weighting, and cotangent weights. Defaults to a tensor of ones. conformal_energy: A `bool` indicating whether each vertex is associated with a scale factor or not. If this parameter is True, scaling information must be encoded in the norm of `quaternions`. If this parameter is False, this function implements the energy described in [2]. aggregate_loss: A `bool` defining whether the returned loss should be an aggregate measure. When True, the mean squared error is returned. When False, returns two losses for every edge of the mesh. name: A name for this op. Defaults to "as_conformal_as_possible_energy". Returns: When aggregate_loss is `True`, returns a tensor of shape `[A1, ..., An]` containing the ACAP energies. When aggregate_loss is `False`, returns a tensor of shape `[A1, ..., An, 2*E]` containing each term of the summation described in the equation 7 of [2]. Raises: ValueError: if the shape of `vertices_rest_pose`, `vertices_deformed_pose`, `quaternions`, `edges`, `vertex_weight`, or `edge_weight` is not supported. """ with tf.compat.v1.name_scope(name, "as_conformal_as_possible_energy", [ vertices_rest_pose, vertices_deformed_pose, quaternions, edges, conformal_energy, vertex_weight, edge_weight ]): vertices_rest_pose = tf.convert_to_tensor(value=vertices_rest_pose) vertices_deformed_pose = tf.convert_to_tensor( value=vertices_deformed_pose) quaternions = tf.convert_to_tensor(value=quaternions) edges = tf.convert_to_tensor(value=edges) if vertex_weight is not None: vertex_weight = tf.convert_to_tensor(value=vertex_weight) if edge_weight is not None: edge_weight = tf.convert_to_tensor(value=edge_weight) shape.check_static(tensor=vertices_rest_pose, tensor_name="vertices_rest_pose", has_rank=2, has_dim_equals=(-1, 3)) shape.check_static(tensor=vertices_deformed_pose, tensor_name="vertices_deformed_pose", has_rank_greater_than=1, has_dim_equals=(-1, 3)) shape.check_static(tensor=quaternions, tensor_name="quaternions", has_rank_greater_than=1, has_dim_equals=(-1, 4)) shape.compare_batch_dimensions(tensors=(vertices_deformed_pose, quaternions), last_axes=(-3, -3), broadcast_compatible=False) shape.check_static(tensor=edges, tensor_name="edges", has_rank=2, has_dim_equals=(-1, 2)) tensors_with_vertices = [ vertices_rest_pose, vertices_deformed_pose, quaternions ] names_with_vertices = [ "vertices_rest_pose", "vertices_deformed_pose", "quaternions" ] axes_with_vertices = [-2, -2, -2] if vertex_weight is not None: shape.check_static(tensor=vertex_weight, tensor_name="vertex_weight", has_rank=1) tensors_with_vertices.append(vertex_weight) names_with_vertices.append("vertex_weight") axes_with_vertices.append(0) shape.compare_dimensions(tensors=tensors_with_vertices, axes=axes_with_vertices, tensor_names=names_with_vertices) if edge_weight is not None: shape.check_static(tensor=edge_weight, tensor_name="edge_weight", has_rank=1) shape.compare_dimensions(tensors=(edges, edge_weight), axes=(0, 0), tensor_names=("edges", "edge_weight")) if not conformal_energy: quaternions = quaternion.normalize(quaternions) # Extracts the indices of vertices. indices_i, indices_j = tf.unstack(edges, axis=-1) # Extracts the vertices we need per term. vertices_i_rest = tf.gather(vertices_rest_pose, indices_i, axis=-2) vertices_j_rest = tf.gather(vertices_rest_pose, indices_j, axis=-2) vertices_i_deformed = tf.gather(vertices_deformed_pose, indices_i, axis=-2) vertices_j_deformed = tf.gather(vertices_deformed_pose, indices_j, axis=-2) # Extracts the weights we need per term. weights_shape = vertices_i_rest.shape.as_list()[-2] if vertex_weight is not None: weight_i = tf.gather(vertex_weight, indices_i) weight_j = tf.gather(vertex_weight, indices_j) else: weight_i = weight_j = tf.ones(weights_shape, dtype=vertices_rest_pose.dtype) weight_i = tf.expand_dims(weight_i, axis=-1) weight_j = tf.expand_dims(weight_j, axis=-1) if edge_weight is not None: weight_ij = edge_weight else: weight_ij = tf.ones(weights_shape, dtype=vertices_rest_pose.dtype) weight_ij = tf.expand_dims(weight_ij, axis=-1) # Extracts the rotation we need per term. quaternion_i = tf.gather(quaternions, indices_i, axis=-2) quaternion_j = tf.gather(quaternions, indices_j, axis=-2) # Computes the energy. deformed_ij = vertices_i_deformed - vertices_j_deformed rotated_rest_ij = quaternion.rotate( (vertices_i_rest - vertices_j_rest), quaternion_i) energy_ij = weight_i * weight_ij * (deformed_ij - rotated_rest_ij) deformed_ji = vertices_j_deformed - vertices_i_deformed rotated_rest_ji = quaternion.rotate( (vertices_j_rest - vertices_i_rest), quaternion_j) energy_ji = weight_j * weight_ij * (deformed_ji - rotated_rest_ji) energy_ij_squared = vector.dot(energy_ij, energy_ij, keepdims=False) energy_ji_squared = vector.dot(energy_ji, energy_ji, keepdims=False) if aggregate_loss: average_energy_ij = tf.reduce_mean(input_tensor=energy_ij_squared, axis=-1) average_energy_ji = tf.reduce_mean(input_tensor=energy_ji_squared, axis=-1) return (average_energy_ij + average_energy_ji) / 2.0 return tf.concat((energy_ij_squared, energy_ji_squared), axis=-1)
def estimate_radiance(point_light_radiance, point_light_position, surface_point_position, surface_point_normal, observation_point, brdf, name="estimate_radiance", reflected_light_fall_off=False): """Estimates the spectral radiance of a point light reflected from the surface point towards the observation point. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. B1 to Bm are optional batch dimensions for the lights, which must be broadcast compatible. Note: In case the light or the observation point are located behind the surface the function will return 0. Note: The gradient of this function is not smooth when the dot product of the normal with the light-to-surface or surface-to-observation vectors is 0. Args: point_light_radiance: A tensor of shape '[B1, ..., Bm, K]', where the last axis represents the radiance of the point light at a specific wave length. point_light_position: A tensor of shape `[B1, ..., Bm, 3]`, where the last axis represents the position of the point light. surface_point_position: A tensor of shape `[A1, ..., An, 3]`, where the last axis represents the position of the surface point. surface_point_normal: A tensor of shape `[A1, ..., An, 3]`, where the last axis represents the normalized surface normal at the given surface point. observation_point: A tensor of shape `[A1, ..., An, 3]`, where the last axis represents the observation point. brdf: The BRDF of the surface as a function of: incoming_light_direction - The incoming light direction as the last axis of a tensor with shape `[A1, ..., An, 3]`. outgoing_light_direction - The outgoing light direction as the last axis of a tensor with shape `[A1, ..., An, 3]`. surface_point_normal - The surface normal as the last axis of a tensor with shape `[A1, ..., An, 3]`. Note - The BRDF should return a tensor of size '[A1, ..., An, K]' where the last axis represents the amount of reflected light in each wave length. name: A name for this op. Defaults to "estimate_radiance". reflected_light_fall_off: A boolean specifying whether or not to include the fall off of the light reflected from the surface towards the observation point in the calculation. Defaults to False. Returns: A tensor of shape `[A1, ..., An, B1, ..., Bm, K]`, where the last axis represents the amount of light received at the observation point after being reflected from the given surface point. Raises: ValueError: if the shape of `point_light_position`, `surface_point_position`, `surface_point_normal`, or `observation_point` is not supported. InvalidArgumentError: if 'surface_point_normal' is not normalized. """ with tf.name_scope(name): point_light_radiance = tf.convert_to_tensor(value=point_light_radiance) point_light_position = tf.convert_to_tensor(value=point_light_position) surface_point_position = tf.convert_to_tensor(value=surface_point_position) surface_point_normal = tf.convert_to_tensor(value=surface_point_normal) observation_point = tf.convert_to_tensor(value=observation_point) shape.check_static( tensor=point_light_position, tensor_name="point_light_position", has_dim_equals=(-1, 3)) shape.check_static( tensor=surface_point_position, tensor_name="surface_point_position", has_dim_equals=(-1, 3)) shape.check_static( tensor=surface_point_normal, tensor_name="surface_point_normal", has_dim_equals=(-1, 3)) shape.check_static( tensor=observation_point, tensor_name="observation_point", has_dim_equals=(-1, 3)) shape.compare_batch_dimensions( tensors=(surface_point_position, surface_point_normal, observation_point), tensor_names=("surface_point_position", "surface_point_normal", "observation_point"), last_axes=-2, broadcast_compatible=True) shape.compare_batch_dimensions( tensors=(point_light_radiance, point_light_position), tensor_names=("point_light_radiance", "point_light_position"), last_axes=-2, broadcast_compatible=True) surface_point_normal = asserts.assert_normalized(surface_point_normal) # Get the number of lights dimensions (B1,...,Bm). lights_num_dimensions = max( len(point_light_radiance.shape), len(point_light_position.shape)) - 1 # Reshape the other parameters so they can be broadcasted to the output of # shape [A1,...,An, B1,...,Bm, K]. surface_point_position = tf.reshape( surface_point_position, surface_point_position.shape[:-1] + (1,) * lights_num_dimensions + (3,)) surface_point_normal = tf.reshape( surface_point_normal, surface_point_normal.shape[:-1] + (1,) * lights_num_dimensions + (3,)) observation_point = tf.reshape( observation_point, observation_point.shape[:-1] + (1,) * lights_num_dimensions + (3,)) light_to_surface_point = surface_point_position - point_light_position distance_light_surface_point = tf.norm( tensor=light_to_surface_point, axis=-1, keepdims=True) incoming_light_direction = tf.math.l2_normalize( light_to_surface_point, axis=-1) surface_to_observation_point = observation_point - surface_point_position outgoing_light_direction = tf.math.l2_normalize( surface_to_observation_point, axis=-1) brdf_value = brdf(incoming_light_direction, outgoing_light_direction, surface_point_normal) incoming_light_dot_surface_normal = vector.dot(-incoming_light_direction, surface_point_normal) outgoing_light_dot_surface_normal = vector.dot(outgoing_light_direction, surface_point_normal) estimated_radiance = (point_light_radiance * \ brdf_value * incoming_light_dot_surface_normal) / \ (4. * math.pi * tf.math.square(distance_light_surface_point)) if reflected_light_fall_off: distance_surface_observation_point = tf.norm( tensor=surface_to_observation_point, axis=-1, keepdims=True) estimated_radiance = estimated_radiance / \ tf.math.square(distance_surface_observation_point) # Create a condition for checking whether the light or observation point are # behind the surface. min_dot = tf.minimum(incoming_light_dot_surface_normal, outgoing_light_dot_surface_normal) common_shape = shape.get_broadcasted_shape(min_dot.shape, estimated_radiance.shape) d_val = lambda dim: 1 if dim is None else tf.compat.dimension_value(dim) common_shape = [d_val(dim) for dim in common_shape] condition = tf.broadcast_to(tf.greater_equal(min_dot, 0.0), common_shape) return tf.where(condition, estimated_radiance, tf.zeros_like(estimated_radiance))
def brdf(direction_incoming_light: type_alias.TensorLike, direction_outgoing_light: type_alias.TensorLike, surface_normal: type_alias.TensorLike, shininess: type_alias.TensorLike, albedo: type_alias.TensorLike, brdf_normalization: bool = True, name: str = "phong_brdf") -> tf.Tensor: """Evaluates the specular brdf of the Phong model. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. Note: The gradient of this function is not smooth when the dot product of the normal with any light is 0.0. Args: direction_incoming_light: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized incoming light vector. direction_outgoing_light: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized outgoing light vector. surface_normal: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized surface normal. shininess: A tensor of shape `[A1, ..., An, 1]`, where the last dimension represents a non-negative shininess coefficient. albedo: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents albedo with values in [0,1]. brdf_normalization: A `bool` indicating whether normalization should be applied to enforce the energy conservation property of BRDFs. Note that `brdf_normalization` must be set to False in order to use the original Blinn specular model. name: A name for this op. Defaults to "phong_brdf". Returns: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents the amount of light reflected in the outgoing light direction. Raises: ValueError: if the shape of `direction_incoming_light`, `direction_outgoing_light`, `surface_normal`, `shininess` or `albedo` is not supported. InvalidArgumentError: if not all of shininess values are non-negative, or if at least one element of `albedo` is outside of [0,1]. """ with tf.name_scope(name): direction_incoming_light = tf.convert_to_tensor( value=direction_incoming_light) direction_outgoing_light = tf.convert_to_tensor( value=direction_outgoing_light) surface_normal = tf.convert_to_tensor(value=surface_normal) shininess = tf.convert_to_tensor(value=shininess) albedo = tf.convert_to_tensor(value=albedo) shape.check_static(tensor=direction_incoming_light, tensor_name="direction_incoming_light", has_dim_equals=(-1, 3)) shape.check_static(tensor=direction_outgoing_light, tensor_name="direction_outgoing_light", has_dim_equals=(-1, 3)) shape.check_static(tensor=surface_normal, tensor_name="surface_normal", has_dim_equals=(-1, 3)) shape.check_static(tensor=shininess, tensor_name="shininess", has_dim_equals=(-1, 1)) shape.check_static(tensor=albedo, tensor_name="albedo", has_dim_equals=(-1, 3)) shape.compare_batch_dimensions( tensors=(direction_incoming_light, direction_outgoing_light, surface_normal, shininess, albedo), tensor_names=("direction_incoming_light", "direction_outgoing_light", "surface_normal", "shininess", "albedo"), last_axes=-2, broadcast_compatible=True) direction_incoming_light = asserts.assert_normalized( direction_incoming_light) direction_outgoing_light = asserts.assert_normalized( direction_outgoing_light) surface_normal = asserts.assert_normalized(surface_normal) albedo = asserts.assert_all_in_range(albedo, 0.0, 1.0, open_bounds=False) shininess = asserts.assert_all_above(shininess, 0.0, open_bound=False) # Checks whether the incoming or outgoing light point behind the surface. dot_incoming_light_surface_normal = vector.dot( -direction_incoming_light, surface_normal) dot_outgoing_light_surface_normal = vector.dot( direction_outgoing_light, surface_normal) min_dot = tf.minimum(dot_incoming_light_surface_normal, dot_outgoing_light_surface_normal) perfect_reflection_direction = vector.reflect(direction_incoming_light, surface_normal) perfect_reflection_direction = tf.math.l2_normalize( perfect_reflection_direction, axis=-1) cos_alpha = vector.dot(perfect_reflection_direction, direction_outgoing_light, axis=-1) cos_alpha = tf.maximum(cos_alpha, tf.zeros_like(cos_alpha)) phong_model = albedo * tf.pow(cos_alpha, shininess) if brdf_normalization: phong_model *= _brdf_normalization_factor(shininess) common_shape = shape.get_broadcasted_shape(min_dot.shape, phong_model.shape) d_val = lambda dim: 1 if dim is None else tf.compat.dimension_value(dim ) common_shape = [d_val(dim) for dim in common_shape] condition = tf.broadcast_to(tf.greater_equal(min_dot, 0.0), common_shape) phong_model = tf.broadcast_to(phong_model, common_shape) return tf.where(condition, phong_model, tf.zeros_like(phong_model))
def right_handed(camera_position, look_at, up_vector, name=None): """Builds a right handed look at view matrix. Note: In the following, A1 to An are optional batch dimensions. Args: camera_position: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents the 3D position of the camera. look_at: 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 'right_handed'. 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, 4, 4]`, containing right handed look at matrices. """ with tf.compat.v1.name_scope(name, "right_handed", [camera_position, look_at, up_vector]): camera_position = tf.convert_to_tensor(value=camera_position) look_at = tf.convert_to_tensor(value=look_at) up_vector = tf.convert_to_tensor(value=up_vector) shape.check_static(tensor=camera_position, tensor_name="camera_position", has_dim_equals=(-1, 3)) shape.check_static(tensor=look_at, tensor_name="look_at", has_dim_equals=(-1, 3)) shape.check_static(tensor=up_vector, tensor_name="up_vector", has_dim_equals=(-1, 3)) shape.compare_batch_dimensions(tensors=(camera_position, look_at, up_vector), last_axes=-2, tensor_names=("camera_position", "look_at", "up_vector"), broadcast_compatible=False) z_axis = tf.linalg.l2_normalize(look_at - camera_position, axis=-1) horizontal_axis = tf.linalg.l2_normalize(vector.cross( z_axis, up_vector), axis=-1) vertical_axis = vector.cross(horizontal_axis, z_axis) batch_shape = tf.shape(input=horizontal_axis)[:-1] zeros = tf.zeros(shape=tf.concat((batch_shape, (3, )), axis=-1), dtype=horizontal_axis.dtype) one = tf.ones(shape=tf.concat((batch_shape, (1, )), axis=-1), dtype=horizontal_axis.dtype) matrix = tf.concat( (horizontal_axis, -vector.dot(horizontal_axis, camera_position), vertical_axis, -vector.dot(vertical_axis, camera_position), -z_axis, vector.dot(z_axis, camera_position), zeros, one), axis=-1) matrix_shape = tf.shape(input=matrix) output_shape = tf.concat((matrix_shape[:-1], (4, 4)), axis=-1) return tf.reshape(matrix, shape=output_shape)
def intersection_ray_sphere(sphere_center, sphere_radius, ray, point_on_ray, name=None): """Finds positions and surface normals where the sphere and the ray intersect. Note: In the following, A1 to An are optional batch dimensions. Args: sphere_center: A tensor of shape `[3]` representing the 3d sphere center. sphere_radius: A tensor of shape `[1]` containing a strictly positive value defining the radius of the sphere. ray: A tensor of shape `[A1, ..., An, 3]` containing normalized 3D vectors. point_on_ray: A tensor of shape `[A1, ..., An, 3]`. name: A name for this op. The default value of None means "ray_intersection_ray_sphere". Returns: A tensor of shape `[2, A1, ..., An, 3]` containing the position of the intersections, and a tensor of shape `[2, A1, ..., An, 3]` the associated surface normals at that point. Both tensors contain NaNs when there is no intersections. The first dimension of the returned tensor provides access to the first and second intersections of the ray with the sphere. Raises: ValueError: if the shape of `sphere_center`, `sphere_radius`, `ray` or `point_on_ray` is not supported. tf.errors.InvalidArgumentError: If `ray` is not normalized. """ with tf.compat.v1.name_scope( name, "ray_intersection_ray_sphere", [sphere_center, sphere_radius, ray, point_on_ray]): sphere_center = tf.convert_to_tensor(value=sphere_center) sphere_radius = tf.convert_to_tensor(value=sphere_radius) ray = tf.convert_to_tensor(value=ray) point_on_ray = tf.convert_to_tensor(value=point_on_ray) shape.check_static( tensor=sphere_center, tensor_name="sphere_center", has_rank=1, has_dim_equals=(0, 3)) shape.check_static( tensor=sphere_radius, tensor_name="sphere_radius", has_rank=1, has_dim_equals=(0, 1)) shape.check_static(tensor=ray, tensor_name="ray", has_dim_equals=(-1, 3)) shape.check_static( tensor=point_on_ray, tensor_name="point_on_ray", has_dim_equals=(-1, 3)) shape.compare_batch_dimensions( tensors=(ray, point_on_ray), last_axes=(-2, -2), broadcast_compatible=False) sphere_radius = asserts.assert_all_above( sphere_radius, 0.0, open_bound=True) ray = asserts.assert_normalized(ray) vector_sphere_center_to_point_on_ray = sphere_center - point_on_ray distance_sphere_center_to_point_on_ray = tf.norm( tensor=vector_sphere_center_to_point_on_ray, axis=-1, keepdims=True) distance_projection_sphere_center_on_ray = vector.dot( vector_sphere_center_to_point_on_ray, ray) closest_distance_sphere_center_to_ray = tf.sqrt( tf.square(distance_sphere_center_to_point_on_ray) - tf.pow(distance_projection_sphere_center_on_ray, 2)) half_secant_length = tf.sqrt( tf.square(sphere_radius) - tf.square(closest_distance_sphere_center_to_ray)) distances = tf.stack( (distance_projection_sphere_center_on_ray - half_secant_length, distance_projection_sphere_center_on_ray + half_secant_length), axis=0) intersections_points = distances * ray + point_on_ray normals = tf.math.l2_normalize( intersections_points - sphere_center, axis=-1) return intersections_points, normals
def func(u, v): return tf.squeeze(vector.dot(u, v), axis=-1)