def project(point_3d, focal, principal_point, name=None):
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' = \frac{f_x}{z}x + c_x, & y' = \frac{f_y}{z}y + c_y,
\end{matrix}
$$
where \\((f_x, f_y)\\) is the focal length and \\((c_x, c_y)\\) the principal
point.
Note:
In the following, A1 to An are optional batch dimensions that must be
broadcast compatible.
Args:
point_3d: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents a 3d point to project.
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.
name: A name for this op that defaults to "perspective_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`, `focal`, or `principal_point` is not
supported.
"""
with tf.compat.v1.name_scope(name, "perspective_project",
[point_3d, focal, principal_point]):
point_3d = tf.convert_to_tensor(value=point_3d)
focal = tf.convert_to_tensor(value=focal)
principal_point = tf.convert_to_tensor(value=principal_point)
shape.check_static(
tensor=point_3d, tensor_name="point_3d", has_dim_equals=(-1, 3))
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.compare_batch_dimensions(
tensors=(point_3d, focal, principal_point),
tensor_names=("point_3d", "focal", "principal_point"),
last_axes=-2,
broadcast_compatible=True)
point_2d, depth = tf.split(point_3d, (2, 1), axis=-1)
point_2d *= safe_ops.safe_signed_div(focal, depth)
point_2d += principal_point
return point_2d
def between_two_vectors_3d(vector1, vector2, name=None):
"""Computes quaternion over the shortest arc between two vectors.
Result quaternion describes shortest geodesic rotation from
vector1 to vector2.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
vector1: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents the first vector.
vector2: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents the second vector.
name: A name for this op that defaults to
"quaternion_between_two_vectors_3d".
Returns:
A tensor of shape `[A1, ..., An, 4]`, where the last dimension represents
a normalized quaternion.
Raises:
ValueError: If the shape of `vector1` or `vector2` is not supported.
"""
with tf.compat.v1.name_scope(name, "quaternion_between_two_vectors_3d",
[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=(-1, 3))
shape.check_static(
tensor=vector2, tensor_name="vector2", has_dim_equals=(-1, 3))
shape.compare_batch_dimensions(
tensors=(vector1, vector2), last_axes=-2, broadcast_compatible=True)
# Make sure that we are dealing with unit vectors.
vector1 = tf.nn.l2_normalize(vector1, axis=-1)
vector2 = tf.nn.l2_normalize(vector2, axis=-1)
cos_theta = vector.dot(vector1, vector2)
real_part = 1.0 + cos_theta
axis = vector.cross(vector1, vector2)
# Compute arbitrary antiparallel axes to rotate around in case of opposite
# vectors.
x, y, z = tf.split(vector1, (1, 1, 1), axis=-1)
x_bigger_z = tf.abs(x) > tf.abs(z)
x_bigger_z = tf.concat([x_bigger_z] * 3, axis=-1)
antiparallel_axis = tf.compat.v1.where(
x_bigger_z, tf.concat((-y, x, tf.zeros_like(z)), axis=-1),
tf.concat((tf.zeros_like(x), -z, y), axis=-1))
# Compute rotation between two vectors.
is_antiparallel = real_part < 1e-6
is_antiparallel = tf.concat([is_antiparallel] * 4, axis=-1)
rot = tf.compat.v1.where(
is_antiparallel,
tf.concat((antiparallel_axis, tf.zeros_like(real_part)), axis=-1),
tf.concat((axis, real_part), axis=-1))
return tf.nn.l2_normalize(rot, axis=-1)
def stratified_geomspace_1d(near: TensorLike,
far: TensorLike,
num_samples: int,
name="stratified_geomspace_1d") -> tf.Tensor:
"""Stratified sampling based on evenly-spaced bins on a geometric progression.
Args:
near: A tensor of shape `[A1, ... An]` containing the starting points of
the sampling interval.
far: A tensor of shape `[A1, ... An]` containing the ending points of
the sampling interval.
num_samples: The number M of points to be sampled.
name: A name for this op that defaults to "stratified".
Returns:
A tensor of shape `[A1, ..., An, M]` indicating the M points on the ray
"""
with tf.name_scope(name):
near = tf.convert_to_tensor(near)
far = tf.convert_to_tensor(far)
shape.compare_batch_dimensions(tensors=(tf.expand_dims(near, axis=-1),
tf.expand_dims(far, axis=-1)),
tensor_names=("near", "far"),
last_axes=-1,
broadcast_compatible=True)
bin_borders = geomspace_1d(near, far, num_samples + 1)
bin_below = bin_borders[..., :-1]
bin_above = bin_borders[..., 1:]
target_shape = tf.concat([tf.shape(near), [num_samples]], axis=-1)
random_point_in_bin = tf.random.uniform(target_shape)
z_values = bin_below + (bin_above - bin_below) * random_point_in_bin
return z_values
def uniform_1d(near: TensorLike,
far: TensorLike,
num_samples: int,
name="uniform_1d") -> tf.Tensor:
"""Sample uniformly numbers on an interval, with the numbers being sorted.
Args:
near: A tensor of shape `[A1, ... An]` containing the starting points of
the sampling interval.
far: A tensor of shape `[A1, ... An]` containing the ending points of
the sampling interval.
num_samples: The number M of points to be sampled.
name: A name for this op that defaults to "uniform_1d".
Returns:
A tensor of shape `[A1, ..., An, M]` indicating the M points on the ray
"""
with tf.name_scope(name):
near = tf.convert_to_tensor(near)
far = tf.convert_to_tensor(far)
shape.compare_batch_dimensions(tensors=(tf.expand_dims(near, axis=-1),
tf.expand_dims(far, axis=-1)),
tensor_names=("near", "far"),
last_axes=-1,
broadcast_compatible=True)
target_shape = tf.concat([tf.shape(near), [num_samples]], axis=-1)
random_samples = tf.random.uniform(target_shape,
minval=tf.expand_dims(near, -1),
maxval=tf.expand_dims(far, -1))
return tf.sort(random_samples, axis=-1)
def logspace_1d(near: TensorLike,
far: TensorLike,
num_samples: int,
base: float = 10.0,
name="logspace_1d") -> tf.Tensor:
"""Sample evenly spaced numbers from an interval on a log scale.
Args:
near: A tensor of shape `[A1, ... An]` containing the starting points of
the sampling interval.
far: A tensor of shape `[A1, ... An]` containing the ending points of
the sampling interval.
num_samples: The number M of points to be sampled.
base: The logarithmic base.
name: A name for this op that defaults to "logspace_1d".
Returns:
A tensor of shape `[A1, ..., An, M]` indicating the M points on the ray
"""
with tf.name_scope(name):
near = tf.convert_to_tensor(near)
far = tf.convert_to_tensor(far)
shape.compare_batch_dimensions(tensors=(tf.expand_dims(near, axis=-1),
tf.expand_dims(far, axis=-1)),
tensor_names=("near", "far"),
last_axes=-1,
broadcast_compatible=True)
linspace = tf.linspace(near, far, num_samples, axis=-1)
return tf.math.pow(base, linspace)
def geomspace_1d(near: TensorLike,
far: TensorLike,
num_samples: int,
name="geomspace_1d") -> tf.Tensor:
"""Sample evenly spaced numbers on a geometric progression (log scale).
Args:
near: A tensor of shape `[A1, ... An]` containing the starting points of
the sampling interval.
far: A tensor of shape `[A1, ... An]` containing the ending points of
the sampling interval.
num_samples: The number M of points to be sampled.
name: A name for this op that defaults to "geomspace_1d".
Returns:
A tensor of shape `[A1, ..., An, M]` indicating the M points on the ray
"""
with tf.name_scope(name):
near = tf.convert_to_tensor(near)
far = tf.convert_to_tensor(far)
shape.compare_batch_dimensions(tensors=(tf.expand_dims(near, axis=-1),
tf.expand_dims(far, axis=-1)),
tensor_names=("near", "far"),
last_axes=-1,
broadcast_compatible=True)
return logspace_1d(_log10(near), _log10(far), num_samples)
def camera_rays_from_extrinsics(rays, rotation_matrix, translation_vector):
"""Transform the rays from a camera located at (0, 0, 0) to ray origins and directions for a camera with given extrinsics.
Args:
rays: A tensor of shape `[A1, ..., An, N, 3]` where N is the number of rays.
rotation_matrix: A tensor of shape `[A1, ..., An, 3, 3]`.
translation_vector: A tensor of shape `[A1, ..., An, 3, 1]`.
Returns:
A tensor of shape `[A1, ..., An, N, 3]` representing the ray origin and
a tensor of shape `[A1, ..., An, N, 3]` representing the ray direction.
"""
shape.check_static(tensor=rays,
tensor_name="pixels",
has_rank_greater_than=1)
shape.compare_batch_dimensions(tensors=(rays, rotation_matrix,
translation_vector),
tensor_names=("points_on_rays",
"rotation_matrix",
"translation_vector"),
last_axes=-3,
broadcast_compatible=False)
rays_org = _move_in_front_of_camera(tf.zeros_like(rays), rotation_matrix,
translation_vector)
rays_dir_ = _move_in_front_of_camera(rays, rotation_matrix,
0 * translation_vector)
rays_dir = rays_dir_ / tf.norm(rays_dir_, axis=-1, keepdims=True)
return rays_org, rays_dir
def from_axis_angle_translation(
axis: type_alias.TensorLike,
angle: type_alias.TensorLike,
translation_vector: type_alias.TensorLike,
name: str = "dual_quat_from_axis_angle_trans"
) -> type_alias.TensorLike:
"""Converts an axis-angle rotation and translation to a dual 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.
translation_vector: A `[A1, ..., An, 3]`-tensor, where the last dimension
represents a translation vector.
name: A name for this op that defaults to "dual_quat_from_axis_angle_trans".
Returns:
A `[A1, ..., An, 8]`-tensor, where the last dimension represents a
normalized dual quaternion.
Raises:
ValueError: If the shape of `axis`, `angle`, or `translation_vector`
is not supported.
"""
with tf.name_scope(name):
axis = tf.convert_to_tensor(value=axis)
angle = tf.convert_to_tensor(value=angle)
translation_vector = tf.convert_to_tensor(value=translation_vector)
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.check_static(tensor=translation_vector,
tensor_name="translation_vector",
has_dim_equals=(-1, 3))
shape.compare_batch_dimensions(tensors=(axis, angle,
translation_vector),
last_axes=-2,
broadcast_compatible=True)
scalar_shape = tf.concat((tf.shape(translation_vector)[:-1], (1, )),
axis=-1)
dtype = translation_vector.dtype
quaternion_rotation = quaternion.from_axis_angle(axis, angle)
quaternion_translation = tf.concat(
(translation_vector, tf.zeros(scalar_shape, dtype)), axis=-1)
dual_quaternion_dual_part = 0.5 * quaternion.multiply(
quaternion_translation, quaternion_rotation)
return tf.concat((quaternion_rotation, dual_quaternion_dual_part),
axis=-1)
def from_axis_angle(axis, angle, name=None):
"""Convert an axis-angle representation to a rotation matrix.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
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 a normalized axis.
name: A name for this op that defaults to
"rotation_matrix_3d_from_axis_angle".
Returns:
A tensor of shape `[A1, ..., An, 3, 3]`, where the last two dimensions
represents a 3d rotation matrix.
Raises:
ValueError: If the shape of `axis` or `angle` is not supported.
"""
with tf.compat.v1.name_scope(name, "rotation_matrix_3d_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),
tensor_names=("axis", "angle"),
last_axes=-2,
broadcast_compatible=True)
axis = asserts.assert_normalized(axis)
sin_axis = tf.sin(angle) * axis
cos_angle = tf.cos(angle)
cos1_axis = (1.0 - cos_angle) * axis
_, axis_y, axis_z = tf.unstack(axis, axis=-1)
cos1_axis_x, cos1_axis_y, _ = tf.unstack(cos1_axis, axis=-1)
sin_axis_x, sin_axis_y, sin_axis_z = tf.unstack(sin_axis, axis=-1)
tmp = cos1_axis_x * axis_y
m01 = tmp - sin_axis_z
m10 = tmp + sin_axis_z
tmp = cos1_axis_x * axis_z
m02 = tmp + sin_axis_y
m20 = tmp - sin_axis_y
tmp = cos1_axis_y * axis_z
m12 = tmp - sin_axis_x
m21 = tmp + sin_axis_x
diag = cos1_axis * axis + cos_angle
diag_x, diag_y, diag_z = tf.unstack(diag, axis=-1)
matrix = tf.stack((diag_x, m01, m02,
m10, diag_y, m12,
m20, m21, diag_z),
axis=-1) # pyformat: disable
output_shape = tf.concat((tf.shape(input=axis)[:-1], (3, 3)), axis=-1)
return tf.reshape(matrix, shape=output_shape)
def _points_from_z_values(ray_org: TensorLike, ray_dir: TensorLike,
z_values: TensorLike) -> tf.Tensor:
"""Sample points on rays given the z values (distances along the rays).
Args:
ray_org: A tensor of shape `[A1, ..., An, 3]`,
where the last dimension represents the 3D position of the ray origin.
ray_dir: A tensor of shape `[A1, ..., An, 3]`,
where the last dimension represents the 3D direction of the ray.
z_values: A tensor of shape `[A1, ..., An, M]` containing the 1D position of
M points along the ray.
Returns:
A tensor of shape `[A1, ..., An, M, 3]`
"""
shape.check_static(tensor=ray_dir,
tensor_name="ray_dir",
has_dim_equals=(-1, 3))
shape.check_static(tensor=ray_org,
tensor_name="ray_org",
has_dim_equals=(-1, 3))
shape.compare_batch_dimensions(tensors=(ray_org, ray_dir, z_values),
tensor_names=("ray_org", "ray_dir",
"z_values"),
last_axes=-2,
broadcast_compatible=False)
points3d = (tf.expand_dims(ray_dir, axis=-2) *
tf.expand_dims(z_values, axis=-1))
points3d = tf.expand_dims(ray_org, -2) + points3d
return points3d
def dot(vector1, vector2, axis=-1, keepdims=True, name=None):
"""Computes the dot 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: Tensor of rank R and shape `[A1, ..., Ai, ..., An]`, where the
dimension i = axis represents a vector.
vector2: Tensor of rank R and shape `[A1, ..., Ai, ..., An]`, where the
dimension i = axis represents a vector.
axis: The dimension along which to compute the dot product.
keepdims: If True, retains reduced dimensions with length 1.
name: A name for this op which defaults to "vector_dot".
Returns:
A tensor of shape `[A1, ..., Ai = 1, ..., An]`, where the dimension i = axis
represents the result of the dot product.
"""
with tf.compat.v1.name_scope(name, "vector_dot", [vector1, vector2]):
vector1 = tf.convert_to_tensor(value=vector1)
vector2 = tf.convert_to_tensor(value=vector2)
shape.compare_batch_dimensions(tensors=(vector1, vector2),
last_axes=-1,
broadcast_compatible=True)
shape.compare_dimensions(tensors=(vector1, vector2),
axes=axis,
tensor_names=("vector1", "vector2"))
return tf.reduce_sum(input_tensor=vector1 * vector2,
axis=axis,
keepdims=keepdims)
def regular_inverse_1d(near: TensorLike,
far: TensorLike,
num_samples: int,
name="regular_inverse_1d") -> tf.Tensor:
"""Sample inverse evenly spaced numbers on an interval.
Args:
near: A tensor of shape `[A1, ... An]` containing the starting points of
the sampling interval.
far: A tensor of shape `[A1, ... An]` containing the ending points of
the sampling interval.
num_samples: The number M of points to be sampled.
name: A name for this op that defaults to "regular_inverse_1d".
Returns:
A tensor of shape `[A1, ..., An, M]` indicating the M points on the ray
"""
with tf.name_scope(name):
near = tf.convert_to_tensor(near)
far = tf.convert_to_tensor(far)
shape.compare_batch_dimensions(tensors=(tf.expand_dims(near, axis=-1),
tf.expand_dims(far, axis=-1)),
tensor_names=("near", "far"),
last_axes=-1,
broadcast_compatible=True)
return 1. / tf.linspace(1. / near, 1. / far, num_samples, 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 evaluate(point_set_a: type_alias.TensorLike,
point_set_b: type_alias.TensorLike,
name: str = "hausdorff_distance_evaluate") -> tf.Tensor:
"""Computes the Hausdorff distance from point_set_a to point_set_b.
Note:
Hausdorff distance from point_set_a to point_set_b is defined as the maximum
of all distances from a point in point_set_a to the closest point in
point_set_b. It is an asymmetric metric.
Note:
This function returns the exact Hausdorff distance and not an approximation.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
Args:
point_set_a: A tensor of shape `[A1, ..., An, N, D]`, where the last axis
represents points in a D dimensional space.
point_set_b: A tensor of shape `[A1, ..., An, M, D]`, where the last axis
represents points in a D dimensional space.
name: A name for this op. Defaults to "hausdorff_distance_evaluate".
Returns:
A tensor of shape `[A1, ..., An]` storing the hausdorff distance from
from point_set_a to point_set_b.
Raises:
ValueError: if the shape of `point_set_a`, `point_set_b` is not supported.
"""
with tf.name_scope(name):
point_set_a = tf.convert_to_tensor(value=point_set_a)
point_set_b = tf.convert_to_tensor(value=point_set_b)
shape.compare_batch_dimensions(
tensors=(point_set_a, point_set_b),
tensor_names=("point_set_a", "point_set_b"),
last_axes=-3,
broadcast_compatible=True)
# Verify that the last axis of the tensors has the same dimension.
dimension = point_set_a.shape.as_list()[-1]
shape.check_static(
tensor=point_set_b,
tensor_name="point_set_b",
has_dim_equals=(-1, dimension))
# Create N x M matrix where the entry i,j corresponds to ai - bj (vector of
# dimension D).
difference = (
tf.expand_dims(point_set_a, axis=-2) -
tf.expand_dims(point_set_b, axis=-3))
# Calculate the square distances between each two points: |ai - bj|^2.
square_distances = tf.einsum("...i,...i->...", difference, difference)
minimum_square_distance_a_to_b = tf.reduce_min(
input_tensor=square_distances, axis=-1)
return tf.sqrt(
tf.reduce_max(input_tensor=minimum_square_distance_a_to_b, axis=-1))
def model_to_eye(point_model_space,
camera_position,
look_at_point,
up_vector,
name=None):
"""Transforms points from model to eye coordinates.
Note:
In the following, A1 to An are optional batch dimensions which must be
broadcast compatible.
Args:
point_model_space: A tensor of shape `[A1, ..., An, 3]`, where the last
dimension represents the 3D points in model space.
camera_position: A tensor of shape `[A1, ..., An, 3]`, where the last
dimension represents the 3D position of the camera.
look_at_point: A tensor of shape `[A1, ..., An, 3]`, with the last dimension
storing the position where the camera is looking at.
up_vector: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
defines the up vector of the camera.
name: A name for this op. Defaults to 'model_to_eye'.
Raises:
ValueError: if the all the inputs are not of the same shape, or if any input
of of an unsupported shape.
Returns:
A tensor of shape `[A1, ..., An, 3]`, containing `point_model_space` in eye
coordinates.
"""
with tf.compat.v1.name_scope(
name, "model_to_eye",
[point_model_space, camera_position, look_at_point, up_vector]):
point_model_space = tf.convert_to_tensor(value=point_model_space)
camera_position = tf.convert_to_tensor(value=camera_position)
look_at_point = tf.convert_to_tensor(value=look_at_point)
up_vector = tf.convert_to_tensor(value=up_vector)
shape.check_static(tensor=point_model_space,
tensor_name="point_model_space",
has_dim_equals=(-1, 3))
shape.compare_batch_dimensions(tensors=(point_model_space,
camera_position),
last_axes=-2,
tensor_names=("point_model_space",
"camera_position"),
broadcast_compatible=True)
model_to_eye_matrix = look_at.right_handed(camera_position,
look_at_point, up_vector)
batch_shape = tf.shape(input=point_model_space)[:-1]
one = tf.ones(shape=tf.concat((batch_shape, (1, )), axis=-1),
dtype=point_model_space.dtype)
point_model_space = tf.concat((point_model_space, one), axis=-1)
point_model_space = tf.expand_dims(point_model_space, axis=-1)
res = tf.squeeze(tf.matmul(model_to_eye_matrix, point_model_space),
axis=-1)
return res[..., :-1]
def vector_weights(vector1: type_alias.TensorLike,
vector2: type_alias.TensorLike,
percent: Union[type_alias.Float, type_alias.TensorLike],
eps: Optional[type_alias.Float] = None,
name: str = "vector_weights") -> Tuple[tf.Tensor, tf.Tensor]:
"""Spherical linear interpolation (slerp) between two unnormalized vectors.
This function applies geometric slerp to unnormalized vectors by first
normalizing them to return the interpolation weights. It reduces to lerp when
input vectors are exactly anti-parallel.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
vector1: A tensor of shape `[A1, ... , An, M]`, which stores a normalized
vector in its last dimension.
vector2: A tensor of shape `[A1, ... , An, M]`, which stores a normalized
vector in its last dimension.
percent: A `float` or tensor with shape broadcastable to the shape of input
vectors.
eps: A small float for operation safety. If left None, its value is
automatically selected using dtype of input vectors.
name: A name for this op. Defaults to "vector_weights".
Raises:
ValueError: if the shape of `vector1`, `vector2`, or `percent` is not
supported.
Returns:
Two tensors of shape `[A1, ... , An, 1]`, representing interpolation weights
for each input vector.
"""
with tf.name_scope(name):
vector1 = tf.convert_to_tensor(value=vector1)
vector2 = tf.convert_to_tensor(value=vector2)
percent = tf.convert_to_tensor(value=percent, dtype=vector1.dtype)
if percent.shape.ndims == 0:
percent = tf.expand_dims(percent, axis=0)
shape.compare_dimensions(
tensors=(vector1, vector2),
axes=-1,
tensor_names=("vector1", "vector2"))
shape.compare_batch_dimensions(
tensors=(vector1, vector2, percent),
last_axes=(-2, -2, -1),
broadcast_compatible=True,
tensor_names=("vector1", "vector2", "percent"))
normalized1 = tf.nn.l2_normalize(vector1, axis=-1)
normalized2 = tf.nn.l2_normalize(vector2, axis=-1)
dot_product = _safe_dot(normalized1, normalized2, eps)
theta = tf.acos(dot_product)
scale1 = safe_ops.safe_sinpx_div_sinx(theta, 1.0 - percent, eps)
scale2 = safe_ops.safe_sinpx_div_sinx(theta, percent, eps)
return scale1, scale2
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 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 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 random_rays(focal: tf.Tensor,
principal_point: tf.Tensor,
height: int,
width: int,
n_rays: int,
margin: int = 0,
name: str = "random_rays") -> Tuple[tf.Tensor, tf.Tensor]:
"""Sample rays at random pixel location from a perspective camera.
Args:
focal: A tensor of shape `[A1, ..., An, 2]` where the last dimension
contains the fx and fy focal length values.
principal_point: A tensor of shape `[A1, ..., An, 2]` where the last
dimension contains the cx and cy principal point values.
height: The height of the image plane in pixels
width: The width of the image plane in pixels.
n_rays: The number M of rays to sample.
margin: The margin around the borders of the image.
name: A name for this op that defaults to "random_rays".
Returns:
A tensor of shape `[A1, ..., An, M, 3]` with the ray directions and
a tensor of shape `[A1, ..., An, M, 2]` with the pixel x, y locations.
"""
with tf.name_scope(name):
focal = tf.convert_to_tensor(value=focal)
principal_point = tf.convert_to_tensor(value=principal_point)
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.compare_batch_dimensions(tensors=(focal, principal_point),
tensor_names=("focal",
"principal_point"),
last_axes=-2,
broadcast_compatible=True)
batch_dims = tf.shape(focal)[:-1]
target_shape = tf.concat([batch_dims, [n_rays]], axis=0)
random_x = tf.random.uniform(target_shape,
minval=margin,
maxval=width - margin,
dtype=tf.int32)
random_y = tf.random.uniform(target_shape,
minval=margin,
maxval=height - margin,
dtype=tf.int32)
pixels = tf.cast(tf.stack((random_x, random_y), axis=-1), tf.float32)
rays = ray(pixels, tf.expand_dims(focal, -2),
tf.expand_dims(principal_point, -2))
return rays, tf.cast(pixels, tf.int32)
def linear_coefficients(matte: type_alias.TensorLike,
pseudo_inverse: type_alias.TensorLike,
name: str = "matting_linear_coefficients"
) -> Tuple[tf.Tensor, tf.Tensor]:
"""Computes the matting linear coefficients.
Computes the matting linear coefficients (a, b) based on the `pseudo_inverse`
generated by the `build_matrices` 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]`.
pseudo_inverse: A tensor of shape `[B, H - pad, W - pad, C + 1, size^2]`
containing the pseudo-inverse matrices computed by the `build_matrices`
function, where `pad` is equal to `size - 1` and `size` is the patch size
used to compute this tensor.
name: A name for this op. Defaults to "matting_linear_coefficients".
Returns:
A tuple contraining two Tensors for the linear coefficients (a, b) of shape
`[B, H, W, C]` and `[B, H, W, 1]`.
Raises:
ValueError: If the last dimension of `matte` is not 1. If `matte` is not
of rank 4. If `pseudo_inverse` is not of rank 5. If `B` is different
between `matte` and `pseudo_inverse`.
"""
with tf.name_scope(name):
matte = tf.convert_to_tensor(value=matte)
pseudo_inverse = tf.convert_to_tensor(value=pseudo_inverse)
pixels = tf.compat.dimension_value(pseudo_inverse.shape[-1])
shape.check_static(matte, has_rank=4, has_dim_equals=(-1, 1))
shape.check_static(pseudo_inverse, has_rank=5)
shape.compare_batch_dimensions(
tensors=(matte, pseudo_inverse),
last_axes=0,
broadcast_compatible=False)
size = np.sqrt(pixels)
# Computes the linear coefficients.
patches = tf.expand_dims(_image_patches(matte, size), axis=-1)
coeffs = tf.squeeze(tf.matmul(pseudo_inverse, patches), axis=-1)
# Averages the linear coefficients over patches.
height = tf.shape(input=coeffs)[1]
width = tf.shape(input=coeffs)[2]
ones = tf.ones(shape=_shape((1,), height, width, 1), dtype=matte.dtype)
height = tf.shape(input=matte)[1] + size - 1
width = tf.shape(input=matte)[2] + size - 1
coeffs = tf.image.resize_with_crop_or_pad(coeffs, height, width)
ones = tf.image.resize_with_crop_or_pad(ones, height, width)
coeffs = _image_average(coeffs, size) / _image_average(ones, size)
return tf.split(coeffs, (-1, 1), axis=-1)
def evaluate(ground_truth: type_alias.TensorLike,
prediction: type_alias.TensorLike,
precision_function: Callable[...,
Any] = precision_module.evaluate,
recall_function: Callable[..., Any] = recall_module.evaluate,
name: str = "fscore_evaluate") -> tf.Tensor:
"""Computes the fscore metric for the given ground truth and predicted labels.
The fscore is calculated as 2 * (precision * recall) / (precision + recall)
where the precision and recall are evaluated by the given function parameters.
The precision and recall functions default to their definition for boolean
labels (see https://en.wikipedia.org/wiki/Precision_and_recall for more
details).
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
Args:
ground_truth: A tensor of shape `[A1, ..., An, N]`, where the last axis
represents the ground truth values.
prediction: A tensor of shape `[A1, ..., An, N]`, where the last axis
represents the predicted values.
precision_function: The function to use for evaluating the precision.
Defaults to the precision evaluation for binary ground-truth and
predictions.
recall_function: The function to use for evaluating the recall. Defaults to
the recall evaluation for binary ground-truth and prediction.
name: A name for this op. Defaults to "fscore_evaluate".
Returns:
A tensor of shape `[A1, ..., An]` that stores the fscore metric for the
given ground truth labels and predictions.
Raises:
ValueError: if the shape of `ground_truth`, `prediction` is
not supported.
"""
with tf.name_scope(name):
ground_truth = tf.convert_to_tensor(value=ground_truth)
prediction = tf.convert_to_tensor(value=prediction)
shape.compare_batch_dimensions(tensors=(ground_truth, prediction),
tensor_names=("ground_truth",
"prediction"),
last_axes=-1,
broadcast_compatible=True)
recall = recall_function(ground_truth, prediction)
precision = precision_function(ground_truth, prediction)
return safe_ops.safe_signed_div(2 * precision * recall,
precision + recall)
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 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 check_valid_graph_unpooling_input(data: type_alias.TensorLike,
pool_map: tf.sparse.SparseTensor,
sizes: type_alias.TensorLike):
"""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.")
if not isinstance(pool_map, tf.sparse.SparseTensor):
raise ValueError("'pool_map' 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=data, tensor_name="data", has_rank_less_than=6)
shape.check_static(tensor=pool_map,
tensor_name="pool_map",
has_rank=data_ndims)
if not _is_dynamic_shape(tensors=(data, pool_map)):
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 evaluate(ground_truth_labels: type_alias.TensorLike,
predicted_labels: type_alias.TensorLike,
grid_size: int = 1,
name: str = "intersection_over_union_evaluate") -> tf.Tensor:
"""Computes the Intersection-Over-Union metric for the given ground truth and predicted labels.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible, and G1 to Gm are the grid dimensions.
Args:
ground_truth_labels: A tensor of shape `[A1, ..., An, G1, ..., Gm]`, where
the last m axes represent a grid of ground truth attributes. Each
attribute can either be 0 or 1.
predicted_labels: A tensor of shape `[A1, ..., An, G1, ..., Gm]`, where the
last m axes represent a grid of predicted attributes. Each attribute can
either be 0 or 1.
grid_size: The number of grid dimensions. Defaults to 1.
name: A name for this op. Defaults to "intersection_over_union_evaluate".
Returns:
A tensor of shape `[A1, ..., An]` that stores the intersection-over-union
metric of the given ground truth labels and predictions.
Raises:
ValueError: if the shape of `ground_truth_labels`, `predicted_labels` is
not supported.
"""
with tf.name_scope(name):
ground_truth_labels = tf.convert_to_tensor(value=ground_truth_labels)
predicted_labels = tf.convert_to_tensor(value=predicted_labels)
shape.compare_batch_dimensions(tensors=(ground_truth_labels,
predicted_labels),
tensor_names=("ground_truth_labels",
"predicted_labels"),
last_axes=-grid_size,
broadcast_compatible=True)
ground_truth_labels = asserts.assert_binary(ground_truth_labels)
predicted_labels = asserts.assert_binary(predicted_labels)
sum_ground_truth = tf.math.reduce_sum(input_tensor=ground_truth_labels,
axis=list(range(-grid_size, 0)))
sum_predictions = tf.math.reduce_sum(input_tensor=predicted_labels,
axis=list(range(-grid_size, 0)))
intersection = tf.math.reduce_sum(input_tensor=ground_truth_labels *
predicted_labels,
axis=list(range(-grid_size, 0)))
union = sum_ground_truth + sum_predictions - intersection
return tf.where(tf.math.equal(union, 0), tf.ones_like(union),
intersection / union)
def inverse_transform_sampling_1d(
bins: TensorLike,
pdf: TensorLike,
num_samples: int,
name="inverse_transform_sampling_1d") -> tf.Tensor:
"""Sampling 1D points from a distribution using the inverse transform.
The target distrubution is defined by its probability density function and
the spatial 1D location of its bins. The new random samples correspond to
the centers of the bins.
Args:
bins: A tensor of shape `[A1, ..., An, M]` containing 1D location of M bins.
For example, a tensor [a, b, c, d] corresponds to
the bin structure |--a--|-b-|--c--|d|.
pdf: A tensor of shape `[A1, ..., An, M]` containing the probability
distribution in M bins.
num_samples: The number N of new samples.
name: A name for this op that defaults to "inverse_transform_sampling_1d".
Returns:
A tensor of shape `[A1, ..., An, N]` indicating the new N random points.
"""
with tf.name_scope(name):
bins = tf.convert_to_tensor(value=bins)
pdf = tf.convert_to_tensor(value=pdf)
shape.check_static(tensor=bins,
tensor_name="bins",
has_rank_greater_than=0)
shape.check_static(tensor=pdf,
tensor_name="pdf",
has_rank_greater_than=0)
shape.compare_batch_dimensions(tensors=(bins, pdf),
tensor_names=("bins", "pdf"),
last_axes=-2,
broadcast_compatible=True)
shape.compare_dimensions(tensors=(bins, pdf),
tensor_names=("bins", "pdf"),
axes=-1)
# Do not consider the last bin, as the cdf contains has +1 dimension.
pdf = _normalize_pdf(pdf[..., :-1])
cdf = _get_cdf(pdf)
batch_shape = tf.shape(pdf)[:-1]
# TODO(krematas): Use dynamic values
batch_dims = tf.get_static_value(tf.rank(pdf) - 1)
target_shape = tf.concat([batch_shape, [num_samples]], axis=-1)
uniform_samples = tf.random.uniform(target_shape)
bin_indices = tf.searchsorted(cdf, uniform_samples, side="right")
bin_indices = tf.maximum(0, bin_indices - 1)
z_values = tf.gather(bins, bin_indices, axis=-1, batch_dims=batch_dims)
return z_values
def gather_faces(vertices: type_alias.TensorLike,
indices: type_alias.TensorLike,
name: str = "normals_gather_faces") -> type_alias.TensorLike:
"""Gather corresponding vertices for each face.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
vertices: A tensor of shape `[A1, ..., An, V, D]`, where `V` is the number
of vertices and `D` the dimensionality of each vertex. The rank of this
tensor should be at least 2.
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. The rank of this tensor
should be at least 2.
name: A name for this op. Defaults to "normals_gather_faces".
Returns:
A tensor of shape `[A1, ..., An, F, M, D]` containing the vertices of each
face.
Raises:
ValueError: If the shape of `vertices` or `indices` is not supported.
"""
with tf.name_scope(name):
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)
shape.check_static(tensor=indices,
tensor_name="indices",
has_rank_greater_than=1)
shape.compare_batch_dimensions(tensors=(vertices, indices),
last_axes=(-3, -3),
broadcast_compatible=False)
if hasattr(tf, "batch_gather"):
expanded_vertices = tf.expand_dims(vertices, axis=-3)
broadcasted_shape = tf.concat(
[tf.shape(input=indices)[:-1],
tf.shape(input=vertices)[-2:]],
axis=-1)
broadcasted_vertices = tf.broadcast_to(expanded_vertices,
broadcasted_shape)
return tf.gather(broadcasted_vertices, indices, batch_dims=-1)
else:
return tf.gather(vertices,
indices,
axis=-2,
batch_dims=indices.shape.ndims - 2)
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 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
