def test_rasterizer_rasterize_exception_raised(self, shapes, dtypes, backend): """Tests that unsupported backends raise exceptions.""" placeholders = self._create_placeholders(shapes, dtypes) with self.assertRaisesRegexp(KeyError, 'Backend is not supported'): rasterization_backend.rasterize(placeholders[0], placeholders[1], placeholders[2], (600, 800), self._enable_cull_face, self._num_layers, backend)
def test_rasterizer_rasterize_exception_not_raised(self, shapes, dtypes, enable_cull_face): """Tests that supported backends do not raise exceptions.""" placeholders = self._create_placeholders(shapes, dtypes) try: rasterization_backend.rasterize(placeholders[0], placeholders[1], placeholders[2], (600, 800), enable_cull_face, self._num_layers, self._backend) except Exception as e: # pylint: disable=broad-except self.fail('Exception raised: %s' % str(e))
def test_differentialbe_barycentrics_close_to_rasterizer( self, rasterization_backend_type): image_width = 32 image_height = 32 initial_vertices = tf.constant([[[0, 0, 0], [0.5, 0, 0], [0.5, 0.5, 0]]], dtype=tf.float32) triangles = [[0, 1, 2]] view_projection_matrix = tf.expand_dims(tf.eye(4), axis=0) # Compute rasterization barycentrics rasterized = rasterization_backend.rasterize( initial_vertices, triangles, view_projection_matrix, (image_width, image_height), num_layers=1, backend=rasterization_backend_type).layer(0) ras_barycentrics = rasterized.barycentrics.value clip_space_vertices = utils.transform_homogeneous(view_projection_matrix, initial_vertices) rasterized_w_diff_barycentrics = barycentrics.differentiable_barycentrics( rasterized, clip_space_vertices, triangles) diff_barycentrics = rasterized_w_diff_barycentrics.barycentrics.value self.assertAllClose(ras_barycentrics, diff_barycentrics, rtol=1e-4)
def test_rasterizer_return_correct_batch_shapes(self, shapes, dtypes, enable_cull_face): """Tests that supported backends return correct shape.""" placeholders = self._create_placeholders(shapes, dtypes) frame_buffer = rasterization_backend.rasterize( placeholders[0], placeholders[1], placeholders[2], (600, 800), enable_cull_face, self._num_layers, self._backend) batch_size = shapes[0][0] self.assertEqual([batch_size], frame_buffer.triangle_id.get_shape().as_list()[:-3]) self.assertEqual( [batch_size], frame_buffer.foreground_mask.get_shape().as_list()[:-3])
def test_rasterizer_all_vertices_visible(self): """Renders simple triangle and asserts that it is fully visible.""" vertices = tf.convert_to_tensor([[[0, 0, 0], [10, 10, 0], [0, 10, 0]]], dtype=tf.float32) triangles = tf.convert_to_tensor([[0, 1, 2]], dtype=tf.int32) view_projection_matrix = tf.expand_dims(tf.eye(4), axis=0) frame_buffer = rasterization_backend.rasterize( vertices, triangles, view_projection_matrix, (100, 100), self._enable_cull_face, self._num_layers, self._backend) self.assertAllEqual(frame_buffer.triangle_id.shape[:-1], frame_buffer.vertex_ids.shape[:-1]) # Assert that triangle is visible. self.assertAllLess(frame_buffer.triangle_id, 2) self.assertAllGreaterEqual(frame_buffer.triangle_id, 0) # Assert that all three vertices are visible. self.assertAllLess(frame_buffer.triangle_id, 3) self.assertAllGreaterEqual(frame_buffer.triangle_id, 0)
def test_renders_colored_cube_multilayer(self, num_layers): """Renders a simple colored cube with multiple layers.""" rasterized = rasterization_backend.rasterize( self.cube_vertex_positions, self.cube_triangles, self.projection, (self.image_width, self.image_height), num_layers=num_layers, enable_cull_face=False, backend=rasterization_backend.RasterizationBackends.CPU) vertex_rgb = (self.cube_vertex_positions * 0.5 + 0.5) vertex_rgba = tf.concat([vertex_rgb, tf.ones([1, 8, 1])], axis=-1) rendered = interpolate.interpolate_vertex_attribute( vertex_rgba, rasterized).value image_shape = [1] + rendered.shape[2:] for layer_index in range(num_layers): baseline_image = rasterization_test_utils.load_baseline_image( 'Unlit_Cube_{}_{}.png'.format(0, layer_index), image_shape) images_near, error_message = rasterization_test_utils.compare_images( self, baseline_image, rendered[:, layer_index, ...]) self.assertTrue(images_near, msg=error_message)
def test_renders_colored_cube(self): """Renders a simple colored cube.""" num_layers = 1 rasterized = rasterization_backend.rasterize( self.cube_vertex_positions, self.cube_triangles, self.projection, (self.image_width, self.image_height), num_layers=num_layers, enable_cull_face=False, backend=rasterization_backend.RasterizationBackends.CPU) vertex_rgb = (self.cube_vertex_positions * 0.5 + 0.5) vertex_rgba = tf.concat([vertex_rgb, tf.ones([1, 8, 1])], axis=-1) rendered = interpolate.interpolate_vertex_attribute(vertex_rgba, rasterized).value baseline_image = rasterization_test_utils.load_baseline_image( 'Unlit_Cube_0_0.png', rendered.shape) images_near, error_message = rasterization_test_utils.compare_images( self, baseline_image, rendered) self.assertTrue(images_near, msg=error_message)
def test_render_simple_triangle(self): """Directly renders a rasterized triangle's barycentric coordinates.""" w_vector = tf.constant([1.0, 1.0, 1.0], dtype=tf.float32) clip_init = tf.constant( [[[-0.5, -0.5, 0.8], [0.0, 0.5, 0.3], [0.5, -0.5, 0.3]]], dtype=tf.float32) clip_coordinates = clip_init * tf.reshape(w_vector, [1, 3, 1]) triangles = tf.constant([[0, 1, 2]], dtype=tf.int32) face_culling_enabled = False framebuffer = rasterization_backend.rasterize( clip_coordinates, triangles, tf.eye(4, batch_shape=[1]), (self.IMAGE_WIDTH, self.IMAGE_HEIGHT), face_culling_enabled, self._num_layers, self._backend).layer(0) ones_image = tf.ones([1, self.IMAGE_HEIGHT, self.IMAGE_WIDTH, 1]) rendered_coordinates = tf.concat( [framebuffer.barycentrics.value, ones_image], axis=-1) baseline_image = rasterization_test_utils.load_baseline_image( 'Simple_Triangle.png', rendered_coordinates.shape) images_near, error_message = rasterization_test_utils.compare_images( self, baseline_image, rendered_coordinates) self.assertTrue(images_near, msg=error_message)
def rasterize(vertices, triangles, attributes, model_to_eye_matrix, perspective_matrix, image_size, backend=rasterization_backend.RasterizationBackends.OPENGL, name=None): """Rasterizes the scene. Note: In the following, A1 to An are optional batch dimensions. 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 `vertices`. attributes: A dictionary of tensors, each of shape `[A1, ..., An, V, K_a]` containing batches of `V` vertices, each associated with K-dimensional attributes. K_a may vary by attribute. model_to_eye_matrix: A tensor of shape `[A1, ..., An, 4, 4]` containing batches of matrices used to transform vertices from model to eye coordinates. perspective_matrix: A tensor of shape `[A1, ..., An, 4, 4]` containing batches of matrices used to project vertices from eye to clip coordinates. image_size: A tuple (height, width) containing the dimensions in pixels of the rasterized image. backend: A rasterization_backend.RasterizationBackends enum containing the backend method to use for rasterization. name: A name for this op. Defaults to 'triangle_rasterizer_rasterize'. Returns: A dictionary. The key "mask" is of shape `[A1, ..., An, height, width, 1]` and stores a value of `0` of the pixel is assciated with the background, and `1` with the foreground. The key "barycentrics" is of shape `[A1, ..., An, height, width, 3]` and stores barycentric weights. Finally, the dictionary contains perspective correct interpolated attributes of shape `[A1, ..., An, height, width, K]` per entry in the `attributes` dictionary. """ with tf.compat.v1.name_scope(name, "triangle_rasterizer_rasterize", (vertices, triangles, attributes, model_to_eye_matrix, perspective_matrix)): vertices = tf.convert_to_tensor(value=vertices) triangles = tf.convert_to_tensor(value=triangles) model_to_eye_matrix = tf.convert_to_tensor(value=model_to_eye_matrix) perspective_matrix = tf.convert_to_tensor(value=perspective_matrix) 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=model_to_eye_matrix, tensor_name="model_to_eye_matrix", has_dim_equals=(((-2, 4), (-1, 4)))) shape.check_static( tensor=perspective_matrix, tensor_name="perspective_matrix", has_dim_equals=(((-2, 4), (-1, 4)))) image_size_float = (float(image_size[0]), float(image_size[1])) image_size_backend = (int(image_size[1]), int(image_size[0])) view_projection_matrix = tf.linalg.matmul(perspective_matrix, model_to_eye_matrix) rasterized = rasterization_backend.rasterize(vertices, triangles, view_projection_matrix, image_size_backend, backend) outputs = { "mask": rasterized.foreground_mask, "triangle_indices": rasterized.triangle_id } # Extract batch shape in order to make sure it is preserved after `gather` # operation. batch_shape = rasterized.triangle_id.shape[:-3] batch_shape = [_dim_value(dim) for dim in batch_shape] vertices_per_pixel = tf.gather( vertices, rasterized.vertex_ids, batch_dims=len(batch_shape)) barycentrics = _perspective_correct_barycentrics(vertices_per_pixel, model_to_eye_matrix, perspective_matrix, image_size_float) mask_float = tf.cast(rasterized.foreground_mask, vertices.dtype) outputs["barycentrics"] = mask_float * barycentrics for key, attribute in attributes.items(): attribute = tf.convert_to_tensor(value=attribute) outputs[key] = mask_float * _perspective_correct_attributes( attribute, barycentrics, triangles, rasterized.triangle_id[..., 0], len(batch_shape)) return outputs
def rasterize(vertices, triangles, attributes, view_projection_matrix, image_size, backend=rasterization_backend.RasterizationBackends.OPENGL, name=None): """Rasterizes the scene. Note: In the following, A1 to An are optional batch dimensions. 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 `vertices`. attributes: A dictionary of tensors, each of shape `[A1, ..., An, V, K_a]` containing batches of `V` vertices, each associated with K-dimensional attributes. K_a may vary by attribute. view_projection_matrix: A tensor of shape `[A1, ..., An, 4, 4]` containing batches of matrices used to transform vertices from model to clip coordinates. image_size: A tuple (height, width) containing the dimensions in pixels of the rasterized image. backend: A rasterization_backend.RasterizationBackends enum containing the backend method to use for rasterization. name: A name for this op. Defaults to 'triangle_rasterizer_rasterize'. Returns: A dictionary. The key "mask" is of shape `[A1, ..., An, height, width, 1]` and stores a value of `0` of the pixel is assciated with the background, and `1` with the foreground. The key "barycentrics" is of shape `[A1, ..., An, height, width, 3]` and stores barycentric weights. Finally, the dictionary contains perspective correct interpolated attributes of shape `[A1, ..., An, height, width, K]` per entry in the `attributes` dictionary. """ with tf.compat.v1.name_scope( name, "triangle_rasterizer_rasterize", (vertices, triangles, attributes, view_projection_matrix)): vertices = tf.convert_to_tensor(value=vertices) triangles = tf.convert_to_tensor(value=triangles) view_projection_matrix = tf.convert_to_tensor( value=view_projection_matrix) 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_matrix, tensor_name="view_projection_matrix", has_dim_equals=(((-2, 4), (-1, 4)))) image_size_backend = (int(image_size[1]), int(image_size[0])) input_batch_shape = vertices.shape[:-2] view_projection_matrix = _merge_batch_dims(view_projection_matrix, last_axis=-2) vertices = _merge_batch_dims(vertices, last_axis=-2) rasterized = rasterization_backend.rasterize(vertices, triangles, view_projection_matrix, image_size_backend, backend=backend) outputs = { "mask": _restore_batch_dims(rasterized.foreground_mask, input_batch_shape), "triangle_indices": _restore_batch_dims(rasterized.triangle_id, input_batch_shape) } # Extract batch shape in order to make sure it is preserved after `gather` # operation. batch_shape = rasterized.triangle_id.shape[:-3] batch_shape = [_dim_value(dim) for dim in batch_shape] clip_space_vertices = utils.transform_homogeneous( view_projection_matrix, vertices) rasterized = barycentrics_module.differentiable_barycentrics( rasterized, clip_space_vertices, triangles) barycentrics = rasterized.barycentrics.value outputs["barycentrics"] = _restore_batch_dims( rasterized.foreground_mask * barycentrics, input_batch_shape) for key, attribute in attributes.items(): attribute = tf.convert_to_tensor(value=attribute) attribute = _merge_batch_dims(attribute, last_axis=-2) masked_attribute = interpolate.interpolate_vertex_attribute( attribute, rasterized) masked_attribute = _restore_batch_dims(masked_attribute.value, input_batch_shape) outputs[key] = masked_attribute return outputs
def rasterize(vertices, triangles, attributes, model_to_eye_matrix, perspective_matrix, image_size, background_attribute, backend=rasterization_backend.RasterizationBackends.OPENGL, name=None): """Rasterizes the scene. Note: In the following, A1 to An are optional batch dimensions. 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 `vertices`. attributes: A dictionary of tensors, each of shape `[A1, ..., An, V, K]` containing batches of `V` vertices, each associated with K-dimensional attributes. model_to_eye_matrix: A tensor of shape `[A1, ..., An, 4, 4]` containing batches of matrices used to transform vertices from model to eye coordinates. perspective_matrix: A tensor of shape `[A1, ..., An, 4, 4]` containing batches of matrices used to project vertices from eye to clip coordinates. image_size: A tuple (height, width) containing the dimensions in pixels of the rasterized image. background_attribute: A tensor of shape `[K]` containing the attribute to use for pixels associated with the background of the rendered scene. backend: A rasterization_backend.RasterizationBackends enum containing the backend method to use for rasterization. name: A name for this op. Defaults to 'triangle_rasterizer_rasterize'. Returns: A dictionary. The key "mask" is of shape `[A1, ..., An, height, width]` and stores a value of `0` of the pixel is assciated with the background, and `1` with the foreground. The key "barycentrics" is of shape `[A1, ..., An, height, width, 3]` and stores barycentric weights. Finally, the dictionary contains perspective correct interpolated attributes of shape `[A1, ..., An, height, width, K]` per entry in the `attributes` dictionary. """ with tf.compat.v1.name_scope( name, "triangle_rasterizer_rasterize", (vertices, triangles, attributes, model_to_eye_matrix, perspective_matrix, background_attribute)): vertices = tf.convert_to_tensor(value=vertices) triangles = tf.convert_to_tensor(value=triangles) model_to_eye_matrix = tf.convert_to_tensor(value=model_to_eye_matrix) perspective_matrix = tf.convert_to_tensor(value=perspective_matrix) background_attribute = tf.convert_to_tensor(value=background_attribute) 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=model_to_eye_matrix, tensor_name="model_to_eye_matrix", has_dim_equals=(((-2, 4), (-1, 4)))) shape.check_static( tensor=perspective_matrix, tensor_name="perspective_matrix", has_dim_equals=(((-2, 4), (-1, 4)))) shape.check_static( tensor=background_attribute, tensor_name="background_attribute", has_rank=1) image_size_float = (float(image_size[0]), float(image_size[1])) image_size_backend = (int(image_size[1]), int(image_size[0])) view_projection_matrix = tf.linalg.matmul(perspective_matrix, model_to_eye_matrix) triangle_index, _, mask = rasterization_backend.rasterize( vertices, triangles, view_projection_matrix, image_size_backend, backend) batch_shape = triangle_index.shape[:-3] batch_shape = [_dim_value(dim) for dim in batch_shape] vertices = tf.gather(vertices, triangles, axis=-2) # Gather does not work on negative indices, which is the case for the pixel # associated to the background. triangle_index = triangle_index * mask vertices_per_pixel = tf.gather( vertices, triangle_index, axis=-3, batch_dims=len(batch_shape)) barycentrics = _perspective_correct_barycentrics(vertices_per_pixel, model_to_eye_matrix, perspective_matrix, image_size_float) mask_float = tf.cast(tf.expand_dims(mask, axis=-1), vertices.dtype) outputs = {"mask": mask, "barycentrics": barycentrics} masked_background_attribute = (1.0 - mask_float) * background_attribute for key, attribute in attributes.items(): attribute = tf.convert_to_tensor(value=attribute) shape.check_static( tensor=attribute, tensor_name=key, has_dim_equals=(-1, background_attribute.shape[-1])) interpolated_attribute = _perspective_correct_attributes( attribute, barycentrics, triangles, triangle_index, len(batch_shape)) interpolated_attribute = ( mask_float * interpolated_attribute) + masked_background_attribute outputs[key] = interpolated_attribute return outputs
def rasterize( vertices: type_alias.TensorLike, triangles: type_alias.TensorLike, attributes: Dict[str, type_alias.TensorLike], view_projection_matrix: type_alias.TensorLike, image_size: Tuple[int, int], enable_cull_face: bool = True, use_vectorized_map: bool = True, backend: enum.Enum = rasterization_backend.RasterizationBackends.OPENGL, name: str = "triangle_rasterizer_rasterize" ) -> Dict[str, type_alias.TensorLike]: """Rasterizes the scene. Note: In the following, A1 to An are optional batch dimensions. 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 `vertices`. attributes: A dictionary of tensors, each of shape `[A1, ..., An, V, K]` containing batches of `V` vertices, each associated with K-dimensional attributes. K may vary by attribute. view_projection_matrix: A tensor of shape `[A1, ..., An, 4, 4]` containing batches of matrices used to transform vertices from model to clip coordinates. image_size: A tuple (height, width) containing the dimensions in pixels of the rasterized image. enable_cull_face: Enables BACK face culling when True, and no culling when False. use_vectorized_map: If true uses vectorized_map for barycentrics computations otherwise uses map_fn. backend: A rasterization_backend.RasterizationBackends enum containing the backend method to use for rasterization. name: A name for this op. Defaults to "triangle_rasterizer_rasterize". Returns: A dictionary. The key "mask" is of shape `[A1, ..., An, height, width, 1]` and stores a value of `0` of the pixel is assciated with the background, and `1` with the foreground. The key "barycentrics" is of shape `[A1, ..., An, height, width, 3]` and stores barycentric weights. Finally, the dictionary contains perspective correct interpolated attributes of shape `[A1, ..., An, height, width, K]` per entry in the `attributes` dictionary. """ with tf.name_scope(name): vertices = tf.convert_to_tensor(value=vertices) triangles = tf.convert_to_tensor(value=triangles) view_projection_matrix = tf.convert_to_tensor(value=view_projection_matrix) 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_matrix, tensor_name="view_projection_matrix", has_dim_equals=(((-2, 4), (-1, 4)))) image_size_backend = (int(image_size[1]), int(image_size[0])) input_batch_shape = vertices.shape[:-2] view_projection_matrix = utils.merge_batch_dims( view_projection_matrix, last_axis=-2) vertices = utils.merge_batch_dims(vertices, last_axis=-2) rasterized = rasterization_backend.rasterize( vertices, triangles, view_projection_matrix, image_size_backend, enable_cull_face=enable_cull_face, backend=backend) rasterized = rasterized.layer(0) outputs = { "mask": utils.restore_batch_dims(rasterized.foreground_mask, input_batch_shape), "triangle_indices": utils.restore_batch_dims(rasterized.triangle_id, input_batch_shape) } clip_space_vertices = utils.transform_homogeneous(view_projection_matrix, vertices) rasterized = barycentrics_module.differentiable_barycentrics( rasterized, clip_space_vertices, triangles, use_vectorized_map) barycentrics = rasterized.barycentrics.value outputs["barycentrics"] = utils.restore_batch_dims( rasterized.foreground_mask * barycentrics, input_batch_shape) for key, attribute in attributes.items(): attribute = tf.convert_to_tensor(value=attribute) attribute = utils.merge_batch_dims(attribute, last_axis=-2) masked_attribute = interpolate.interpolate_vertex_attribute( attribute, rasterized) masked_attribute = utils.restore_batch_dims(masked_attribute.value, input_batch_shape) outputs[key] = masked_attribute return outputs
def rasterize_then_splat( vertices: type_alias.TensorLike, triangles: type_alias.TensorLike, attributes: Dict[str, type_alias.TensorLike], view_projection_matrix: type_alias.TensorLike, image_size: Tuple[int, int], shading_function: Callable[[Dict[str, tf.Tensor]], tf.Tensor], num_layers=1, return_extra_buffers=False, backend: enum.Enum = rasterization_backend.RasterizationBackends.CPU, name='rasterize_then_splat'): """Rasterization with differentiable occlusion using rasterize-then-splat. Rasterizes the input triangles to produce surface point samples, applies a user-specified shading function, then splats the shaded point samples onto the pixel grid. The attributes are arbitrary per-vertex quantities (colors, normals, texture coordinates, etc.). The rasterization step interpolates these attributes across triangles to produce a dictionary of per-pixel interpolated attributes buffers with shapes `[H, W, K]` where `K` is the number of channels of the input attribute. This dictionary is passed to the user-provided `shading_function`, which performs shading and outputs a `[H, W, 4]` buffer of RGBA colors. The result of the shader is replaced with (0,0,0,0) for background pixels. In the common case that the attributes are RGBA vertex colors, the shading function may just pass the rasterized attributes through (i.e., `shading_function = lambda x: x['color']` where `color` is an RGBA attribute). 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 `vertices`. attributes: A dictionary of tensors, each of shape `[A1, ..., An, V, K]` containing batches of `V` vertices, each associated with `K`-dimensional attributes. `K` may vary by attribute. view_projection_matrix: A tensor of shape `[A1, ..., An, 4, 4]` containing batches of matrices used to transform vertices from model to clip coordinates. image_size: A tuple (height, width) containing the dimensions in pixels of the rasterized image. shading_function: a function that takes a dictionary of `[H, W, K]` rasterized attribute tensors and returns a `[H, W, 4]` RGBA tensor. num_layers: int specifying number of depth layers to composite. return_extra_buffers: if True, the function will return raw accumulation buffers for visualization. backend: A rasterization_backend.RasterizationBackends enum containing the backend method to use for rasterization. name: A name for this op. Defaults to "rasterize_then_splat". Returns: a `[A1, ..., An, H, W, 4]` tensor of RGBA values. """ with tf.name_scope(name): vertices = tf.convert_to_tensor(vertices) triangles = tf.convert_to_tensor(triangles) view_projection_matrix = tf.convert_to_tensor(view_projection_matrix) 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_matrix, tensor_name='view_projection_matrix', has_dim_equals=(((-2, 4), (-1, 4)))) input_batch_shape = vertices.shape[:-2] view_projection_matrix = utils.merge_batch_dims(view_projection_matrix, last_axis=-2) vertices = utils.merge_batch_dims(vertices, last_axis=-2) image_size_backend = (int(image_size[1]), int(image_size[0])) # We don't need derivatives of barycentric coordinates for RtS, so use # rasterization_backend directly. # Back face culling is necessary when rendering multiple layers so that # back faces aren't counted as occluding layers. rasterized = rasterization_backend.rasterize(vertices, triangles, view_projection_matrix, image_size_backend, enable_cull_face=True, num_layers=num_layers, backend=backend) # TODO(fcole): check if any of these keys already exist in attributes shader_dict = { 'mask': rasterized.foreground_mask, 'triangle_indices': rasterized.triangle_id, 'barycentrics': rasterized.barycentrics.value } for key, attribute in attributes.items(): attribute = tf.convert_to_tensor(value=attribute) attribute = utils.merge_batch_dims(attribute, last_axis=-2) interpolated = interpolate.interpolate_vertex_attribute( attribute, rasterized) shader_dict[key] = interpolated.value # Nested vectorized map over batch and layer dimensions. shaded_buffer = tf.vectorized_map( lambda l: tf.vectorized_map(shading_function, l), shader_dict) # Zero out shader result outside of foreground mask. shaded_buffer = shaded_buffer * rasterized.foreground_mask clip_space_vertices = utils.transform_homogeneous( view_projection_matrix, vertices) clip_space_buffer = interpolate.interpolate_vertex_attribute( clip_space_vertices, rasterized, (0, 0, 1, 1)).value ndc_xyz = clip_space_buffer[..., :3] / clip_space_buffer[..., 3:4] image_height, image_width = image_size viewport_xyz = (ndc_xyz + 1.0) * tf.constant( [image_width, image_height, 1], dtype=tf.float32, shape=[1, 1, 1, 1, 3]) * 0.5 output, accum, norm_accum = tf.vectorized_map( splat_at_pixel_centers, (viewport_xyz, shaded_buffer)) if return_extra_buffers: return output, accum, norm_accum output = utils.restore_batch_dims(output, input_batch_shape) return output