tfg.rendering.splat.rasterize_then_splat

Rasterization with differentiable occlusion using rasterize-then-splat.

tfg.rendering.splat.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 = tfg.rendering.rasterization_backend.RasterizationBackends.CPU,
    name='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).

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".

a [A1, ..., An, H, W, 4] tensor of RGBA values.