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
