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 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 test_get_broadcasted_shape(self, shape_x, shape_y, broadcasted_shape): """Checks if the get_broadcasted_shape function works as expected.""" if tf.executing_eagerly(): if (shape_x is None or shape_y is None or None in shape_x or None in shape_y): return shape_x = tf.compat.v1.placeholder_with_default( tf.zeros(shape_x, dtype=tf.float32), shape=shape_x).shape shape_y = tf.compat.v1.placeholder_with_default( tf.zeros(shape_y, dtype=tf.float32), shape=shape_y).shape else: shape_x = tf.compat.v1.placeholder(shape=shape_x, dtype=tf.float32).shape shape_y = tf.compat.v1.placeholder(shape=shape_y, dtype=tf.float32).shape self.assertAllEqual( shape.get_broadcasted_shape(shape_x, shape_y), broadcasted_shape)
def rasterize(vertices, triangles, view_projection_matrices, image_size, name=None): """Rasterizes the scene. This rasterizer estimates which triangle is associated with each pixel using OpenGL. Note: In the following, A1 to An are optional batch dimensions which must be broadcast compatible for inputs `vertices` and `view_projection_matrices`. Args: vertices: A tensor of shape `[A1, ..., An, V, 3]` containing batches of `V` vertices, each defined by a 3D point. triangles: A tensor of shape `[T, 3]` containing `T` triangles, each associated with 3 vertices from `scene_vertices` view_projection_matrices: A tensor of shape `[A1, ..., An, 4, 4]` containing batches of view projection matrices image_size: An tuple of integers (height, width) containing the dimensions in pixels of the rasterized image. name: A name for this op. Defaults to 'rasterization_backend_rasterize'. Returns: A tuple of 3 elements. The first one of shape `[A1, ..., An, H, W, 1]` representing the triangle index associated with each pixel. If no triangle is associated to a pixel, the index is set to -1. The second element in the tuple is of shape `[A1, ..., An, H, W, 3]` and correspond to barycentric coordinates per pixel. The last element in the tuple is of shape `[A1, ..., An, H, W, 1]` and stores a value of `0` of the pixel is assciated with the background, and `1` with the foreground """ with tf.compat.v1.name_scope(name, "rasterization_backend_rasterize", (vertices, triangles, view_projection_matrices)): vertices = tf.convert_to_tensor(value=vertices) triangles = tf.convert_to_tensor(value=triangles) view_projection_matrices = tf.convert_to_tensor( value=view_projection_matrices) shape.check_static( tensor=vertices, tensor_name="vertices", has_rank_greater_than=1, has_dim_equals=((-1, 3))) shape.check_static( tensor=triangles, tensor_name="triangles", has_rank=2, has_dim_equals=((-1, 3))) shape.check_static( tensor=view_projection_matrices, tensor_name="view_projection_matrices", has_rank_greater_than=1, has_dim_equals=((-1, 4), (-2, 4))) shape.compare_batch_dimensions( tensors=(vertices, view_projection_matrices), tensor_names=("vertices", "view_projection_matrices"), last_axes=(-3, -3), broadcast_compatible=True) common_batch_shape = shape.get_broadcasted_shape( vertices.shape[:-2], view_projection_matrices.shape[:-2]) common_batch_shape = [_dim_value(dim) for dim in common_batch_shape] vertices = tf.broadcast_to(vertices, common_batch_shape + vertices.shape[-2:]) view_projection_matrices = tf.broadcast_to(view_projection_matrices, common_batch_shape + [4, 4]) geometry = tf.gather(vertices, triangles, axis=-2) rasterized = render_ops.rasterize( num_points=geometry.shape[-3], variable_names=("view_projection_matrix", "triangular_mesh"), variable_kinds=("mat", "buffer"), variable_values=(view_projection_matrices, tf.reshape(geometry, shape=common_batch_shape + [-1])), output_resolution=image_size, vertex_shader=vertex_shader, geometry_shader=geometry_shader, fragment_shader=fragment_shader) triangle_index = tf.cast(rasterized[..., 0] - 1, tf.int32) barycentric_coordinates = rasterized[..., 1:3] barycentric_coordinates = tf.concat( (barycentric_coordinates, 1.0 - barycentric_coordinates[..., 0:1] - barycentric_coordinates[..., 1:2]), axis=-1) mask = tf.cast(triangle_index >= 0, tf.int32) return triangle_index, barycentric_coordinates, mask
def blend(points, skinning_weights, bone_rotations, bone_translations, name=None): """Transforms the points using Linear Blend Skinning. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible and allow transforming full 3D shapes at once. In the following, B1 to Bm are optional batch dimensions, which allow transforming multiple poses at once. Args: points: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a 3d point. skinning_weights: A tensor of shape `[A1, ..., An, W]`, where the last dimension represents the skinning weights of each bone. bone_rotations: A tensor of shape `[B1, ..., Bm, W, 3, 3]`, which represents the 3d rotations applied to each bone. bone_translations: A tensor of shape `[B1, ..., Bm, W, 3]`, which represents the 3d translation vectors applied to each bone. name: A name for this op that defaults to "linear_blend_skinning_blend". Returns: A tensor of shape `[B1, ..., Bm, A1, ..., An, 3]`, where the last dimension represents a 3d point. Raises: ValueError: If the shape of the input tensors are not supported. """ with tf.compat.v1.name_scope( name, "linear_blend_skinning_blend", [points, skinning_weights, bone_rotations, bone_translations]): points = tf.convert_to_tensor(value=points) skinning_weights = tf.convert_to_tensor(value=skinning_weights) bone_rotations = tf.convert_to_tensor(value=bone_rotations) bone_translations = tf.convert_to_tensor(value=bone_translations) shape.check_static( tensor=points, tensor_name="points", has_dim_equals=(-1, 3)) shape.check_static( tensor=bone_rotations, tensor_name="bone_rotations", has_rank_greater_than=2, has_dim_equals=((-2, 3), (-1, 3))) shape.check_static( tensor=bone_translations, tensor_name="bone_translations", has_rank_greater_than=1, has_dim_equals=(-1, 3)) shape.compare_dimensions( tensors=(skinning_weights, bone_rotations), tensor_names=("skinning_weights", "bone_rotations"), axes=(-1, -3)) shape.compare_dimensions( tensors=(skinning_weights, bone_translations), tensor_names=("skinning_weights", "bone_translations"), axes=(-1, -2)) shape.compare_batch_dimensions( tensors=(points, skinning_weights), tensor_names=("points", "skinning_weights"), last_axes=(-2, -2), broadcast_compatible=True) shape.compare_batch_dimensions( tensors=(bone_rotations, bone_translations), tensor_names=("bone_rotations", "bone_translations"), last_axes=(-3, -2), broadcast_compatible=True) num_bones = skinning_weights.shape[-1] def dim_value(dim): return 1 if dim is None else tf.compat.v1.dimension_value(dim) # TODO(b/148362025): factorize this block out points_batch_shape = shape.get_broadcasted_shape( points.shape[:-1], skinning_weights.shape[:-1]) points_batch_shape = [dim_value(dim) for dim in points_batch_shape] points = tf.broadcast_to(points, points_batch_shape + [3]) skinning_weights = tf.broadcast_to(skinning_weights, points_batch_shape + [num_bones]) bones_batch_shape = shape.get_broadcasted_shape( bone_rotations.shape[:-3], bone_translations.shape[:-2]) bones_batch_shape = [dim_value(dim) for dim in bones_batch_shape] bone_rotations = tf.broadcast_to(bone_rotations, bones_batch_shape + [num_bones, 3, 3]) bone_translations = tf.broadcast_to(bone_translations, bones_batch_shape + [num_bones, 3]) points_batch_dims = points.shape.ndims - 1 bones_batch_dims = bone_rotations.shape.ndims - 3 points = tf.reshape(points, [1] * bones_batch_dims + points_batch_shape + [1, 3]) skinning_weights = tf.reshape(skinning_weights, [1] * bones_batch_dims + points_batch_shape + [num_bones, 1]) bone_rotations = tf.reshape( bone_rotations, bones_batch_shape + [1] * points_batch_dims + [num_bones, 3, 3]) bone_translations = tf.reshape( bone_translations, bones_batch_shape + [1] * points_batch_dims + [num_bones, 3]) transformed_points = rotation_matrix_3d.rotate( points, bone_rotations) + bone_translations weighted_points = tf.multiply(skinning_weights, transformed_points) blended_points = tf.reduce_sum(input_tensor=weighted_points, axis=-2) return blended_points
def 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 vertex_normals(vertices, indices, clockwise=True, name=None): """Computes vertex normals from a mesh. This function computes vertex normals as the weighted sum of the adjacent face normals, where the weights correspond to the area of each face. This function supports planar convex polygon faces. For non-triangular meshes, this function converts them into triangular meshes to calculate vertex normals. Note: In the following, A1 to An are optional batch dimensions. Args: vertices: A tensor of shape `[A1, ..., An, V, 3]`, where V is the number of vertices. 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. clockwise: Winding order to determine front-facing faces. The order of vertices should be either clockwise or counterclockwise. name: A name for this op. Defaults to "normals_vertex_normals". Returns: A tensor of shape `[A1, ..., An, V, 3]` containing vertex normals. If vertices and indices have different batch dimensions, this function broadcasts them into the same batch dimensions and the output batch dimensions are the broadcasted. Raises: ValueError: If the shape of `vertices`, `indices` is not supported. """ with tf.compat.v1.name_scope(name, "normals_vertex_normals", [vertices, indices]): 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, has_dim_equals=(-1, 3)) shape.check_static(tensor=indices, tensor_name="indices", has_rank_greater_than=1, has_dim_greater_than=(-1, 2)) shape.compare_batch_dimensions(tensors=(vertices, indices), last_axes=(-3, -3), broadcast_compatible=True) shape_indices = indices.shape.as_list() if None in shape_indices[:-2]: raise ValueError("'indices' must have specified batch dimensions.") common_batch_dims = shape.get_broadcasted_shape( vertices.shape[:-2], indices.shape[:-2]) vertices_repeat = [ common_batch_dims[x] // vertices.shape.as_list()[x] for x in range(len(common_batch_dims)) ] indices_repeat = [ common_batch_dims[x] // shape_indices[x] for x in range(len(common_batch_dims)) ] vertices = tf.tile(vertices, vertices_repeat + [1, 1], name="vertices_broadcast") indices = tf.tile(indices, indices_repeat + [1, 1], name="indices_broadcast") # Triangulate non-triangular faces. if shape_indices[-1] > 3: triangle_indices = [] for i in range(1, shape_indices[-1] - 1): triangle_indices.append( tf.concat((indices[..., 0:1], indices[..., i:i + 2]), axis=-1)) indices = tf.concat(triangle_indices, axis=-2) shape_indices = indices.shape.as_list() face_vertices = gather_faces(vertices, indices) # Use unnormalized face normals to scale normals by area. mesh_face_normals = face_normals(face_vertices, clockwise=clockwise, normalize=False) if vertices.shape.ndims > 2: outer_indices = np.meshgrid( *[np.arange(i) for i in shape_indices[:-2]], sparse=False, indexing="ij") outer_indices = [np.expand_dims(i, axis=-1) for i in outer_indices] outer_indices = np.concatenate(outer_indices, axis=-1) outer_indices = np.expand_dims(outer_indices, axis=-2) outer_indices = tf.constant(outer_indices, dtype=tf.int32) outer_indices = tf.tile(outer_indices, [1] * len(shape_indices[:-2]) + [tf.shape(input=indices)[-2]] + [1]) unnormalized_vertex_normals = tf.zeros_like(vertices) for i in range(shape_indices[-1]): scatter_indices = tf.concat( [outer_indices, indices[..., i:i + 1]], axis=-1) unnormalized_vertex_normals = tf.compat.v1.tensor_scatter_add( unnormalized_vertex_normals, scatter_indices, mesh_face_normals) else: unnormalized_vertex_normals = tf.zeros_like(vertices) for i in range(shape_indices[-1]): unnormalized_vertex_normals = tf.compat.v1.tensor_scatter_add( unnormalized_vertex_normals, indices[..., i:i + 1], mesh_face_normals) vector_norms = tf.sqrt( tf.reduce_sum(input_tensor=unnormalized_vertex_normals**2, axis=-1, keepdims=True)) return safe_ops.safe_unsigned_div(unnormalized_vertex_normals, vector_norms)
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 matrix_from_intrinsics(focal, principal_point, skew=(0.0, ), name="perspective_matrix_from_intrinsics"): r"""Builds calibration matrix from intrinsic parameters. Builds the camera calibration matrix as $$ \mathbf{C} = \begin{bmatrix} f_x & sc & c_x \\ 0 & f_y & c_y \\ 0 & 0 & 1 \\ \end{bmatrix} $$ from the focal length \\((f_x, f_y)\\) and the principal point \\((c_x, c_y)\\). Note: In the following, A1 to An are optional batch dimensions. Args: 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. skew: A tensor of shape `[A1, ..., An, 1]`, where the last dimension represents a skew coefficient. name: A name for this op that defaults to "perspective_matrix_from_intrinsics". Returns: A tensor of shape `[A1, ..., An, 3, 3]`, where the last two dimensions represent a camera calibration matrix. Raises: ValueError: If the shape of `focal`, or `principal_point` is not supported. """ with tf.name_scope(name): focal = tf.convert_to_tensor(value=focal) principal_point = tf.convert_to_tensor(value=principal_point) skew = tf.convert_to_tensor(value=skew) common_batch_shape = shape.get_broadcasted_shape( focal.shape[:-1], skew.shape[:-1]) 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] skew = tf.broadcast_to(skew, common_batch_shape + [1]) 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.check_static( tensor=skew, tensor_name="skew", has_dim_equals=(-1, 1), ) shape.compare_batch_dimensions(tensors=(focal, principal_point, skew), tensor_names=("focal", "principal_point", "skew"), last_axes=-2, broadcast_compatible=False) fx, fy = tf.unstack(focal, axis=-1) cx, cy = tf.unstack(principal_point, axis=-1) zero = tf.zeros_like(fx) one = tf.ones_like(fx) skew = tf.reshape(skew, tf.shape(fx)) matrix = tf.stack((fx, skew, cx, zero, fy, cy, zero, zero, one), axis=-1) # pyformat: disable matrix_shape = tf.shape(input=matrix) output_shape = tf.concat((matrix_shape[:-1], (3, 3)), axis=-1) return tf.reshape(matrix, shape=output_shape)
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 rasterize(vertices, triangles, view_projection_matrices, image_size, name=None): """Rasterizes the scene. This rasterizer estimates which triangle is associated with each pixel using OpenGL. Note: In the following, A1 to An are optional batch dimensions which must be broadcast compatible for inputs `vertices` and `view_projection_matrices`. Args: vertices: A tensor of shape `[A1, ..., An, V, 3]` containing batches of `V` vertices, each defined by a 3D point. triangles: A tensor of shape `[T, 3]` containing `T` triangles, each associated with 3 vertices from `scene_vertices` view_projection_matrices: A tensor of shape `[A1, ..., An, 4, 4]` containing batches of view projection matrices image_size: An tuple of integers (width, height) containing the dimensions in pixels of the rasterized image. name: A name for this op. Defaults to 'rasterization_backend_rasterize'. Returns: A Framebuffer containing the rasterized values: barycentrics, triangle_id, foreground_mask, vertex_ids. Returned Tensors have shape [batch, num_layers, height, width, channels] Note: triangle_id contains the triangle id value for each pixel in the output image. For pixels within the mesh, this is the integer value in the range [0, num_vertices] from triangles. For vertices outside the mesh this is 0; 0 can either indicate belonging to triangle 0, or being outside the mesh. This ensures all returned triangle ids will validly index into the vertex array, enabling the use of tf.gather with indices from this tensor. The barycentric coordinates can be used to determine pixel validity instead. See framebuffer.py for a description of the Framebuffer fields. """ with tf.compat.v1.name_scope( name, "rasterization_backend_rasterize", (vertices, triangles, view_projection_matrices)): vertices = tf.convert_to_tensor(value=vertices) triangles = tf.convert_to_tensor(value=triangles) view_projection_matrices = tf.convert_to_tensor( value=view_projection_matrices) shape.check_static(tensor=vertices, tensor_name="vertices", has_rank_greater_than=1, has_dim_equals=((-1, 3))) shape.check_static(tensor=triangles, tensor_name="triangles", has_rank=2, has_dim_equals=((-1, 3))) shape.check_static(tensor=view_projection_matrices, tensor_name="view_projection_matrices", has_rank_greater_than=1, has_dim_equals=((-1, 4), (-2, 4))) shape.compare_batch_dimensions( tensors=(vertices, view_projection_matrices), tensor_names=("vertices", "view_projection_matrices"), last_axes=(-3, -3), broadcast_compatible=True) common_batch_shape = shape.get_broadcasted_shape( vertices.shape[:-2], view_projection_matrices.shape[:-2]) common_batch_shape = [_dim_value(dim) for dim in common_batch_shape] vertices = tf.broadcast_to(vertices, common_batch_shape + vertices.shape[-2:]) view_projection_matrices = tf.broadcast_to(view_projection_matrices, common_batch_shape + [4, 4]) geometry = tf.gather(vertices, triangles, axis=-2) rasterized = render_ops.rasterize( num_points=geometry.shape[-3], alpha_clear=0.0, enable_cull_face=True, variable_names=("view_projection_matrix", "triangular_mesh"), variable_kinds=("mat", "buffer"), variable_values=(view_projection_matrices, tf.reshape(geometry, shape=common_batch_shape + [-1])), output_resolution=image_size, vertex_shader=vertex_shader, geometry_shader=geometry_shader, fragment_shader=fragment_shader) triangle_index = tf.cast(rasterized[..., 0], tf.int32) # Slicing of the tensor will result in all batch dimensions being # `None` for tensorflow graph mode, therefore we have to fix it in order to # have explicit shape. width, height = image_size triangle_index = tf.reshape(triangle_index, common_batch_shape + [height, width, 1]) barycentric_coordinates = rasterized[..., 1:3] barycentric_coordinates = tf.concat( (barycentric_coordinates, 1.0 - barycentric_coordinates[..., 0:1] - barycentric_coordinates[..., 1:2]), axis=-1) mask = tf.cast(rasterized[..., 3], tf.int32) mask = tf.reshape(mask, common_batch_shape + [height, width, 1]) triangles_batch = tf.broadcast_to(triangles, common_batch_shape + triangles.shape) vertex_ids = tf.gather(triangles_batch, triangle_index[..., 0], batch_dims=len(common_batch_shape)) return fb.Framebuffer(foreground_mask=mask, triangle_id=triangle_index, vertex_ids=vertex_ids, barycentrics=fb.RasterizedAttribute( value=barycentric_coordinates, d_dx=None, d_dy=None))