def project(point_3d: type_alias.TensorLike,
name: str = "orthographic_project") -> tf.Tensor:
r"""Projects a 3d point onto the 2d camera plane.
Projects a 3d point \\((x, y, z)\\) to a 2d point \\((x', y')\\) onto the
image plane, with
$$
\begin{matrix}
x' = x, & y' = y.
\end{matrix}
$$
Note:
In the following, A1 to An are optional batch dimensions.
Args:
point_3d: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents a 3d point to project.
name: A name for this op that defaults to "orthographic_project".
Returns:
A tensor of shape `[A1, ..., An, 2]`, where the last dimension represents
a 2d point.
Raises:
ValueError: If the shape of `point_3d` is not supported.
"""
with tf.name_scope(name):
point_3d = tf.convert_to_tensor(value=point_3d)
shape.check_static(
tensor=point_3d, tensor_name="point_3d", has_dim_equals=(-1, 3))
point_xy, _ = tf.split(point_3d, (2, 1), axis=-1)
return point_xy
def loss(matte, weights, name=None):
"""Computes the matting loss function based on the matting Laplacian weights.
Computes the matting loss function based on the `weights` generated by the
`laplacian_weights` function which implements the approach proposed by Levin
et al. in "A Closed Form Solution to Natural Image Matting".
Args:
matte: A tensor of shape `[B, H, W, 1]`.
weights: A tensor containing the Laplacian weights computed by the
`laplacian_weights` function.
name: A name for this op. Defaults to "matting_loss".
Returns:
A tensor containing a scalar value defining the matting loss.
Raises:
ValueError: If the last dimension of `matte` is not 1. If `matte` is not
of rank 4. If the last two dimensions of `weights` are not of the
same size. If `weights` is not of rank 5. If `B` is different
between `matte` and `weights`.
"""
with tf.compat.v1.name_scope(name, "matting_loss", [matte, weights]):
matte = tf.convert_to_tensor(value=matte)
weights = tf.convert_to_tensor(value=weights)
pixels = tf.compat.dimension_value(weights.shape[-1])
shape.check_static(matte, has_rank=4, has_dim_equals=(-1, 1))
shape.check_static(weights, has_rank=5, has_dim_equals=(-2, pixels))
shape.compare_batch_dimensions(
tensors=(matte, weights), last_axes=0, broadcast_compatible=False)
size = np.sqrt(pixels)
patches = tf.expand_dims(_image_patches(matte, size), axis=-2)
losses = _quadratic_form(weights, patches)
return tf.reduce_mean(input_tensor=losses)
def match_intermediate_batch_dimensions(tensor1, tensor2):
"""Match the batch dimensions.
Args:
tensor1: A tensor of shape `[A1, M]`.
tensor2: A tensor of shape `[A1, ..., An, N]`.
Returns:
A tensor of shape `[A1, ..., An, M]`.
"""
shape.check_static(tensor=tensor1, tensor_name="tensor1", has_rank=2)
shape.check_static(tensor=tensor2,
tensor_name="tensor2",
has_rank_greater_than=1)
shape.compare_dimensions(tensors=(tensor1, tensor2),
tensor_names=("tensor1", "tensor2"),
axes=0)
shape1 = tf.shape(tensor1)
shape2 = tf.shape(tensor2)
shape_diff = len(shape2) - len(shape1)
new_shape = tf.concat([[shape1[0]], [1] * shape_diff, [shape1[-1]]],
axis=-1)
target_shape = tf.concat([shape2[:-1], [shape1[-1]]], axis=-1)
return tf.broadcast_to(tf.reshape(tensor1, new_shape), target_shape)
def inverse(axis, angle, name=None):
"""Computes the axis-angle that is the inverse of the input axis-angle.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
axis: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents a normalized axis.
angle: A tensor of shape `[A1, ..., An, 1]` where the last dimension
represents an angle.
name: A name for this op that defaults to "axis_angle_inverse".
Returns:
A tuple of two tensors, respectively of shape `[A1, ..., An, 3]` and
`[A1, ..., An, 1]`, where the first tensor represents the axis, and the
second represents the angle. The resulting axis is a normalized vector.
Raises:
ValueError: If the shape of `axis` or `angle` is not supported.
"""
with tf.compat.v1.name_scope(name, "axis_angle_inverse", [axis, angle]):
axis = tf.convert_to_tensor(value=axis)
angle = tf.convert_to_tensor(value=angle)
shape.check_static(tensor=axis, tensor_name="axis", has_dim_equals=(-1, 3))
shape.check_static(
tensor=angle, tensor_name="angle", has_dim_equals=(-1, 1))
shape.compare_batch_dimensions(
tensors=(axis, angle),
tensor_names=("axis", "angle"),
last_axes=-2,
broadcast_compatible=True)
axis = asserts.assert_normalized(axis)
return axis, -angle
def from_axis_angle(axis, angle, name=None):
"""Converts an axis-angle representation to a quaternion.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
axis: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents a normalized axis.
angle: A tensor of shape `[A1, ..., An, 1]`, where the last dimension
represents an angle.
name: A name for this op that defaults to "quaternion_from_axis_angle".
Returns:
A tensor of shape `[A1, ..., An, 4]`, where the last dimension represents
a normalized quaternion.
Raises:
ValueError: If the shape of `axis` or `angle` is not supported.
"""
with tf.compat.v1.name_scope(name, "quaternion_from_axis_angle",
[axis, angle]):
axis = tf.convert_to_tensor(value=axis)
angle = tf.convert_to_tensor(value=angle)
shape.check_static(tensor=axis, tensor_name="axis", has_dim_equals=(-1, 3))
shape.check_static(
tensor=angle, tensor_name="angle", has_dim_equals=(-1, 1))
shape.compare_batch_dimensions(
tensors=(axis, angle), last_axes=-2, broadcast_compatible=True)
axis = asserts.assert_normalized(axis)
half_angle = 0.5 * angle
w = tf.cos(half_angle)
xyz = tf.sin(half_angle) * axis
return tf.concat((xyz, w), axis=-1)
def upsample(image, num_levels, name=None):
"""Generates the different levels of the pyramid (upsampling).
Args:
image: A tensor of shape `[B, H, W, C]`, where `B` is the batch size, `H`
the height of the image, `W` the width of the image, and `C` the number of
channels of the image.
num_levels: The number of levels to generate.
name: A name for this op that defaults to "pyramid_upsample".
Returns:
A list containing `num_levels` tensors of shape `[B, H_i, W_i, C]`, where
`H_i` and `W_i` are the height and width of the upsampled image for the
level i.
Raises:
ValueError: If the shape of `image` is not supported.
"""
with tf.compat.v1.name_scope(name, "pyramid_upsample", [image]):
image = tf.convert_to_tensor(value=image)
shape.check_static(tensor=image, tensor_name="image", has_rank=4)
return _build_pyramid(image, _upsample, num_levels)
def multiply(quaternion1, quaternion2, name=None):
"""Multiplies two quaternions.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
quaternion1: A tensor of shape `[A1, ..., An, 4]`, where the last dimension
represents a quaternion.
quaternion2: A tensor of shape `[A1, ..., An, 4]`, where the last dimension
represents a quaternion.
name: A name for this op that defaults to "quaternion_multiply".
Returns:
A tensor of shape `[A1, ..., An, 4]` representing quaternions.
Raises:
ValueError: If the shape of `quaternion1` or `quaternion2` is not supported.
"""
with tf.compat.v1.name_scope(name, "quaternion_multiply",
[quaternion1, quaternion2]):
quaternion1 = tf.convert_to_tensor(value=quaternion1)
quaternion2 = tf.convert_to_tensor(value=quaternion2)
shape.check_static(
tensor=quaternion1, tensor_name="quaternion1", has_dim_equals=(-1, 4))
shape.check_static(
tensor=quaternion2, tensor_name="quaternion2", has_dim_equals=(-1, 4))
x1, y1, z1, w1 = tf.unstack(quaternion1, axis=-1)
x2, y2, z2, w2 = tf.unstack(quaternion2, axis=-1)
x = x1 * w2 + y1 * z2 - z1 * y2 + w1 * x2
y = -x1 * z2 + y1 * w2 + z1 * x2 + w1 * y2
z = x1 * y2 - y1 * x2 + z1 * w2 + w1 * z2
w = -x1 * x2 - y1 * y2 - z1 * z2 + w1 * w2
return tf.stack((x, y, z, w), axis=-1)
def from_euler_with_small_angles_approximation(angles, name=None):
r"""Convert an Euler angle representation to a rotation matrix.
The resulting matrix is $$\mathbf{R} = \mathbf{R}_z\mathbf{R}_y\mathbf{R}_x$$.
Under the small angle assumption, $$\sin(x)$$ and $$\cos(x)$$ can be
approximated by their second order Taylor expansions, where
$$\sin(x) \approx x$$ and $$\cos(x) \approx 1 - \frac{x^2}{2}$$.
In the current implementation, the smallness of the angles is not verified.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
angles: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents the three small Euler angles. `[A1, ..., An, 0]` is the angle
about `x` in radians, `[A1, ..., An, 1]` is the angle about `y` in radians
and `[A1, ..., An, 2]` is the angle about `z` in radians.
name: A name for this op that defaults to "rotation_matrix_3d_from_euler".
Returns:
A tensor of shape `[A1, ..., An, 3, 3]`, where the last two dimensions
represent a 3d rotation matrix.
Raises:
ValueError: If the shape of `angles` is not supported.
"""
with tf.compat.v1.name_scope(
name, "rotation_matrix_3d_from_euler_with_small_angles", [angles]):
angles = tf.convert_to_tensor(value=angles)
shape.check_static(
tensor=angles, tensor_name="angles", has_dim_equals=(-1, 3))
sin_angles = angles
cos_angles = 1.0 - 0.5 * tf.square(angles)
return _build_matrix_from_sines_and_cosines(sin_angles, cos_angles)
def spherical_to_cartesian_coordinates(point_spherical, name=None):
"""Function to transform Cartesian coordinates to spherical coordinates.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
point_spherical: A tensor of shape `[A1, ..., An, 3]`. The last dimension
contains r, theta, and phi that respectively correspond to the radius,
polar angle and azimuthal angle; r must be non-negative.
name: A name for this op. Defaults to 'spherical_to_cartesian_coordinates'.
Raises:
tf.errors.InvalidArgumentError: If r, theta or phi contains out of range
data.
Returns:
A tensor of shape `[A1, ..., An, 3]`, where the last dimension contains the
cartesian coordinates in x,y,z order.
"""
with tf.compat.v1.name_scope(name, "spherical_to_cartesian_coordinates",
[point_spherical]):
point_spherical = tf.convert_to_tensor(value=point_spherical)
shape.check_static(
tensor=point_spherical,
tensor_name="point_spherical",
has_dim_equals=(-1, 3))
r, theta, phi = tf.unstack(point_spherical, axis=-1)
r = asserts.assert_all_above(r, 0)
tmp = r * tf.sin(theta)
x = tmp * tf.cos(phi)
y = tmp * tf.sin(phi)
z = r * tf.cos(theta)
return tf.stack((x, y, z), axis=-1)
def is_normalized(axis, angle, atol=1e-3, name=None):
"""Determines if the axis-angle is normalized or not.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
axis: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents a normalized axis.
angle: A tensor of shape `[A1, ..., An, 1]` where the last dimension
represents an angle.
atol: The absolute tolerance parameter.
name: A name for this op that defaults to "axis_angle_is_normalized".
Returns:
A tensor of shape `[A1, ..., An, 1]`, where False indicates that the axis is
not normalized.
"""
with tf.compat.v1.name_scope(name, "axis_angle_is_normalized",
[axis, angle]):
axis = tf.convert_to_tensor(value=axis)
angle = tf.convert_to_tensor(value=angle)
shape.check_static(tensor=axis,
tensor_name="axis",
has_dim_equals=(-1, 3))
shape.check_static(tensor=angle,
tensor_name="angle",
has_dim_equals=(-1, 1))
shape.compare_batch_dimensions(tensors=(axis, angle),
tensor_names=("axis", "angle"),
last_axes=-2,
broadcast_compatible=True)
norms = tf.norm(tensor=axis, axis=-1, keepdims=True)
return tf.abs(norms - 1.) < atol
def cross(vector1, vector2, axis=-1, name=None):
"""Computes the cross product between two tensors along an axis.
Note:
In the following, A1 to An are optional batch dimensions, which should be
broadcast compatible.
Args:
vector1: A tensor of shape `[A1, ..., Ai = 3, ..., An]`, where the dimension
i = axis represents a 3d vector.
vector2: A tensor of shape `[A1, ..., Ai = 3, ..., An]`, where the dimension
i = axis represents a 3d vector.
axis: The dimension along which to compute the cross product.
name: A name for this op which defaults to "vector_cross".
Returns:
A tensor of shape `[A1, ..., Ai = 3, ..., An]`, where the dimension i = axis
represents the result of the cross product.
"""
with tf.compat.v1.name_scope(name, "vector_cross", [vector1, vector2]):
vector1 = tf.convert_to_tensor(value=vector1)
vector2 = tf.convert_to_tensor(value=vector2)
shape.check_static(
tensor=vector1, tensor_name="vector1", has_dim_equals=(axis, 3))
shape.check_static(
tensor=vector2, tensor_name="vector2", has_dim_equals=(axis, 3))
shape.compare_batch_dimensions(
tensors=(vector1, vector2), last_axes=-1, broadcast_compatible=True)
vector1_x, vector1_y, vector1_z = tf.unstack(vector1, axis=axis)
vector2_x, vector2_y, vector2_z = tf.unstack(vector2, axis=axis)
n_x = vector1_y * vector2_z - vector1_z * vector2_y
n_y = vector1_z * vector2_x - vector1_x * vector2_z
n_z = vector1_x * vector2_y - vector1_y * vector2_x
return tf.stack((n_x, n_y, n_z), axis=axis)
def check_valid_graph_unpooling_input(data, pool_map, sizes):
"""Checks that the inputs are valid for graph unpooling.
Note:
In the following, A1 to A3 are optional batch dimensions.
Args:
data: A `float` tensor with shape `[A1, ..., A3, V1, C]`.
pool_map: A `SparseTensor` with the same type as `data` and with shape
`[A1, ..., A3, V1, V2]`.
sizes: An `int` tensor of shape `[A1, ..., A3, 2]`. Can be `None`.
Raises:
TypeError: if the input types are invalid.
ValueError: if the input dimensions are invalid.
"""
if not data.dtype.is_floating:
raise TypeError("'data' must have a float type.")
if pool_map.dtype != data.dtype:
raise TypeError("'pool_map' and 'data' must have the same type.")
if sizes is not None and not sizes.dtype.is_integer:
raise TypeError("'sizes' must have an integer type.")
data_ndims = data.shape.ndims
shape.check_static(tensor=data, tensor_name="data", has_rank_greater_than=1)
shape.check_static(tensor=data, tensor_name="data", has_rank_less_than=6)
shape.check_static(
tensor=pool_map, tensor_name="pool_map", has_rank=data_ndims)
shape.compare_dimensions(
tensors=(data, pool_map),
tensor_names=("data", "pool_map"),
axes=(-2, -2))
if sizes is None:
shape.compare_batch_dimensions(
tensors=(data, pool_map),
tensor_names=("data", "pool_map"),
last_axes=-3,
broadcast_compatible=False)
else:
shape.check_static(
tensor=sizes, tensor_name="sizes", has_rank=data_ndims - 1)
shape.compare_batch_dimensions(
tensors=(data, pool_map, sizes),
tensor_names=("data", "pool_map", "sizes"),
last_axes=(-3, -3, -2),
broadcast_compatible=False)
def compute_radiance(
rgba_values: type_alias.TensorLike,
distances: type_alias.TensorLike,
name: str = "ray_radiance") -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
"""Renders the rgba values for points along a ray, as described in ["NeRF Representing Scenes as Neural Radiance Fields for View Synthesis"](https://github.com/bmild/nerf).
Note:
In the following, A1 to An are optional batch dimensions.
Args:
rgba_values: A tensor of shape `[A1, ..., An, N, 4]`, where N are the
samples on the ray.
distances: A tensor of shape `[A1, ..., An, N]` containing the distances
between the samples, where N are the samples on the ray.
name: A name for this op. Defaults to "ray_radiance".
Returns:
A tensor of shape `[A1, ..., An, 3]` for the estimated rgb values,
a tensor of shape `[A1, ..., An, 1]` for the estimated density values,
and a tensor of shape `[A1, ..., An, N]` for the sample weights.
"""
with tf.name_scope(name):
rgba_values = tf.convert_to_tensor(value=rgba_values)
distances = tf.convert_to_tensor(value=distances)
distances = tf.expand_dims(distances, -1)
shape.check_static(tensor=rgba_values,
tensor_name="rgba_values",
has_dim_equals=(-1, 4))
shape.check_static(tensor=rgba_values,
tensor_name="rgba_values",
has_rank_greater_than=1)
shape.check_static(tensor=distances,
tensor_name="distances",
has_rank_greater_than=1)
shape.compare_batch_dimensions(tensors=(rgba_values, distances),
tensor_names=("ray_values", "dists"),
last_axes=-3,
broadcast_compatible=True)
shape.compare_dimensions(tensors=(rgba_values, distances),
tensor_names=("ray_values", "dists"),
axes=-2)
rgb, density = tf.split(rgba_values, [3, 1], axis=-1)
alpha = 1. - tf.exp(-density * distances)
alpha = tf.squeeze(alpha, -1)
ray_sample_weights = alpha * tf.math.cumprod(
1. - alpha + 1e-10, -1, exclusive=True)
ray_rgb = tf.reduce_sum(
input_tensor=tf.expand_dims(ray_sample_weights, -1) * rgb, axis=-2)
ray_alpha = tf.expand_dims(tf.reduce_sum(
input_tensor=ray_sample_weights, axis=-1),
axis=-1)
return ray_rgb, ray_alpha, ray_sample_weights
def rasterizer_bounding_box(triangle_vertices_2d,
image_width,
image_height,
name=None):
"""Compute 2D bounding boxes for 2D triangles inside an image.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
triangle_vertices_2d: A tensor of shape `[A1, ..., An, 3, 2], where the last
two dimensions represent the x, y coordinates for each vertex in a
triangle.
image_width: A scalar tensor or an `int`.
image_height: A scalar tensor or an `int`.
name: A name for this op that defaults to "rasterizer_bounding_box".
Returns:
bottom_right_corner: A tensor of shape `[A1, ..., An, 2], where the last
dimension represents the x_max, y_max of the 2D bounding boxes.
top_left_corner: A tensor of shape `[A1, ..., An, 2], where the last
dimension represents the x_min, y_min of the 2D bounding boxes.
"""
with tf.compat.v1.name_scope(
name, "rasterizer_bounding_box",
[triangle_vertices_2d, image_width, image_height]):
triangle_vertices_2d = tf.convert_to_tensor(value=triangle_vertices_2d)
image_width = tf.convert_to_tensor(value=image_width,
dtype=triangle_vertices_2d.dtype)
image_height = tf.convert_to_tensor(value=image_height,
dtype=triangle_vertices_2d.dtype)
shape.check_static(tensor=triangle_vertices_2d,
tensor_name="triangle_vertices_2d",
has_dim_equals=((-1, 2), (-2, 3)))
shape.check_static(tensor=image_width,
tensor_name="image_width",
has_rank=0)
shape.check_static(tensor=image_height,
tensor_name="image_height",
has_rank=0)
# clipping by the image size
x, y = tf.unstack(triangle_vertices_2d, axis=-1)
x = tf.clip_by_value(x, -image_width / 2, image_width / 2 - 1)
y = tf.clip_by_value(y, -image_height / 2, image_height / 2 - 1)
xy = tf.stack((x, y), axis=-1)
# calculate the bottom right and top left corners of the 2d bounding box
bottom_right_corner = tf.reduce_max(input_tensor=xy, axis=-2)
bottom_right_corner = tf.math.ceil(bottom_right_corner)
top_left_corner = tf.reduce_min(input_tensor=xy, axis=-2)
top_left_corner = tf.math.floor(top_left_corner)
return bottom_right_corner, top_left_corner
def get_bounding_box(triangle_vertices, image_width, image_height, name=None):
"""Computes 2D bounding boxes for 2D triangles inside an image.
This function rounds the estimated bounding box corners and therefore has zero
derivative everywhere with respect to the vertices. Note that this does not
prevent the rasterizer from being differentiable.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
triangle_vertices: A tensor of shape `[A1, ..., An, 3, 2], where the last
two dimensions represent the `x` and `y` coordinates of each vertex in a
triangle.
image_width: A scalar tensor or a `float`.
image_height: A scalar tensor or a `float`.
name: A name for this op that defaults to "rasterizer_get_bounding_box".
Returns:
bottom_right_corner: A tensor of shape `[A1, ..., An, 2], where the last
dimension represents the `x_max` and `y_max` of the 2D bounding boxes.
top_left_corner: A tensor of shape `[A1, ..., An, 2], where the last
dimension represents the `x_min` and `y_min` of the 2D bounding boxes.
"""
with tf.compat.v1.name_scope(
name, "rasterizer_get_bounding_box",
[triangle_vertices, image_width, image_height]):
triangle_vertices = tf.convert_to_tensor(value=triangle_vertices)
image_width = tf.convert_to_tensor(value=image_width,
dtype=triangle_vertices.dtype)
image_height = tf.convert_to_tensor(value=image_height,
dtype=triangle_vertices.dtype)
shape.check_static(tensor=triangle_vertices,
tensor_name="triangle_vertices",
has_dim_equals=((-1, 2), (-2, 3)))
shape.check_static(tensor=image_width,
tensor_name="image_width",
has_rank=0)
shape.check_static(tensor=image_height,
tensor_name="image_height",
has_rank=0)
max_clip_values = tf.stack((image_width / 2, image_height / 2),
axis=-1)
min_clip_values = -max_clip_values
clipped_vertices = tf.clip_by_value(triangle_vertices, min_clip_values,
max_clip_values)
bottom_right_corner = tf.reduce_max(input_tensor=clipped_vertices,
axis=-2)
bottom_right_corner = tf.math.ceil(bottom_right_corner)
top_left_corner = tf.reduce_min(input_tensor=clipped_vertices, axis=-2)
top_left_corner = tf.math.floor(top_left_corner)
return bottom_right_corner, top_left_corner
def rotate_zonal_harmonics(
zonal_coeffs: TensorLike,
theta: TensorLike,
phi: TensorLike,
name: str = "spherical_harmonics_rotate_zonal_harmonics"
) -> TensorLike:
"""Rotates zonal harmonics.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
zonal_coeffs: A tensor of shape `[C]` storing zonal harmonics coefficients.
theta: A tensor of shape `[A1, ..., An, 1]` storing polar angles.
phi: A tensor of shape `[A1, ..., An, 1]` storing azimuthal angles.
name: A name for this op. Defaults to
"spherical_harmonics_rotate_zonal_harmonics".
Returns:
A tensor of shape `[A1, ..., An, C*C]` storing coefficients of the rotated
harmonics.
Raises:
ValueError: If the shape of `zonal_coeffs`, `theta` or `phi` is not
supported.
"""
with tf.name_scope(name):
zonal_coeffs = tf.convert_to_tensor(value=zonal_coeffs)
theta = tf.convert_to_tensor(value=theta)
phi = tf.convert_to_tensor(value=phi)
shape.check_static(tensor=zonal_coeffs,
tensor_name="zonal_coeffs",
has_rank=1)
shape.check_static(tensor=phi,
tensor_name="phi",
has_dim_equals=(-1, 1))
shape.check_static(tensor=theta,
tensor_name="theta",
has_dim_equals=(-1, 1))
shape.compare_batch_dimensions(tensors=(theta, phi),
last_axes=-2,
tensor_names=("theta", "phi"),
broadcast_compatible=False)
tiled_zonal_coeffs = tile_zonal_coefficients(zonal_coeffs)
max_band = zonal_coeffs.shape.as_list()[-1]
l, m = generate_l_m_permutations(max_band - 1)
broadcast_shape = theta.shape.as_list()[:-1] + l.shape.as_list()
l_broadcasted = tf.broadcast_to(l, broadcast_shape)
m_broadcasted = tf.broadcast_to(m, broadcast_shape)
n_star = tf.sqrt(4.0 * np.pi /
(2.0 * tf.cast(l, dtype=theta.dtype) + 1.0))
return n_star * tiled_zonal_coeffs * evaluate_spherical_harmonics(
l_broadcasted, m_broadcasted, theta, phi)
def rotate(point, axis, angle, name=None):
r"""Rotates a 3d point using an axis-angle by applying the Rodrigues' formula.
Rotates a vector $$\mathbf{v} \in {\mathbb{R}^3}$$ into a vector
$$\mathbf{v}' \in {\mathbb{R}^3}$$ using the Rodrigues' rotation formula:
$$\mathbf{v}'=\mathbf{v}\cos(\theta)+(\mathbf{a}\times\mathbf{v})\sin(\theta)
+\mathbf{a}(\mathbf{a}\cdot\mathbf{v})(1-\cos(\theta)).$$
Note:
In the following, A1 to An are optional batch dimensions.
Args:
point: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents a 3d point to rotate.
axis: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents a normalized axis.
angle: A tensor of shape `[A1, ..., An, 1]`, where the last dimension
represents an angle.
name: A name for this op that defaults to "axis_angle_rotate".
Returns:
A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents
a 3d point.
Raises:
ValueError: If `point`, `axis`, or `angle` are of different shape or if
their respective shape is not supported.
"""
with tf.compat.v1.name_scope(name, "axis_angle_rotate",
[point, axis, angle]):
point = tf.convert_to_tensor(value=point)
axis = tf.convert_to_tensor(value=axis)
angle = tf.convert_to_tensor(value=angle)
shape.check_static(tensor=point,
tensor_name="point",
has_dim_equals=(-1, 3))
shape.check_static(tensor=axis,
tensor_name="axis",
has_dim_equals=(-1, 3))
shape.check_static(tensor=angle,
tensor_name="angle",
has_dim_equals=(-1, 1))
shape.compare_batch_dimensions(tensors=(point, axis, angle),
tensor_names=("point", "axis", "angle"),
last_axes=-2,
broadcast_compatible=True)
axis = asserts.assert_normalized(axis)
cos_angle = tf.cos(angle)
axis_dot_point = vector.dot(axis, point)
return point * cos_angle + vector.cross(axis, point) * tf.sin(
angle) + axis * axis_dot_point * (1.0 - cos_angle)
def check_valid_graph_convolution_input(data: type_alias.TensorLike,
neighbors: tf.sparse.SparseTensor,
sizes: type_alias.TensorLike):
"""Checks that the inputs are valid for graph convolution ops.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
data: A `float` tensor with shape `[A1, ..., An, V1, V2]`.
neighbors: A SparseTensor with the same type as `data` and with shape `[A1,
..., An, V1, V1]`.
sizes: An `int` tensor of shape `[A1, ..., An]`. Optional, can be `None`.
Raises:
TypeError: if the input types are invalid.
ValueError: if the input dimensions are invalid.
"""
if not data.dtype.is_floating:
raise TypeError("'data' must have a float type.")
if neighbors.dtype != data.dtype:
raise TypeError("'neighbors' and 'data' must have the same type.")
if sizes is not None and not sizes.dtype.is_integer:
raise TypeError("'sizes' must have an integer type.")
if not isinstance(neighbors, tf.sparse.SparseTensor):
raise ValueError("'neighbors' must be a SparseTensor.")
data_ndims = data.shape.ndims
shape.check_static(tensor=data,
tensor_name="data",
has_rank_greater_than=1)
shape.check_static(tensor=neighbors,
tensor_name="neighbors",
has_rank=data_ndims)
if not _is_dynamic_shape(tensors=(data, neighbors)):
shape.compare_dimensions(tensors=(data, neighbors, neighbors),
tensor_names=("data", "neighbors",
"neighbors"),
axes=(-2, -2, -1))
if sizes is None:
shape.compare_batch_dimensions(tensors=(data, neighbors),
tensor_names=("data", "neighbors"),
last_axes=-3,
broadcast_compatible=False)
else:
shape.check_static(tensor=sizes,
tensor_name="sizes",
has_rank=data_ndims - 2)
shape.compare_batch_dimensions(tensors=(data, neighbors, sizes),
tensor_names=("data", "neighbors",
"sizes"),
last_axes=(-3, -3, -1),
broadcast_compatible=False)
def partition_sums_2d(data, group_ids, row_weights=None, name=None):
"""Sum over subsets of rows in a 2-D tensor.
Args:
data: 2-D tensor with shape `[D1, D2]`.
group_ids: 1-D `int` tensor with shape `[D1]`.
row_weights: 1-D tensor with shape `[D1]`. Can be `None`.
name: A name for this op. Defaults to 'utils_partition_sums_2d'.
Returns:
A 2-D tensor with shape `[max(group_ids) + 1, D2]` where
`output[i, :] = sum(data[j, :] * weight[j] * 1(group_ids[j] == i)), 1(.)`
is the indicator function.
Raises:
ValueError: if the inputs have invalid dimensions or types.
"""
with tf.compat.v1.name_scope(name, "utils_partition_sums_2d",
[data, group_ids, row_weights]):
data = tf.convert_to_tensor(value=data)
group_ids = tf.convert_to_tensor(value=group_ids)
if not group_ids.dtype.is_integer:
raise TypeError("'group_ids' must be an integer tensor.")
elif group_ids.dtype != tf.int64:
group_ids = tf.cast(group_ids, dtype=tf.int64)
if row_weights is None:
row_weights = tf.ones_like(group_ids, dtype=data.dtype)
else:
row_weights = tf.convert_to_tensor(value=row_weights)
if row_weights.dtype != data.dtype:
raise TypeError(
"'data' and 'row_weights' must have the same type.")
shape.check_static(tensor=data, tensor_name="data", has_rank=2)
shape.check_static(tensor=group_ids,
tensor_name="group_ids",
has_rank=1)
shape.check_static(tensor=row_weights,
tensor_name="row_weights",
has_rank=1)
shape.compare_dimensions(tensors=(data, group_ids, row_weights),
tensor_names=("data", "group_ids",
"row_weights"),
axes=0)
num_rows = tf.size(input=group_ids, out_type=tf.int64)
sparse_indices = tf.stack((group_ids, tf.range(num_rows)), axis=1)
out_shape = (tf.reduce_max(input_tensor=group_ids) + 1, num_rows)
sparse = tf.SparseTensor(sparse_indices,
row_weights,
dense_shape=out_shape)
return tf.sparse.sparse_dense_matmul(sparse, data)
def normal(v0: type_alias.TensorLike,
v1: type_alias.TensorLike,
v2: type_alias.TensorLike,
clockwise: bool = False,
normalize: bool = True,
name: str = "triangle_normal") -> tf.Tensor:
"""Computes face normals (triangles).
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
Args:
v0: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents the first vertex of a triangle.
v1: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents the second vertex of a triangle.
v2: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents the third vertex of a triangle.
clockwise: Winding order to determine front-facing triangles.
normalize: A `bool` indicating whether output normals should be normalized
by the function.
name: A name for this op. Defaults to "triangle_normal".
Returns:
A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents
a normalized vector.
Raises:
ValueError: If the shape of `v0`, `v1`, or `v2` is not supported.
"""
with tf.name_scope(name):
v0 = tf.convert_to_tensor(value=v0)
v1 = tf.convert_to_tensor(value=v1)
v2 = tf.convert_to_tensor(value=v2)
shape.check_static(tensor=v0, tensor_name="v0", has_dim_equals=(-1, 3))
shape.check_static(tensor=v1, tensor_name="v1", has_dim_equals=(-1, 3))
shape.check_static(tensor=v2, tensor_name="v2", has_dim_equals=(-1, 3))
shape.compare_batch_dimensions(tensors=(v0, v1, v2),
last_axes=-2,
broadcast_compatible=True)
normal_vector = vector.cross(v1 - v0, v2 - v0, axis=-1)
normal_vector = asserts.assert_nonzero_norm(normal_vector)
if not clockwise:
normal_vector *= -1.0
if normalize:
return tf.nn.l2_normalize(normal_vector, axis=-1)
return normal_vector
def compute_density(density_values, distances, name=None):
"""Renders the density values (alpha) for points along a ray, as described in ["NeRF Representing Scenes as Neural Radiance Fields for View Synthesis"](https://github.com/bmild/nerf).
Note:
In the following, A1 to An are optional batch dimensions.
Args:
density_values: A tensor of shape `[A1, ..., An, N, 1]`,
where N are the samples on the ray.
distances: A tensor of shape `[A1, ..., An, N]` containing the distances
between the samples, where N are the samples on the ray.
name: A name for this op. Defaults to "ray_radiance".
Returns:
A tensor of shape `[A1, ..., An, 1]` for the estimated density values,
and a tensor of shape `[A1, ..., An, N]` for the sample weights.
"""
with tf.compat.v1.name_scope(name, "ray_density",
[density_values, distances]):
density_values = tf.convert_to_tensor(value=density_values)
distances = tf.convert_to_tensor(value=distances)
distances = tf.expand_dims(distances, -1)
shape.check_static(tensor=density_values,
tensor_name="density_values",
has_dim_equals=(-1, 1))
shape.check_static(tensor=density_values,
tensor_name="density_values",
has_rank_greater_than=1)
shape.check_static(tensor=distances,
tensor_name="distances",
has_rank_greater_than=1)
shape.compare_batch_dimensions(tensors=(density_values, distances),
tensor_names=("density_values",
"dists"),
last_axes=-3,
broadcast_compatible=True)
shape.compare_dimensions(tensors=(density_values, distances),
tensor_names=("density_values", "dists"),
axes=-2)
alpha = 1. - tf.exp(-density_values * distances)
alpha = tf.squeeze(alpha, -1)
ray_sample_weights = alpha * tf.math.cumprod(
1. - alpha + 1e-10, -1, exclusive=True)
ray_alpha = tf.expand_dims(tf.reduce_sum(ray_sample_weights, -1),
axis=-1)
return ray_alpha, ray_sample_weights
def area(v0: type_alias.TensorLike,
v1: type_alias.TensorLike,
v2: type_alias.TensorLike,
name: str = "triangle_area") -> tf.Tensor:
"""Computes triangle areas.
Note: Computed triangle area = 0.5 * | e1 x e2 | where e1 and e2 are edges
of triangle. A degenerate triangle will return 0 area, whereas the normal
for a degenerate triangle is not defined.
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
Args:
v0: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents the first vertex of a triangle.
v1: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents the second vertex of a triangle.
v2: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents the third vertex of a triangle.
name: A name for this op. Defaults to "triangle_area".
Returns:
A tensor of shape `[A1, ..., An, 1]`, where the last dimension represents
a normalized vector.
"""
with tf.name_scope(name):
v0 = tf.convert_to_tensor(value=v0)
v1 = tf.convert_to_tensor(value=v1)
v2 = tf.convert_to_tensor(value=v2)
shape.check_static(tensor=v0, tensor_name="v0", has_dim_equals=(-1, 3))
shape.check_static(tensor=v1, tensor_name="v1", has_dim_equals=(-1, 3))
shape.check_static(tensor=v2, tensor_name="v2", has_dim_equals=(-1, 3))
shape.compare_batch_dimensions(tensors=(v0, v1, v2),
last_axes=-2,
broadcast_compatible=True)
normals = vector.cross(v1 - v0, v2 - v0, axis=-1)
return 0.5 * tf.linalg.norm(tensor=normals, axis=-1, keepdims=True)
def reconstruct(image, coeff_mul, coeff_add, name=None):
"""Reconstruct the matte from the image using the linear coefficients.
Reconstruct the matte from the image using the linear coefficients (a, b)
returned by the linear_coefficients function.
Args:
image: A tensor of shape `[B, H, W, C]` .
coeff_mul: A tensor of shape `[B, H, W, C]` representing the multiplicative
part of the linear coefficients.
coeff_add: A tensor of shape `[B, H, W, 1]` representing the additive part
of the linear coefficients.
name: A name for this op. Defaults to "matting_reconstruct".
Returns:
A tensor of shape `[B, H, W, 1]` containing the mattes.
Raises:
ValueError: If `image`, `coeff_mul`, or `coeff_add` are not of rank 4. If
the last dimension of `coeff_add` is not 1. If the batch dimensions of
`image`, `coeff_mul`, and `coeff_add` do not match.
"""
with tf.compat.v1.name_scope(name, "matting_reconstruct",
[image, coeff_mul, coeff_add]):
image = tf.convert_to_tensor(value=image)
coeff_mul = tf.convert_to_tensor(value=coeff_mul)
coeff_add = tf.convert_to_tensor(value=coeff_add)
shape.check_static(image, has_rank=4)
shape.check_static(coeff_mul, has_rank=4)
shape.check_static(coeff_add, has_rank=4, has_dim_equals=(-1, 1))
shape.compare_batch_dimensions(tensors=(image, coeff_mul),
last_axes=-1,
broadcast_compatible=False)
shape.compare_batch_dimensions(tensors=(image, coeff_add),
last_axes=-2,
broadcast_compatible=False)
return tfg_vector.dot(coeff_mul, image) + coeff_add
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,
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 area_weighted_random_sample_triangle_mesh(vertex_attributes,
faces,
num_samples,
vertex_positions=None,
seed=None,
stateless=False,
name=None):
"""Performs a face area weighted random sampling of a tri mesh.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
vertex_attributes: A `float` tensor of shape `[A1, ..., An, V, D]`, where V
is the number of vertices, and D is dimensionality of a feature defined on
each vertex. If `vertex_positions` is not provided, then first 3
dimensions of `vertex_attributes` denote the vertex positions.
faces: A `int` tensor of shape `[A1, ..., An, F, 3]`, where F is the number
of faces.
num_samples: An `int` scalar denoting number of samples to be drawn from
each mesh.
vertex_positions: An optional `float` tensor of shape `[A1, ..., An, V, 3]`,
where V is the number of vertices. If None, then vertex_attributes[...,
:3] is used as vertex positions.
seed: Optional random seed.
stateless: Optional flag to use stateless random sampler. If stateless=True,
then seed must be provided as shape `[2]` int tensor. Stateless random
sampling is useful for testing to generate same sequence across calls.
name: Name for op. Defaults to "area_weighted_random_sample_triangle_mesh".
Returns:
sample_pts: A `float` tensor of shape `[A1, ..., An, num_samples, D]`,
where D is dimensionality of each sampled point.
sample_face_indices: A `int` tensor of shape `[A1, ..., An, num_samples]`.
"""
with tf.compat.v1.name_scope(name,
"area_weighted_random_sample_triangle_mesh"):
faces = tf.convert_to_tensor(value=faces)
vertex_attributes = tf.convert_to_tensor(value=vertex_attributes)
num_samples = tf.convert_to_tensor(value=num_samples)
shape.check_static(
tensor=vertex_attributes,
tensor_name="vertex_attributes",
has_rank_greater_than=1)
shape.check_static(
tensor=vertex_attributes,
tensor_name="vertex_attributes",
has_dim_greater_than=(-1, 2))
if vertex_positions is not None:
vertex_positions = tf.convert_to_tensor(value=vertex_positions)
else:
vertex_positions = vertex_attributes[..., :3]
shape.check_static(
tensor=vertex_positions,
tensor_name="vertex_positions",
has_rank_greater_than=1)
shape.check_static(
tensor=vertex_positions,
tensor_name="vertex_positions",
has_dim_equals=(-1, 3))
triangle_vertex_positions = normals.gather_faces(vertex_positions, faces)
triangle_areas = triangle_area(triangle_vertex_positions[..., 0, :],
triangle_vertex_positions[..., 1, :],
triangle_vertex_positions[..., 2, :])
return weighted_random_sample_triangle_mesh(
vertex_attributes,
faces,
num_samples,
face_weights=triangle_areas,
seed=seed,
stateless=stateless)
def weighted_random_sample_triangle_mesh(vertex_attributes,
faces,
num_samples,
face_weights,
seed=None,
stateless=False,
name=None):
"""Performs a face probability weighted random sampling of a tri mesh.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
vertex_attributes: A `float` tensor of shape `[A1, ..., An, V, D]`, where V
is the number of vertices, and D is dimensionality of each vertex.
faces: A `int` tensor of shape `[A1, ..., An, F, 3]`, where F is the number
of faces.
num_samples: A `int` 0-D tensor denoting number of samples to be drawn from
each mesh.
face_weights: A `float` tensor of shape ``[A1, ..., An, F]`, denoting
unnormalized sampling probability of each face, where F is the number of
faces.
seed: Optional random seed.
stateless: Optional flag to use stateless random sampler. If stateless=True,
then seed must be provided as shape `[2]` int tensor. Stateless random
sampling is useful for testing to generate same sequence across calls.
name: Name for op. Defaults to "weighted_random_sample_triangle_mesh".
Returns:
sample_points: A `float` tensor of shape `[A1, ..., An, num_samples, D]`,
where D is dimensionality of each sampled point.
sample_face_indices: A `int` tensor of shape `[A1, ..., An, num_samples]`.
"""
with tf.compat.v1.name_scope(name, "weighted_random_sample_triangle_mesh"):
faces = tf.convert_to_tensor(value=faces)
vertex_attributes = tf.convert_to_tensor(value=vertex_attributes)
face_weights = tf.convert_to_tensor(value=face_weights)
num_samples = tf.convert_to_tensor(value=num_samples)
shape.check_static(
tensor=vertex_attributes,
tensor_name="vertex_attributes",
has_rank_greater_than=1)
shape.check_static(
tensor=faces, tensor_name="faces", has_rank_greater_than=1)
shape.check_static(
tensor=face_weights,
tensor_name="face_weights",
has_rank_greater_than=0)
shape.compare_batch_dimensions(
tensors=(faces, face_weights),
last_axes=(-2, -1),
tensor_names=("faces", "face_weights"),
broadcast_compatible=False)
shape.compare_batch_dimensions(
tensors=(vertex_attributes, faces, face_weights),
last_axes=(-3, -3, -2),
tensor_names=("vertex_attributes", "faces", "face_weights"),
broadcast_compatible=False)
asserts.assert_all_above(face_weights, 0)
batch_dims = faces.shape.ndims - 2
batch_shape = faces.shape.as_list()[:-2]
sample_shape = tf.concat(
(batch_shape, tf.convert_to_tensor(
value=(num_samples,), dtype=tf.int32)),
axis=0)
sample_face_indices = generate_random_face_indices(
num_samples, face_weights, seed=seed, stateless=stateless)
sample_vertex_indices = tf.gather(
faces, sample_face_indices, batch_dims=batch_dims)
sample_vertices = tf.gather(
vertex_attributes, sample_vertex_indices, batch_dims=batch_dims)
barycentric = generate_random_barycentric_coordinates(
sample_shape,
dtype=vertex_attributes.dtype,
seed=seed,
stateless=stateless)
barycentric = tf.expand_dims(barycentric, axis=-1)
sample_points = tf.math.multiply(sample_vertices, barycentric)
sample_points = tf.reduce_sum(input_tensor=sample_points, axis=-2)
return sample_points, sample_face_indices
def sample(image: type_alias.TensorLike,
warp: type_alias.TensorLike,
resampling_type: ResamplingType = ResamplingType.BILINEAR,
border_type: BorderType = BorderType.ZERO,
pixel_type: PixelType = PixelType.HALF_INTEGER,
name: Optional[str] = "sample") -> tf.Tensor:
"""Samples an image at user defined coordinates.
Note:
The warp maps target to source. In the following, A1 to An are optional
batch dimensions.
Args:
image: A tensor of shape `[B, H_i, W_i, C]`, where `B` is the batch size,
`H_i` the height of the image, `W_i` the width of the image, and `C` the
number of channels of the image.
warp: A tensor of shape `[B, A_1, ..., A_n, 2]` containing the x and y
coordinates at which sampling will be performed. The last dimension must
be 2, representing the (x, y) coordinate where x is the index for width
and y is the index for height.
resampling_type: Resampling mode. Supported values are
`ResamplingType.NEAREST` and `ResamplingType.BILINEAR`.
border_type: Border mode. Supported values are `BorderType.ZERO` and
`BorderType.DUPLICATE`.
pixel_type: Pixel mode. Supported values are `PixelType.INTEGER` and
`PixelType.HALF_INTEGER`.
name: A name for this op. Defaults to "sample".
Returns:
Tensor of sampled values from `image`. The output tensor shape
is `[B, A_1, ..., A_n, C]`.
Raises:
ValueError: If `image` has rank != 4. If `warp` has rank < 2 or its last
dimension is not 2. If `image` and `warp` batch dimension does not match.
"""
with tf.name_scope(name):
image = tf.convert_to_tensor(value=image, name="image")
warp = tf.convert_to_tensor(value=warp, name="warp")
shape.check_static(image, tensor_name="image", has_rank=4)
shape.check_static(
warp,
tensor_name="warp",
has_rank_greater_than=1,
has_dim_equals=(-1, 2))
shape.compare_batch_dimensions(
tensors=(image, warp), last_axes=0, broadcast_compatible=False)
if pixel_type == PixelType.HALF_INTEGER:
warp -= 0.5
if resampling_type == ResamplingType.NEAREST:
warp = tf.math.round(warp)
if border_type == BorderType.DUPLICATE:
image_size = tf.cast(tf.shape(input=image)[1:3], dtype=warp.dtype)
height, width = tf.unstack(image_size, axis=-1)
warp_x, warp_y = tf.unstack(warp, axis=-1)
warp_x = tf.clip_by_value(warp_x, 0.0, width - 1.0)
warp_y = tf.clip_by_value(warp_y, 0.0, height - 1.0)
warp = tf.stack((warp_x, warp_y), axis=-1)
return tfa_image.resampler(image, warp)
def perspective_transform(
image: type_alias.TensorLike,
transform_matrix: type_alias.TensorLike,
output_shape: Optional[type_alias.TensorLike] = None,
resampling_type: ResamplingType = ResamplingType.BILINEAR,
border_type: BorderType = BorderType.ZERO,
pixel_type: PixelType = PixelType.HALF_INTEGER,
name: Optional[str] = "perspective_transform",
) -> tf.Tensor:
"""Applies a projective transformation to an image.
The projective transformation is represented by a 3 x 3 matrix
[[a0, a1, a2], [b0, b1, b2], [c0, c1, c2]], mapping a point `[x, y]` to a
transformed point
`[x', y'] = [(a0 x + a1 y + a2) / k, (b0 x + b1 y + b2) / k]`, where
`k = c0 x + c1 y + c2`.
Note:
The transformation matrix maps target to source by transforming output
points to input points.
Args:
image: A tensor of shape `[B, H_i, W_i, C]`, where `B` is the batch size,
`H_i` the height of the image, `W_i` the width of the image, and `C` the
number of channels of the image.
transform_matrix: A tensor of shape `[B, 3, 3]` containing projective
transform matrices. The transformation maps target to source by
transforming output points to input points.
output_shape: The heigh `H_o` and width `W_o` output dimensions after the
transform. If None, output is the same size as input image.
resampling_type: Resampling mode. Supported values are
`ResamplingType.NEAREST` and `ResamplingType.BILINEAR`.
border_type: Border mode. Supported values are `BorderType.ZERO` and
`BorderType.DUPLICATE`.
pixel_type: Pixel mode. Supported values are `PixelType.INTEGER` and
`PixelType.HALF_INTEGER`.
name: A name for this op. Defaults to "perspective_transform".
Returns:
A tensor of shape `[B, H_o, W_o, C]` containing transformed images.
Raises:
ValueError: If `image` has rank != 4. If `transform_matrix` has rank < 3 or
its last two dimensions are not 3. If `image` and `transform_matrix` batch
dimension does not match.
"""
with tf.name_scope(name):
image = tf.convert_to_tensor(value=image, name="image")
transform_matrix = tf.convert_to_tensor(
value=transform_matrix, name="transform_matrix")
output_shape = tf.shape(
input=image)[-3:-1] if output_shape is None else tf.convert_to_tensor(
value=output_shape, name="output_shape")
shape.check_static(image, tensor_name="image", has_rank=4)
shape.check_static(
transform_matrix,
tensor_name="transform_matrix",
has_rank=3,
has_dim_equals=((-1, 3), (-2, 3)))
shape.compare_batch_dimensions(
tensors=(image, transform_matrix),
last_axes=0,
broadcast_compatible=False)
dtype = image.dtype
zero = tf.cast(0.0, dtype)
height, width = tf.unstack(output_shape, axis=-1)
warp = grid.generate(
starts=(zero, zero),
stops=(tf.cast(width, dtype) - 1.0, tf.cast(height, dtype) - 1.0),
nums=(width, height))
warp = tf.transpose(a=warp, perm=[1, 0, 2])
if pixel_type == PixelType.HALF_INTEGER:
warp += 0.5
padding = [[0, 0] for _ in range(warp.shape.ndims)]
padding[-1][-1] = 1
warp = tf.pad(
tensor=warp, paddings=padding, mode="CONSTANT", constant_values=1.0)
warp = warp[..., tf.newaxis]
transform_matrix = transform_matrix[:, tf.newaxis, tf.newaxis, ...]
warp = tf.linalg.matmul(transform_matrix, warp)
warp = warp[..., 0:2, 0] / warp[..., 2, :]
return sample(image, warp, resampling_type, border_type, pixel_type)
def energy(vertices_rest_pose,
vertices_deformed_pose,
quaternions,
edges,
vertex_weight=None,
edge_weight=None,
conformal_energy=True,
aggregate_loss=True,
name=None):
"""Estimates an As Conformal As Possible (ACAP) fitting energy.
For a given mesh in rest pose, this function evaluates a variant of the ACAP
[1] fitting energy for a batch of deformed meshes. The vertex weights and edge
weights are defined on the rest pose.
The method implemented here is similar to [2], but with an added free variable
capturing a scale factor per vertex.
[1]: Yusuke Yoshiyasu, Wan-Chun Ma, Eiichi Yoshida, and Fumio Kanehiro.
"As-Conformal-As-Possible Surface Registration." Computer Graphics Forum. Vol.
33. No. 5. 2014.</br>
[2]: Olga Sorkine, and Marc Alexa.
"As-rigid-as-possible surface modeling". Symposium on Geometry Processing.
Vol. 4. 2007.
Note:
In the description of the arguments, V corresponds to
the number of vertices in the mesh, and E to the number of edges in this
mesh.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
vertices_rest_pose: A tensor of shape `[V, 3]` containing the position of
all the vertices of the mesh in rest pose.
vertices_deformed_pose: A tensor of shape `[A1, ..., An, V, 3]` containing
the position of all the vertices of the mesh in deformed pose.
quaternions: A tensor of shape `[A1, ..., An, V, 4]` defining a rigid
transformation to apply to each vertex of the rest pose. See Section 2
from [1] for further details.
edges: A tensor of shape `[E, 2]` defining indices of vertices that are
connected by an edge.
vertex_weight: An optional tensor of shape `[V]` defining the weight
associated with each vertex. Defaults to a tensor of ones.
edge_weight: A tensor of shape `[E]` defining the weight of edges. Common
choices for these weights include uniform weighting, and cotangent
weights. Defaults to a tensor of ones.
conformal_energy: A `bool` indicating whether each vertex is associated with
a scale factor or not. If this parameter is True, scaling information must
be encoded in the norm of `quaternions`. If this parameter is False, this
function implements the energy described in [2].
aggregate_loss: A `bool` defining whether the returned loss should be an
aggregate measure. When True, the mean squared error is returned. When
False, returns two losses for every edge of the mesh.
name: A name for this op. Defaults to "as_conformal_as_possible_energy".
Returns:
When aggregate_loss is `True`, returns a tensor of shape `[A1, ..., An]`
containing the ACAP energies. When aggregate_loss is `False`, returns a
tensor of shape `[A1, ..., An, 2*E]` containing each term of the summation
described in the equation 7 of [2].
Raises:
ValueError: if the shape of `vertices_rest_pose`, `vertices_deformed_pose`,
`quaternions`, `edges`, `vertex_weight`, or `edge_weight` is not supported.
"""
with tf.compat.v1.name_scope(name, "as_conformal_as_possible_energy", [
vertices_rest_pose, vertices_deformed_pose, quaternions, edges,
conformal_energy, vertex_weight, edge_weight
]):
vertices_rest_pose = tf.convert_to_tensor(value=vertices_rest_pose)
vertices_deformed_pose = tf.convert_to_tensor(
value=vertices_deformed_pose)
quaternions = tf.convert_to_tensor(value=quaternions)
edges = tf.convert_to_tensor(value=edges)
if vertex_weight is not None:
vertex_weight = tf.convert_to_tensor(value=vertex_weight)
if edge_weight is not None:
edge_weight = tf.convert_to_tensor(value=edge_weight)
shape.check_static(tensor=vertices_rest_pose,
tensor_name="vertices_rest_pose",
has_rank=2,
has_dim_equals=(-1, 3))
shape.check_static(tensor=vertices_deformed_pose,
tensor_name="vertices_deformed_pose",
has_rank_greater_than=1,
has_dim_equals=(-1, 3))
shape.check_static(tensor=quaternions,
tensor_name="quaternions",
has_rank_greater_than=1,
has_dim_equals=(-1, 4))
shape.compare_batch_dimensions(tensors=(vertices_deformed_pose,
quaternions),
last_axes=(-3, -3),
broadcast_compatible=False)
shape.check_static(tensor=edges,
tensor_name="edges",
has_rank=2,
has_dim_equals=(-1, 2))
tensors_with_vertices = [
vertices_rest_pose, vertices_deformed_pose, quaternions
]
names_with_vertices = [
"vertices_rest_pose", "vertices_deformed_pose", "quaternions"
]
axes_with_vertices = [-2, -2, -2]
if vertex_weight is not None:
shape.check_static(tensor=vertex_weight,
tensor_name="vertex_weight",
has_rank=1)
tensors_with_vertices.append(vertex_weight)
names_with_vertices.append("vertex_weight")
axes_with_vertices.append(0)
shape.compare_dimensions(tensors=tensors_with_vertices,
axes=axes_with_vertices,
tensor_names=names_with_vertices)
if edge_weight is not None:
shape.check_static(tensor=edge_weight,
tensor_name="edge_weight",
has_rank=1)
shape.compare_dimensions(tensors=(edges, edge_weight),
axes=(0, 0),
tensor_names=("edges", "edge_weight"))
if not conformal_energy:
quaternions = quaternion.normalize(quaternions)
# Extracts the indices of vertices.
indices_i, indices_j = tf.unstack(edges, axis=-1)
# Extracts the vertices we need per term.
vertices_i_rest = tf.gather(vertices_rest_pose, indices_i, axis=-2)
vertices_j_rest = tf.gather(vertices_rest_pose, indices_j, axis=-2)
vertices_i_deformed = tf.gather(vertices_deformed_pose,
indices_i,
axis=-2)
vertices_j_deformed = tf.gather(vertices_deformed_pose,
indices_j,
axis=-2)
# Extracts the weights we need per term.
weights_shape = vertices_i_rest.shape.as_list()[-2]
if vertex_weight is not None:
weight_i = tf.gather(vertex_weight, indices_i)
weight_j = tf.gather(vertex_weight, indices_j)
else:
weight_i = weight_j = tf.ones(weights_shape,
dtype=vertices_rest_pose.dtype)
weight_i = tf.expand_dims(weight_i, axis=-1)
weight_j = tf.expand_dims(weight_j, axis=-1)
if edge_weight is not None:
weight_ij = edge_weight
else:
weight_ij = tf.ones(weights_shape, dtype=vertices_rest_pose.dtype)
weight_ij = tf.expand_dims(weight_ij, axis=-1)
# Extracts the rotation we need per term.
quaternion_i = tf.gather(quaternions, indices_i, axis=-2)
quaternion_j = tf.gather(quaternions, indices_j, axis=-2)
# Computes the energy.
deformed_ij = vertices_i_deformed - vertices_j_deformed
rotated_rest_ij = quaternion.rotate(
(vertices_i_rest - vertices_j_rest), quaternion_i)
energy_ij = weight_i * weight_ij * (deformed_ij - rotated_rest_ij)
deformed_ji = vertices_j_deformed - vertices_i_deformed
rotated_rest_ji = quaternion.rotate(
(vertices_j_rest - vertices_i_rest), quaternion_j)
energy_ji = weight_j * weight_ij * (deformed_ji - rotated_rest_ji)
energy_ij_squared = vector.dot(energy_ij, energy_ij, keepdims=False)
energy_ji_squared = vector.dot(energy_ji, energy_ji, keepdims=False)
if aggregate_loss:
average_energy_ij = tf.reduce_mean(input_tensor=energy_ij_squared,
axis=-1)
average_energy_ji = tf.reduce_mean(input_tensor=energy_ji_squared,
axis=-1)
return (average_energy_ij + average_energy_ji) / 2.0
return tf.concat((energy_ij_squared, energy_ji_squared), axis=-1)
