def test_initialize_framebuffer_with_wrong_shapes(self): with self.assertRaisesRegex(tf.errors.InvalidArgumentError, "Expected all input shapes to be the same"): fb.Framebuffer( fb.RasterizedAttribute(tf.ones([1, 1, 4, 4, 3])), tf.ones([1, 1, 4, 4, 1]), tf.ones([1, 1, 4, 4, 3]), tf.ones([1, 1, 4, 4, 1]), {"an_attr": fb.RasterizedAttribute(tf.ones([1, 1, 4, 3, 4]))})
def test_initialize_framebuffer_with_wrong_rank(self): with self.assertRaisesRegex(ValueError, "Expected all inputs to have the same rank"): fb.Framebuffer( fb.RasterizedAttribute(tf.ones([1, 1, 4, 4, 1])), tf.ones([4, 3]), tf.ones([3, 4, 4, 5, 5]), tf.ones([3, 4, 4, 5, 5]))
def differentiable_barycentrics(framebuffer: fb.Framebuffer, clip_space_vertices: tf.Tensor, triangles: tf.Tensor) -> fb.Framebuffer: """Computes differentiable barycentric coordinates from a Framebuffer. The barycentric coordinates will be differentiable w.r.t. the input vertices. Later, we may support derivatives w.r.t. pixel position for mip-mapping. Args: framebuffer: a multi-layer Framebuffer containing triangle ids and a foreground mask with shape [batch, num_layers, height, width, 1] clip_space_vertices: a 2-D float32 tensor with shape [vertex_count, 4] or a 3-D tensor with shape [batch, vertex_count, 4] containing homogenous vertex positions (xyzw). triangles: a 2-D int32 tensor with shape [triangle_count, 3] or a 3-D tensor with shape [batch, triangle_count, 3] containing per-triangle vertex indices in counter-clockwise order. Returns: a copy of `framebuffer`, but the differentiable barycentric coordinates will replace any barycentric coordinates already in the `framebuffer`. """ rank = lambda t: len(t.shape) clip_space_vertices = tf.convert_to_tensor(clip_space_vertices) shape.check_static( tensor=clip_space_vertices, tensor_name="clip_space_vertices", has_rank_greater_than=1, has_rank_less_than=4) if rank(clip_space_vertices) == 2: clip_space_vertices = tf.expand_dims(clip_space_vertices, axis=0) triangles = tf.convert_to_tensor(triangles) shape.check_static( tensor=triangles, tensor_name="triangles", has_rank_greater_than=1, has_rank_less_than=4) if rank(triangles) == 2: triangles = tf.expand_dims(triangles, axis=0) shape.compare_batch_dimensions( tensors=(clip_space_vertices, triangles, framebuffer.triangle_id), last_axes=(-3, -3, -4), broadcast_compatible=False) # Compute image pixel coordinates. px, py = normalized_pixel_coordinates(framebuffer.width, framebuffer.height) def compute_barycentrics_fn( slices: Tuple[tf.Tensor, tf.Tensor, tf.Tensor]) -> tf.Tensor: clip_vertices_slice, triangle_slice, triangle_id_slice = slices triangle_id_slice = triangle_id_slice[..., 0] if rank(triangle_id_slice) == 2: # There is no layer dimension. triangle_id_slice = tf.expand_dims(triangle_id_slice, axis=0) # Compute per-triangle inverse matrices. triangle_matrices = compute_triangle_matrices(clip_vertices_slice, triangle_slice) # Compute per-pixel barycentric coordinates. barycentric_coords = compute_barycentric_coordinates( triangle_id_slice, triangle_matrices, px, py) barycentric_coords = tf.transpose(barycentric_coords, perm=[1, 2, 3, 0]) return barycentric_coords per_image_barycentrics = tf.vectorized_map( compute_barycentrics_fn, (clip_space_vertices, triangles, framebuffer.triangle_id)) barycentric_coords = tf.stack(per_image_barycentrics, axis=0) # After stacking barycentrics will have layers dimension no matter what. # In order to make sure we return differentiable barycentrics of the same # shape - reshape the tensor using original shape. barycentric_coords = tf.reshape( barycentric_coords, shape=framebuffer.barycentrics.value.shape) # Mask out barycentrics for background pixels. barycentric_coords = barycentric_coords * framebuffer.foreground_mask return fb.Framebuffer( triangle_id=framebuffer.triangle_id, vertex_ids=framebuffer.vertex_ids, foreground_mask=framebuffer.foreground_mask, attributes=framebuffer.attributes, barycentrics=fb.RasterizedAttribute(barycentric_coords, None, None))
def rasterize( vertices: type_alias.TensorLike, triangles: type_alias.TensorLike, view_projection_matrices: type_alias.TensorLike, image_size: Tuple[int, int], enable_cull_face: bool, num_layers: int, name: str = "rasterization_backend_cpu_rasterize") -> fb.Framebuffer: """Rasterizes the scene. This rasterizer estimates which triangle is associated with each pixel using the C++ software rasterizer. Args: vertices: A tensor of shape `[batch, num_vertices, 3]` containing batches of vertices, each defined by a 3D point. triangles: A tensor of shape `[num_triangles, 3]` containing triangles, each associated with 3 vertices from `scene_vertices`. view_projection_matrices: A tensor of shape `[batch, 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. enable_cull_face: A boolean, which will enable BACK face culling when True and no face culling when False. num_layers: Number of depth layers to render. name: A name for this op. Defaults to "rasterization_backend_cpu_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.name_scope(name): 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=3, 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=3, 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) if not num_layers > 0: raise ValueError("num_layers must be > 0.") vertices = utils.transform_homogeneous(view_projection_matrices, vertices) batch_size = tf.compat.dimension_value(vertices.shape[0]) per_image_barycentrics = [] per_image_triangle_ids = [] per_image_masks = [] image_width, image_height = image_size face_culling_mode = FaceCullingMode.BACK if enable_cull_face else FaceCullingMode.NONE for batch_index in range(batch_size): clip_vertices_slice = vertices[batch_index, ...] barycentrics, triangle_ids, z_buffer = ( render_ops.rasterize_triangles(clip_vertices_slice, triangles, image_width, image_height, num_layers, face_culling_mode)) # Shape [num_layers, image_height, image_width, 3] barycentrics = tf.stop_gradient(barycentrics) barycentrics = tf.ensure_shape( barycentrics, [num_layers, image_height, image_width, 3]) triangle_ids = tf.ensure_shape( triangle_ids, [num_layers, image_height, image_width]) z_buffer = tf.ensure_shape(z_buffer, [num_layers, image_height, image_width]) mask = tf.cast(tf.not_equal(z_buffer, 1.0), tf.float32) per_image_barycentrics.append(barycentrics) per_image_triangle_ids.append(triangle_ids) per_image_masks.append(mask) # Shape: [batch_size, num_layers, image_height, image_width, 1] triangle_id = tf.expand_dims(tf.stack(per_image_triangle_ids, axis=0), axis=-1) # Shape: [batch_size, num_layers, image_height, image_width, 3] vertex_ids = tf.gather(triangles, triangle_id[..., 0]) # Shape: [batch_size, num_layers, image_height, image_width, 1] mask = tf.expand_dims(tf.stack(per_image_masks, axis=0), axis=-1) # Shape: [batch_size, num_layers, image_height, image_width, 3] barycentrics = tf.stack(per_image_barycentrics, axis=0) return fb.Framebuffer(foreground_mask=mask, triangle_id=triangle_id, vertex_ids=vertex_ids, barycentrics=fb.RasterizedAttribute( value=barycentrics, d_dx=None, d_dy=None))
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))
def rasterize(vertices, triangles, view_projection_matrices, image_size, enable_cull_face, num_layers, 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 `[batch, num_vertices, 3]` containing batches vertices, each defined by a 3D point. triangles: A tensor of shape `[num_triangles, 3]` each associated with 3 vertices from `scene_vertices` view_projection_matrices: A tensor of shape `[batch, 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. enable_cull_face: A boolean, which will enable BACK face culling when True and no face culling when False. Default is True. num_layers: Number of depth layers to render. Not supported by current backend yet, but exists for interface compatibility. 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)): if num_layers != 1: raise ValueError("OpenGL rasterizer only supports single layer.") 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=3, 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=3, 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) geometry = tf.gather(vertices, triangles, axis=-2) # Extract batch size in order to make sure it is preserved after `gather` # operation. batch_size = _dim_value(vertices.shape[0]) rasterized = render_ops.rasterize( num_points=geometry.shape[-3], alpha_clear=0.0, enable_cull_face=enable_cull_face, variable_names=("view_projection_matrix", "triangular_mesh"), variable_kinds=("mat", "buffer"), variable_values=(view_projection_matrices, tf.reshape(geometry, shape=[batch_size, -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, [batch_size, 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 = rasterized[..., 3] mask = tf.reshape(mask, [batch_size, height, width, 1]) barycentric_coordinates = mask * barycentric_coordinates vertex_ids = tf.gather(triangles, triangle_index[..., 0], batch_dims=0) # Stop gradient for tensors coming out of custom op in order to avoid # confusing Tensorflow that they are differentiable. barycentric_coordinates = tf.stop_gradient(barycentric_coordinates) mask = tf.stop_gradient(mask) 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))