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_rasterized_attribute_with_wrong_shapes(self):
with self.assertRaisesRegex(
tf.errors.InvalidArgumentError,
"Expected all input shapes to be the same"):
fb.RasterizedAttribute(tf.ones([2, 4, 4, 1]), tf.ones([2, 4, 3,
1]),
tf.ones([2, 4, 3, 5]))
def test_initialize_rasterized_attribute_and_derivatives_with_wrong_rank(
self):
with self.assertRaisesRegex(
ValueError,
"Expected value and derivatives to be of the same rank"):
fb.RasterizedAttribute(tf.ones([4, 4, 1]), tf.ones([4, 3]),
tf.ones([3, 4, 4, 5, 5]))
def interpolate_vertex_attribute(
attribute: tf.Tensor,
framebuffer: fb.Framebuffer,
background_value: Optional[Union[tf.Tensor, Iterable[float]]] = None
) -> fb.RasterizedAttribute:
"""Interpolate a single vertex attribute across the input framebuffer.
Args:
attribute: 2-D or 3-D vertex attribute Tensor with shape [batch,
num_vertices, num_channels] or [num_vertices, num_channels].
framebuffer: Framebuffer to interpolate across. Expected to contain
barycentrics, vertex_ids, and foreground_mask.
background_value: 1-D Tensor (or convertible value) with shape
[num_channels] containing the value to use for background pixels. If None,
defaults to zero.
Returns:
A RasterizedAttribute containing the per-pixel interpolated values.
"""
vertex_ids = tf.reshape(framebuffer.vertex_ids,
[framebuffer.batch_size, framebuffer.pixel_count, 3])
num_channels = tf.compat.dimension_value(tf.shape(attribute)[-1])
# Creates indices with each pixel's clip-space triangle's extrema (the pixel's
# 'corner points') ids to look up the attributes for each pixel's triangle.
# Handles batched or unbatched attributes. In either case, corner_attribute
# will be shaped [batch_size, pixel_count, 3, num_channels] (with
# batch_size = 1 if input is unbatched).
if len(attribute.shape) == 3:
corner_attribute = tf.gather(attribute, vertex_ids, batch_dims=1)
else:
vertex_ids = tf.reshape(vertex_ids, (-1, 3))
corner_attribute = tf.gather(attribute, vertex_ids, batch_dims=0)
corner_attribute = tf.reshape(
corner_attribute,
(framebuffer.batch_size, framebuffer.pixel_count, 3, num_channels))
# Computes the pixel attributes by interpolating the known attributes at the
# corner points of the triangle interpolated with the barycentric
# coordinates.
reshaped_barycentrics = tf.reshape(
framebuffer.barycentrics.value,
(framebuffer.batch_size, framebuffer.pixel_count, 3, 1))
weighted_vertex_attribute = tf.multiply(corner_attribute,
reshaped_barycentrics)
summed_attribute = tf.reduce_sum(weighted_vertex_attribute, axis=2)
out_shape = (framebuffer.batch_size, framebuffer.height, framebuffer.width,
num_channels)
attribute_image = tf.reshape(summed_attribute, out_shape)
if background_value is None:
background_value = tf.zeros((num_channels), dtype=attribute.dtype)
else:
background_value = tf.convert_to_tensor(
background_value, dtype=attribute.dtype)
attribute_image = (
framebuffer.foreground_mask * attribute_image +
framebuffer.background_mask * background_value)
return fb.RasterizedAttribute(value=attribute_image, d_dx=None, d_dy=None)
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 test_initialize_rasterized_attribute_with_wrong_rank(self):
with self.assertRaisesRegex(ValueError,
"Expected input value to be rank 4"):
fb.RasterizedAttribute(tf.ones([4, 4, 1]))
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))