def rotate(point: type_alias.TensorLike,
matrix: type_alias.TensorLike,
name: str = "rotation_matrix_3d_rotate") -> tf.Tensor:
"""Rotate a point using a rotation matrix 3d.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
Args:
point: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents a 3d point.
matrix: A tensor of shape `[A1, ..., An, 3,3]`, where the last dimension
represents a 3d rotation matrix.
name: A name for this op that defaults to "rotation_matrix_3d_rotate".
Returns:
A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents
a 3d point.
Raises:
ValueError: If the shape of `point` or `rotation_matrix_3d` is not
supported.
"""
with tf.name_scope(name):
point = tf.convert_to_tensor(value=point)
matrix = tf.convert_to_tensor(value=matrix)
shape.check_static(tensor=point,
tensor_name="point",
has_dim_equals=(-1, 3))
shape.check_static(tensor=matrix,
tensor_name="matrix",
has_rank_greater_than=1,
has_dim_equals=((-2, 3), (-1, 3)))
shape.compare_batch_dimensions(tensors=(point, matrix),
tensor_names=("point", "matrix"),
last_axes=(-2, -3),
broadcast_compatible=True)
matrix = assert_rotation_matrix_normalized(matrix)
point = tf.expand_dims(point, axis=-1)
common_batch_shape = shape.get_broadcasted_shape(
point.shape[:-2], matrix.shape[:-2])
def dim_value(dim):
return 1 if dim is None else tf.compat.dimension_value(dim)
common_batch_shape = [dim_value(dim) for dim in common_batch_shape]
point = tf.broadcast_to(point, common_batch_shape + [3, 1])
matrix = tf.broadcast_to(matrix, common_batch_shape + [3, 3])
rotated_point = tf.matmul(matrix, point)
return tf.squeeze(rotated_point, axis=-1)
def rotate(point, matrix, name=None):
"""Rotates a 2d point using a 2d rotation matrix.
Note:
In the following, A1 to An are optional batch dimensions, which must be
identical.
Args:
point: A tensor of shape `[A1, ..., An, 2]`, where the last dimension
represents a 2d point.
matrix: A tensor of shape `[A1, ..., An, 2, 2]`, where the last two
dimensions represent a 2d rotation matrix.
name: A name for this op that defaults to "rotation_matrix_2d_rotate".
Returns:
A tensor of shape `[A1, ..., An, 2]`, where the last dimension
represents a 2d point.
Raises:
ValueError: If the shape of `point` or `matrix` is not supported.
"""
with tf.compat.v1.name_scope(name, "rotation_matrix_2d_rotate",
[point, matrix]):
point = tf.convert_to_tensor(value=point)
matrix = tf.convert_to_tensor(value=matrix)
shape.check_static(tensor=point,
tensor_name="point",
has_dim_equals=(-1, 2))
shape.check_static(tensor=matrix,
tensor_name="matrix",
has_rank_greater_than=1,
has_dim_equals=((-2, 2), (-1, 2)))
shape.compare_batch_dimensions(tensors=(point, matrix),
tensor_names=("point", "matrix"),
last_axes=(-2, -3),
broadcast_compatible=True)
point = tf.expand_dims(point, axis=-1)
common_batch_shape = shape.get_broadcasted_shape(
point.shape[:-2], matrix.shape[:-2])
def dim_value(dim):
return 1 if dim is None else tf.compat.v1.dimension_value(dim)
common_batch_shape = [dim_value(dim) for dim in common_batch_shape]
point = tf.broadcast_to(point, common_batch_shape + [2, 1])
matrix = tf.broadcast_to(matrix, common_batch_shape + [2, 2])
rotated_point = tf.matmul(matrix, point)
return tf.squeeze(rotated_point, axis=-1)
def test_get_broadcasted_shape(self, shape_x, shape_y, broadcasted_shape):
"""Checks if the get_broadcasted_shape function works as expected."""
if tf.executing_eagerly():
if (shape_x is None or shape_y is None or None in shape_x or
None in shape_y):
return
shape_x = tf.compat.v1.placeholder_with_default(
tf.zeros(shape_x, dtype=tf.float32), shape=shape_x).shape
shape_y = tf.compat.v1.placeholder_with_default(
tf.zeros(shape_y, dtype=tf.float32), shape=shape_y).shape
else:
shape_x = tf.compat.v1.placeholder(shape=shape_x, dtype=tf.float32).shape
shape_y = tf.compat.v1.placeholder(shape=shape_y, dtype=tf.float32).shape
self.assertAllEqual(
shape.get_broadcasted_shape(shape_x, shape_y), broadcasted_shape)
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 (height, width) containing the dimensions
in pixels of the rasterized image.
name: A name for this op. Defaults to 'rasterization_backend_rasterize'.
Returns:
A tuple of 3 elements. The first one of shape `[A1, ..., An, H, W, 1]`
representing the triangle index associated with each pixel. If no triangle
is associated to a pixel, the index is set to -1.
The second element in the tuple is of shape `[A1, ..., An, H, W, 3]` and
correspond to barycentric coordinates per pixel. The last element in the
tuple is of shape `[A1, ..., An, H, W, 1]` and stores a value of `0` of the
pixel is assciated with the background, and `1` with the foreground
"""
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],
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] - 1, tf.int32)
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(triangle_index >= 0, tf.int32)
return triangle_index, barycentric_coordinates, mask
def blend(points,
skinning_weights,
bone_rotations,
bone_translations,
name=None):
"""Transforms the points using Linear Blend Skinning.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible and allow transforming full 3D shapes at once.
In the following, B1 to Bm are optional batch dimensions, which allow
transforming multiple poses at once.
Args:
points: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents a 3d point.
skinning_weights: A tensor of shape `[A1, ..., An, W]`, where the last
dimension represents the skinning weights of each bone.
bone_rotations: A tensor of shape `[B1, ..., Bm, W, 3, 3]`, which represents
the 3d rotations applied to each bone.
bone_translations: A tensor of shape `[B1, ..., Bm, W, 3]`, which represents
the 3d translation vectors applied to each bone.
name: A name for this op that defaults to "linear_blend_skinning_blend".
Returns:
A tensor of shape `[B1, ..., Bm, A1, ..., An, 3]`, where the last dimension
represents a 3d point.
Raises:
ValueError: If the shape of the input tensors are not supported.
"""
with tf.compat.v1.name_scope(
name, "linear_blend_skinning_blend",
[points, skinning_weights, bone_rotations, bone_translations]):
points = tf.convert_to_tensor(value=points)
skinning_weights = tf.convert_to_tensor(value=skinning_weights)
bone_rotations = tf.convert_to_tensor(value=bone_rotations)
bone_translations = tf.convert_to_tensor(value=bone_translations)
shape.check_static(
tensor=points, tensor_name="points", has_dim_equals=(-1, 3))
shape.check_static(
tensor=bone_rotations,
tensor_name="bone_rotations",
has_rank_greater_than=2,
has_dim_equals=((-2, 3), (-1, 3)))
shape.check_static(
tensor=bone_translations,
tensor_name="bone_translations",
has_rank_greater_than=1,
has_dim_equals=(-1, 3))
shape.compare_dimensions(
tensors=(skinning_weights, bone_rotations),
tensor_names=("skinning_weights", "bone_rotations"),
axes=(-1, -3))
shape.compare_dimensions(
tensors=(skinning_weights, bone_translations),
tensor_names=("skinning_weights", "bone_translations"),
axes=(-1, -2))
shape.compare_batch_dimensions(
tensors=(points, skinning_weights),
tensor_names=("points", "skinning_weights"),
last_axes=(-2, -2),
broadcast_compatible=True)
shape.compare_batch_dimensions(
tensors=(bone_rotations, bone_translations),
tensor_names=("bone_rotations", "bone_translations"),
last_axes=(-3, -2),
broadcast_compatible=True)
num_bones = skinning_weights.shape[-1]
def dim_value(dim):
return 1 if dim is None else tf.compat.v1.dimension_value(dim)
# TODO(b/148362025): factorize this block out
points_batch_shape = shape.get_broadcasted_shape(
points.shape[:-1], skinning_weights.shape[:-1])
points_batch_shape = [dim_value(dim) for dim in points_batch_shape]
points = tf.broadcast_to(points, points_batch_shape + [3])
skinning_weights = tf.broadcast_to(skinning_weights,
points_batch_shape + [num_bones])
bones_batch_shape = shape.get_broadcasted_shape(
bone_rotations.shape[:-3], bone_translations.shape[:-2])
bones_batch_shape = [dim_value(dim) for dim in bones_batch_shape]
bone_rotations = tf.broadcast_to(bone_rotations,
bones_batch_shape + [num_bones, 3, 3])
bone_translations = tf.broadcast_to(bone_translations,
bones_batch_shape + [num_bones, 3])
points_batch_dims = points.shape.ndims - 1
bones_batch_dims = bone_rotations.shape.ndims - 3
points = tf.reshape(points,
[1] * bones_batch_dims + points_batch_shape + [1, 3])
skinning_weights = tf.reshape(skinning_weights, [1] * bones_batch_dims +
points_batch_shape + [num_bones, 1])
bone_rotations = tf.reshape(
bone_rotations,
bones_batch_shape + [1] * points_batch_dims + [num_bones, 3, 3])
bone_translations = tf.reshape(
bone_translations,
bones_batch_shape + [1] * points_batch_dims + [num_bones, 3])
transformed_points = rotation_matrix_3d.rotate(
points, bone_rotations) + bone_translations
weighted_points = tf.multiply(skinning_weights, transformed_points)
blended_points = tf.reduce_sum(input_tensor=weighted_points, axis=-2)
return blended_points
def brdf(direction_incoming_light: type_alias.TensorLike,
direction_outgoing_light: type_alias.TensorLike,
surface_normal: type_alias.TensorLike,
albedo: type_alias.TensorLike,
name: str = "lambertian_brdf") -> tf.Tensor:
"""Evaluates the brdf of a Lambertian surface.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
Note:
The gradient of this function is not smooth when the dot product of the
normal with any light is 0.0.
Args:
direction_incoming_light: A tensor of shape `[A1, ..., An, 3]`, where the
last dimension represents a normalized incoming light vector.
direction_outgoing_light: A tensor of shape `[A1, ..., An, 3]`, where the
last dimension represents a normalized outgoing light vector.
surface_normal: A tensor of shape `[A1, ..., An, 3]`, where the last
dimension represents a normalized surface normal.
albedo: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents albedo with values in [0,1].
name: A name for this op. Defaults to "lambertian_brdf".
Returns:
A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents
the amount of reflected light in any outgoing direction.
Raises:
ValueError: if the shape of `direction_incoming_light`,
`direction_outgoing_light`, `surface_normal`, `shininess` or `albedo` is not
supported.
InvalidArgumentError: if at least one element of `albedo` is outside of
[0,1].
"""
with tf.name_scope(name):
direction_incoming_light = tf.convert_to_tensor(
value=direction_incoming_light)
direction_outgoing_light = tf.convert_to_tensor(
value=direction_outgoing_light)
surface_normal = tf.convert_to_tensor(value=surface_normal)
albedo = tf.convert_to_tensor(value=albedo)
shape.check_static(tensor=direction_incoming_light,
tensor_name="direction_incoming_light",
has_dim_equals=(-1, 3))
shape.check_static(tensor=direction_outgoing_light,
tensor_name="direction_outgoing_light",
has_dim_equals=(-1, 3))
shape.check_static(tensor=surface_normal,
tensor_name="surface_normal",
has_dim_equals=(-1, 3))
shape.check_static(tensor=albedo,
tensor_name="albedo",
has_dim_equals=(-1, 3))
shape.compare_batch_dimensions(
tensors=(direction_incoming_light, direction_outgoing_light,
surface_normal, albedo),
tensor_names=("direction_incoming_light",
"direction_outgoing_light", "surface_normal",
"albedo"),
last_axes=-2,
broadcast_compatible=True)
direction_incoming_light = asserts.assert_normalized(
direction_incoming_light)
direction_outgoing_light = asserts.assert_normalized(
direction_outgoing_light)
surface_normal = asserts.assert_normalized(surface_normal)
albedo = asserts.assert_all_in_range(albedo,
0.0,
1.0,
open_bounds=False)
# Checks whether the incoming or outgoing light point behind the surface.
dot_incoming_light_surface_normal = vector.dot(
-direction_incoming_light, surface_normal)
dot_outgoing_light_surface_normal = vector.dot(
direction_outgoing_light, surface_normal)
min_dot = tf.minimum(dot_incoming_light_surface_normal,
dot_outgoing_light_surface_normal)
common_shape = shape.get_broadcasted_shape(min_dot.shape, albedo.shape)
d_val = lambda dim: 1 if dim is None else tf.compat.dimension_value(dim
)
common_shape = [d_val(dim) for dim in common_shape]
condition = tf.broadcast_to(tf.greater_equal(min_dot, 0.0),
common_shape)
albedo = tf.broadcast_to(albedo, common_shape)
return tf.where(condition, albedo / math.pi, tf.zeros_like(albedo))
def vertex_normals(vertices, indices, clockwise=True, name=None):
"""Computes vertex normals from a mesh.
This function computes vertex normals as the weighted sum of the adjacent
face normals, where the weights correspond to the area of each face. This
function supports planar convex polygon faces. For non-triangular meshes,
this function converts them into triangular meshes to calculate vertex
normals.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
vertices: A tensor of shape `[A1, ..., An, V, 3]`, where V is the number of
vertices.
indices: A tensor of shape `[A1, ..., An, F, M]`, where F is the number of
faces and M is the number of vertices per face.
clockwise: Winding order to determine front-facing faces. The order of
vertices should be either clockwise or counterclockwise.
name: A name for this op. Defaults to "normals_vertex_normals".
Returns:
A tensor of shape `[A1, ..., An, V, 3]` containing vertex normals. If
vertices and indices have different batch dimensions, this function
broadcasts them into the same batch dimensions and the output batch
dimensions are the broadcasted.
Raises:
ValueError: If the shape of `vertices`, `indices` is not supported.
"""
with tf.compat.v1.name_scope(name, "normals_vertex_normals",
[vertices, indices]):
vertices = tf.convert_to_tensor(value=vertices)
indices = tf.convert_to_tensor(value=indices)
shape.check_static(tensor=vertices,
tensor_name="vertices",
has_rank_greater_than=1,
has_dim_equals=(-1, 3))
shape.check_static(tensor=indices,
tensor_name="indices",
has_rank_greater_than=1,
has_dim_greater_than=(-1, 2))
shape.compare_batch_dimensions(tensors=(vertices, indices),
last_axes=(-3, -3),
broadcast_compatible=True)
shape_indices = indices.shape.as_list()
if None in shape_indices[:-2]:
raise ValueError("'indices' must have specified batch dimensions.")
common_batch_dims = shape.get_broadcasted_shape(
vertices.shape[:-2], indices.shape[:-2])
vertices_repeat = [
common_batch_dims[x] // vertices.shape.as_list()[x]
for x in range(len(common_batch_dims))
]
indices_repeat = [
common_batch_dims[x] // shape_indices[x]
for x in range(len(common_batch_dims))
]
vertices = tf.tile(vertices,
vertices_repeat + [1, 1],
name="vertices_broadcast")
indices = tf.tile(indices,
indices_repeat + [1, 1],
name="indices_broadcast")
# Triangulate non-triangular faces.
if shape_indices[-1] > 3:
triangle_indices = []
for i in range(1, shape_indices[-1] - 1):
triangle_indices.append(
tf.concat((indices[..., 0:1], indices[..., i:i + 2]),
axis=-1))
indices = tf.concat(triangle_indices, axis=-2)
shape_indices = indices.shape.as_list()
face_vertices = gather_faces(vertices, indices)
# Use unnormalized face normals to scale normals by area.
mesh_face_normals = face_normals(face_vertices,
clockwise=clockwise,
normalize=False)
if vertices.shape.ndims > 2:
outer_indices = np.meshgrid(
*[np.arange(i) for i in shape_indices[:-2]],
sparse=False,
indexing="ij")
outer_indices = [np.expand_dims(i, axis=-1) for i in outer_indices]
outer_indices = np.concatenate(outer_indices, axis=-1)
outer_indices = np.expand_dims(outer_indices, axis=-2)
outer_indices = tf.constant(outer_indices, dtype=tf.int32)
outer_indices = tf.tile(outer_indices,
[1] * len(shape_indices[:-2]) +
[tf.shape(input=indices)[-2]] + [1])
unnormalized_vertex_normals = tf.zeros_like(vertices)
for i in range(shape_indices[-1]):
scatter_indices = tf.concat(
[outer_indices, indices[..., i:i + 1]], axis=-1)
unnormalized_vertex_normals = tf.compat.v1.tensor_scatter_add(
unnormalized_vertex_normals, scatter_indices,
mesh_face_normals)
else:
unnormalized_vertex_normals = tf.zeros_like(vertices)
for i in range(shape_indices[-1]):
unnormalized_vertex_normals = tf.compat.v1.tensor_scatter_add(
unnormalized_vertex_normals, indices[..., i:i + 1],
mesh_face_normals)
vector_norms = tf.sqrt(
tf.reduce_sum(input_tensor=unnormalized_vertex_normals**2,
axis=-1,
keepdims=True))
return safe_ops.safe_unsigned_div(unnormalized_vertex_normals,
vector_norms)
def brdf(direction_incoming_light: type_alias.TensorLike,
direction_outgoing_light: type_alias.TensorLike,
surface_normal: type_alias.TensorLike,
shininess: type_alias.TensorLike,
albedo: type_alias.TensorLike,
brdf_normalization: bool = True,
name: str = "phong_brdf") -> tf.Tensor:
"""Evaluates the specular brdf of the Phong model.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
Note:
The gradient of this function is not smooth when the dot product of the
normal with any light is 0.0.
Args:
direction_incoming_light: A tensor of shape `[A1, ..., An, 3]`, where the
last dimension represents a normalized incoming light vector.
direction_outgoing_light: A tensor of shape `[A1, ..., An, 3]`, where the
last dimension represents a normalized outgoing light vector.
surface_normal: A tensor of shape `[A1, ..., An, 3]`, where the last
dimension represents a normalized surface normal.
shininess: A tensor of shape `[A1, ..., An, 1]`, where the last dimension
represents a non-negative shininess coefficient.
albedo: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents albedo with values in [0,1].
brdf_normalization: A `bool` indicating whether normalization should be
applied to enforce the energy conservation property of BRDFs. Note that
`brdf_normalization` must be set to False in order to use the original
Blinn specular model.
name: A name for this op. Defaults to "phong_brdf".
Returns:
A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents
the amount of light reflected in the outgoing light direction.
Raises:
ValueError: if the shape of `direction_incoming_light`,
`direction_outgoing_light`, `surface_normal`, `shininess` or `albedo` is not
supported.
InvalidArgumentError: if not all of shininess values are non-negative, or if
at least one element of `albedo` is outside of [0,1].
"""
with tf.name_scope(name):
direction_incoming_light = tf.convert_to_tensor(
value=direction_incoming_light)
direction_outgoing_light = tf.convert_to_tensor(
value=direction_outgoing_light)
surface_normal = tf.convert_to_tensor(value=surface_normal)
shininess = tf.convert_to_tensor(value=shininess)
albedo = tf.convert_to_tensor(value=albedo)
shape.check_static(tensor=direction_incoming_light,
tensor_name="direction_incoming_light",
has_dim_equals=(-1, 3))
shape.check_static(tensor=direction_outgoing_light,
tensor_name="direction_outgoing_light",
has_dim_equals=(-1, 3))
shape.check_static(tensor=surface_normal,
tensor_name="surface_normal",
has_dim_equals=(-1, 3))
shape.check_static(tensor=shininess,
tensor_name="shininess",
has_dim_equals=(-1, 1))
shape.check_static(tensor=albedo,
tensor_name="albedo",
has_dim_equals=(-1, 3))
shape.compare_batch_dimensions(
tensors=(direction_incoming_light, direction_outgoing_light,
surface_normal, shininess, albedo),
tensor_names=("direction_incoming_light",
"direction_outgoing_light", "surface_normal",
"shininess", "albedo"),
last_axes=-2,
broadcast_compatible=True)
direction_incoming_light = asserts.assert_normalized(
direction_incoming_light)
direction_outgoing_light = asserts.assert_normalized(
direction_outgoing_light)
surface_normal = asserts.assert_normalized(surface_normal)
albedo = asserts.assert_all_in_range(albedo,
0.0,
1.0,
open_bounds=False)
shininess = asserts.assert_all_above(shininess, 0.0, open_bound=False)
# Checks whether the incoming or outgoing light point behind the surface.
dot_incoming_light_surface_normal = vector.dot(
-direction_incoming_light, surface_normal)
dot_outgoing_light_surface_normal = vector.dot(
direction_outgoing_light, surface_normal)
min_dot = tf.minimum(dot_incoming_light_surface_normal,
dot_outgoing_light_surface_normal)
perfect_reflection_direction = vector.reflect(direction_incoming_light,
surface_normal)
perfect_reflection_direction = tf.math.l2_normalize(
perfect_reflection_direction, axis=-1)
cos_alpha = vector.dot(perfect_reflection_direction,
direction_outgoing_light,
axis=-1)
cos_alpha = tf.maximum(cos_alpha, tf.zeros_like(cos_alpha))
phong_model = albedo * tf.pow(cos_alpha, shininess)
if brdf_normalization:
phong_model *= _brdf_normalization_factor(shininess)
common_shape = shape.get_broadcasted_shape(min_dot.shape,
phong_model.shape)
d_val = lambda dim: 1 if dim is None else tf.compat.dimension_value(dim
)
common_shape = [d_val(dim) for dim in common_shape]
condition = tf.broadcast_to(tf.greater_equal(min_dot, 0.0),
common_shape)
phong_model = tf.broadcast_to(phong_model, common_shape)
return tf.where(condition, phong_model, tf.zeros_like(phong_model))
def matrix_from_intrinsics(focal,
principal_point,
skew=(0.0, ),
name="perspective_matrix_from_intrinsics"):
r"""Builds calibration matrix from intrinsic parameters.
Builds the camera calibration matrix as
$$
\mathbf{C} =
\begin{bmatrix}
f_x & sc & c_x \\
0 & f_y & c_y \\
0 & 0 & 1 \\
\end{bmatrix}
$$
from the focal length \\((f_x, f_y)\\) and the principal point
\\((c_x, c_y)\\).
Note:
In the following, A1 to An are optional batch dimensions.
Args:
focal: A tensor of shape `[A1, ..., An, 2]`, where the last dimension
represents a camera focal length.
principal_point: A tensor of shape `[A1, ..., An, 2]`, where the last
dimension represents a camera principal point.
skew: A tensor of shape `[A1, ..., An, 1]`, where the last dimension
represents a skew coefficient.
name: A name for this op that defaults to
"perspective_matrix_from_intrinsics".
Returns:
A tensor of shape `[A1, ..., An, 3, 3]`, where the last two dimensions
represent a camera calibration matrix.
Raises:
ValueError: If the shape of `focal`, or `principal_point` is not
supported.
"""
with tf.name_scope(name):
focal = tf.convert_to_tensor(value=focal)
principal_point = tf.convert_to_tensor(value=principal_point)
skew = tf.convert_to_tensor(value=skew)
common_batch_shape = shape.get_broadcasted_shape(
focal.shape[:-1], skew.shape[:-1])
def dim_value(dim):
return 1 if dim is None else tf.compat.dimension_value(dim)
common_batch_shape = [dim_value(dim) for dim in common_batch_shape]
skew = tf.broadcast_to(skew, common_batch_shape + [1])
shape.check_static(tensor=focal,
tensor_name="focal",
has_dim_equals=(-1, 2))
shape.check_static(tensor=principal_point,
tensor_name="principal_point",
has_dim_equals=(-1, 2))
shape.check_static(
tensor=skew,
tensor_name="skew",
has_dim_equals=(-1, 1),
)
shape.compare_batch_dimensions(tensors=(focal, principal_point, skew),
tensor_names=("focal",
"principal_point",
"skew"),
last_axes=-2,
broadcast_compatible=False)
fx, fy = tf.unstack(focal, axis=-1)
cx, cy = tf.unstack(principal_point, axis=-1)
zero = tf.zeros_like(fx)
one = tf.ones_like(fx)
skew = tf.reshape(skew, tf.shape(fx))
matrix = tf.stack((fx, skew, cx, zero, fy, cy, zero, zero, one),
axis=-1) # pyformat: disable
matrix_shape = tf.shape(input=matrix)
output_shape = tf.concat((matrix_shape[:-1], (3, 3)), axis=-1)
return tf.reshape(matrix, shape=output_shape)
def estimate_radiance(point_light_radiance,
point_light_position,
surface_point_position,
surface_point_normal,
observation_point,
brdf,
name="estimate_radiance",
reflected_light_fall_off=False):
"""Estimates the spectral radiance of a point light reflected from the surface point towards the observation point.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
B1 to Bm are optional batch dimensions for the lights, which must be
broadcast compatible.
Note:
In case the light or the observation point are located behind the surface
the function will return 0.
Note:
The gradient of this function is not smooth when the dot product of the
normal with the light-to-surface or surface-to-observation vectors is 0.
Args:
point_light_radiance: A tensor of shape '[B1, ..., Bm, K]', where the last
axis represents the radiance of the point light at a specific wave length.
point_light_position: A tensor of shape `[B1, ..., Bm, 3]`, where the last
axis represents the position of the point light.
surface_point_position: A tensor of shape `[A1, ..., An, 3]`, where the last
axis represents the position of the surface point.
surface_point_normal: A tensor of shape `[A1, ..., An, 3]`, where the last
axis represents the normalized surface normal at the given surface point.
observation_point: A tensor of shape `[A1, ..., An, 3]`, where the last axis
represents the observation point.
brdf: The BRDF of the surface as a function of: incoming_light_direction -
The incoming light direction as the last axis of a tensor with shape `[A1,
..., An, 3]`. outgoing_light_direction - The outgoing light direction as
the last axis of a tensor with shape `[A1, ..., An, 3]`.
surface_point_normal - The surface normal as the last axis of a tensor
with shape `[A1, ..., An, 3]`. Note - The BRDF should return a tensor of
size '[A1, ..., An, K]' where the last axis represents the amount of
reflected light in each wave length.
name: A name for this op. Defaults to "estimate_radiance".
reflected_light_fall_off: A boolean specifying whether or not to include the
fall off of the light reflected from the surface towards the observation
point in the calculation. Defaults to False.
Returns:
A tensor of shape `[A1, ..., An, B1, ..., Bm, K]`, where the last
axis represents the amount of light received at the observation point
after being reflected from the given surface point.
Raises:
ValueError: if the shape of `point_light_position`,
`surface_point_position`, `surface_point_normal`, or `observation_point` is
not supported.
InvalidArgumentError: if 'surface_point_normal' is not normalized.
"""
with tf.name_scope(name):
point_light_radiance = tf.convert_to_tensor(value=point_light_radiance)
point_light_position = tf.convert_to_tensor(value=point_light_position)
surface_point_position = tf.convert_to_tensor(value=surface_point_position)
surface_point_normal = tf.convert_to_tensor(value=surface_point_normal)
observation_point = tf.convert_to_tensor(value=observation_point)
shape.check_static(
tensor=point_light_position,
tensor_name="point_light_position",
has_dim_equals=(-1, 3))
shape.check_static(
tensor=surface_point_position,
tensor_name="surface_point_position",
has_dim_equals=(-1, 3))
shape.check_static(
tensor=surface_point_normal,
tensor_name="surface_point_normal",
has_dim_equals=(-1, 3))
shape.check_static(
tensor=observation_point,
tensor_name="observation_point",
has_dim_equals=(-1, 3))
shape.compare_batch_dimensions(
tensors=(surface_point_position, surface_point_normal,
observation_point),
tensor_names=("surface_point_position", "surface_point_normal",
"observation_point"),
last_axes=-2,
broadcast_compatible=True)
shape.compare_batch_dimensions(
tensors=(point_light_radiance, point_light_position),
tensor_names=("point_light_radiance", "point_light_position"),
last_axes=-2,
broadcast_compatible=True)
surface_point_normal = asserts.assert_normalized(surface_point_normal)
# Get the number of lights dimensions (B1,...,Bm).
lights_num_dimensions = max(
len(point_light_radiance.shape), len(point_light_position.shape)) - 1
# Reshape the other parameters so they can be broadcasted to the output of
# shape [A1,...,An, B1,...,Bm, K].
surface_point_position = tf.reshape(
surface_point_position,
surface_point_position.shape[:-1] + (1,) * lights_num_dimensions + (3,))
surface_point_normal = tf.reshape(
surface_point_normal,
surface_point_normal.shape[:-1] + (1,) * lights_num_dimensions + (3,))
observation_point = tf.reshape(
observation_point,
observation_point.shape[:-1] + (1,) * lights_num_dimensions + (3,))
light_to_surface_point = surface_point_position - point_light_position
distance_light_surface_point = tf.norm(
tensor=light_to_surface_point, axis=-1, keepdims=True)
incoming_light_direction = tf.math.l2_normalize(
light_to_surface_point, axis=-1)
surface_to_observation_point = observation_point - surface_point_position
outgoing_light_direction = tf.math.l2_normalize(
surface_to_observation_point, axis=-1)
brdf_value = brdf(incoming_light_direction, outgoing_light_direction,
surface_point_normal)
incoming_light_dot_surface_normal = vector.dot(-incoming_light_direction,
surface_point_normal)
outgoing_light_dot_surface_normal = vector.dot(outgoing_light_direction,
surface_point_normal)
estimated_radiance = (point_light_radiance * \
brdf_value * incoming_light_dot_surface_normal) / \
(4. * math.pi * tf.math.square(distance_light_surface_point))
if reflected_light_fall_off:
distance_surface_observation_point = tf.norm(
tensor=surface_to_observation_point, axis=-1, keepdims=True)
estimated_radiance = estimated_radiance / \
tf.math.square(distance_surface_observation_point)
# Create a condition for checking whether the light or observation point are
# behind the surface.
min_dot = tf.minimum(incoming_light_dot_surface_normal,
outgoing_light_dot_surface_normal)
common_shape = shape.get_broadcasted_shape(min_dot.shape,
estimated_radiance.shape)
d_val = lambda dim: 1 if dim is None else tf.compat.dimension_value(dim)
common_shape = [d_val(dim) for dim in common_shape]
condition = tf.broadcast_to(tf.greater_equal(min_dot, 0.0), common_shape)
return tf.where(condition, estimated_radiance,
tf.zeros_like(estimated_radiance))
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))
