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))
