def test_assert_all_above_exception_raised(self, dtype):
"""Checks that assert_all_above raises exceptions for invalid input."""
vector = _pick_random_vector()
vector = tf.convert_to_tensor(value=vector, dtype=dtype)
vector = vector * vector
vector /= -tf.reduce_max(input_tensor=vector, axis=-1, keepdims=True)
eps = asserts.select_eps_for_addition(dtype)
outside_vector = vector - eps
ones_vector = -tf.ones_like(vector)
with self.subTest(name="outside_and_open_bounds"):
with self.assertRaises(tf.errors.InvalidArgumentError):
self.evaluate(
asserts.assert_all_above(outside_vector,
-1.0,
open_bound=True))
with self.subTest(name="outside_and_close_bounds"):
with self.assertRaises(tf.errors.InvalidArgumentError):
self.evaluate(
asserts.assert_all_above(outside_vector,
-1.0,
open_bound=False))
with self.subTest(name="exact_and_open_bounds"):
with self.assertRaises(tf.errors.InvalidArgumentError):
self.evaluate(
asserts.assert_all_above(ones_vector,
-1.0,
open_bound=True))
def _random_categorical_sample(
num_samples: int,
weights: type_alias.TensorLike,
seed: Optional[type_alias.TensorLike] = None,
stateless: bool = False,
name: str = "random_categorical_sample",
sample_dtype: tf.DType = tf.int32) -> type_alias.TensorLike:
"""Samples from a categorical distribution with arbitrary batch dimensions.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
num_samples: An `int32` scalar denoting the number of samples to generate
per mesh.
weights: A `float` tensor of shape `[A1, ..., An, F]` where F is number of
faces.
All weights must be > 0.
seed: Optional random seed, value depends on `stateless`.
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 the same reproducible sequence
across calls. If stateless=False, then a stateful random number generator
is used (default behavior).
name: Name for op. Defaults to "random_categorical_sample".
sample_dtype: Type of output samples.
Returns:
A `sample_dtype` tensor of shape `[A1, ..., An, num_samples]`.
"""
with tf.name_scope(name):
asserts.assert_all_above(weights, 0)
logits = tf.math.log(weights)
num_faces = tf.shape(input=logits)[-1]
batch_shape = tf.shape(input=logits)[:-1]
logits_2d = tf.reshape(logits, [-1, num_faces])
if stateless:
seed = tf.convert_to_tensor(value=seed)
shape.check_static(tensor=seed,
tensor_name="seed",
has_dim_equals=(-1, 2))
sample_fn = tf.random.stateless_categorical
else:
sample_fn = tf.random.categorical
draws = sample_fn(logits=logits_2d,
num_samples=num_samples,
dtype=sample_dtype,
seed=seed)
samples = tf.reshape(draws,
shape=tf.concat((batch_shape, (num_samples, )),
axis=0))
return samples
def test_assert_all_above_passthrough(self):
"""Checks that the assert is a passthrough when the flag is False."""
vector_input = _pick_random_vector()
vector_output = asserts.assert_all_above(vector_input, 1.0)
self.assertIs(vector_input, vector_output)
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 generate_random_face_indices(
num_samples: int,
face_weights: type_alias.TensorLike,
seed: Optional[type_alias.TensorLike] = None,
stateless: bool = False,
name: str = "generate_random_face_indices") -> type_alias.TensorLike:
"""Generate a sample of face ids given per face probability.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
num_samples: An `int32` scalar denoting the number of samples to generate
per mesh.
face_weights: A `float` tensor of shape `[A1, ..., An, F]` where F is
number of faces. All weights must be > 0.
seed: Optional seed for the random number generator.
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 the same reproducible sequence
across calls. If stateless=False, then a stateful random number generator
is used (default behavior).
name: Name for op. Defaults to "generate_random_face_indices".
Returns:
An `int32` tensor of shape `[A1, ..., An, num_samples]` denoting sampled
face indices.
"""
with tf.name_scope(name):
num_samples = tf.convert_to_tensor(value=num_samples)
face_weights = tf.convert_to_tensor(value=face_weights)
shape.check_static(tensor=face_weights,
tensor_name="face_weights",
has_rank_greater_than=0)
shape.check_static(tensor=num_samples,
tensor_name="num_samples",
has_rank=0)
face_weights = asserts.assert_all_above(face_weights, minval=0.0)
eps = asserts.select_eps_for_division(face_weights.dtype)
face_weights = face_weights + eps
sampled_face_indices = _random_categorical_sample(
num_samples=num_samples,
weights=face_weights,
seed=seed,
stateless=stateless)
return sampled_face_indices
def evaluate_legendre_polynomial(degree_l, order_m, x):
"""Evaluates the Legendre polynomial of degree l and order m at x.
Note:
This function is implementing the algorithm described in p. 10 of `Spherical
Harmonic Lighting: The Gritty Details`.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
degree_l: An integer tensor of shape `[A1, ..., An]` corresponding to the
degree of the associated Legendre polynomial. Note that `degree_l` must be
non-negative.
order_m: An integer tensor of shape `[A1, ..., An]` corresponding to the
order of the associated Legendre polynomial. Note that `order_m` must
satisfy `0 <= order_m <= l`.
x: A tensor of shape `[A1, ..., An]` with values in [-1,1].
Returns:
A tensor of shape `[A1, ..., An]` containing the evaluation of the legendre
polynomial.
"""
degree_l = tf.convert_to_tensor(value=degree_l)
order_m = tf.convert_to_tensor(value=order_m)
x = tf.convert_to_tensor(value=x)
if not degree_l.dtype.is_integer:
raise ValueError("`degree_l` must be of an integer type.")
if not order_m.dtype.is_integer:
raise ValueError("`order_m` must be of an integer type.")
shape.compare_batch_dimensions(
tensors=(degree_l, order_m, x),
last_axes=-1,
tensor_names=("degree_l", "order_m", "x"),
broadcast_compatible=True)
degree_l = asserts.assert_all_above(degree_l, 0)
order_m = asserts.assert_all_in_range(order_m, 0, degree_l)
x = asserts.assert_all_in_range(x, -1.0, 1.0)
pmm = _evaluate_legendre_polynomial_pmm_eval(order_m, x)
return tf.compat.v1.where(
tf.equal(degree_l, order_m), pmm,
_evaluate_legendre_polynomial_branch(degree_l, order_m, x, pmm))
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 perspective_right_handed(vertical_field_of_view,
aspect_ratio,
near,
far,
name=None):
"""Generates the matrix for a right handed perspective projection.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
vertical_field_of_view: A tensor of shape `[A1, ..., An]`, where the last
dimension represents the vertical field of view of the frustum expressed
in radians. Note that values for `vertical_field_of_view` must be in the
range (0,pi).
aspect_ratio: A tensor of shape `[A1, ..., An, C]`, where the last dimension
stores the width over height ratio of the frustum. Note that values for
`aspect_ratio` must be non-negative.
near: A tensor of shape `[A1, ..., An, C]`, where the last dimension
captures the distance between the viewer and the near clipping plane. Note
that values for `near` must be non-negative.
far: A tensor of shape `[A1, ..., An, C]`, where the last dimension
captures the distance between the viewer and the far clipping plane. Note
that values for `far` must be greater than those of `near`.
name: A name for this op. Defaults to 'perspective_rh'.
Raises:
InvalidArgumentError: if any input contains data not in the specified range
of valid values.
ValueError: if the all the inputs are not of the same shape.
Returns:
A tensor of shape `[A1, ..., An, 4, 4]`, containing matrices of right
handed perspective-view frustum.
"""
with tf.compat.v1.name_scope(
name, "perspective_rh",
[vertical_field_of_view, aspect_ratio, near, far]):
vertical_field_of_view = tf.convert_to_tensor(value=vertical_field_of_view)
aspect_ratio = tf.convert_to_tensor(value=aspect_ratio)
near = tf.convert_to_tensor(value=near)
far = tf.convert_to_tensor(value=far)
shape.compare_batch_dimensions(
tensors=(vertical_field_of_view, aspect_ratio, near, far),
last_axes=-1,
tensor_names=("vertical_field_of_view", "aspect_ratio", "near", "far"),
broadcast_compatible=False)
vertical_field_of_view = asserts.assert_all_in_range(
vertical_field_of_view, 0.0, math.pi, open_bounds=True)
aspect_ratio = asserts.assert_all_above(aspect_ratio, 0.0, open_bound=True)
near = asserts.assert_all_above(near, 0.0, open_bound=True)
far = asserts.assert_all_above(far, near, open_bound=True)
inverse_tan_half_vertical_field_of_view = 1.0 / tf.tan(
vertical_field_of_view * 0.5)
zero = tf.zeros_like(inverse_tan_half_vertical_field_of_view)
one = tf.ones_like(inverse_tan_half_vertical_field_of_view)
x = tf.stack((inverse_tan_half_vertical_field_of_view / aspect_ratio, zero,
zero, zero),
axis=-1)
y = tf.stack((zero, inverse_tan_half_vertical_field_of_view, zero, zero),
axis=-1)
near_minus_far = near - far
z = tf.stack(
(zero, zero,
(far + near) / near_minus_far, 2.0 * far * near / near_minus_far),
axis=-1)
w = tf.stack((zero, zero, -one, zero), axis=-1)
return tf.stack((x, y, z, w), axis=-2)
def brdf(direction_incoming_light: type_alias.TensorLike,
direction_outgoing_light: type_alias.TensorLike,
surface_normal: type_alias.TensorLike,
shininess: type_alias.TensorLike,
albedo: type_alias.TensorLike,
brdf_normalization: bool = True,
name: str = "phong_brdf") -> tf.Tensor:
"""Evaluates the specular brdf of the Phong model.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
Note:
The gradient of this function is not smooth when the dot product of the
normal with any light is 0.0.
Args:
direction_incoming_light: A tensor of shape `[A1, ..., An, 3]`, where the
last dimension represents a normalized incoming light vector.
direction_outgoing_light: A tensor of shape `[A1, ..., An, 3]`, where the
last dimension represents a normalized outgoing light vector.
surface_normal: A tensor of shape `[A1, ..., An, 3]`, where the last
dimension represents a normalized surface normal.
shininess: A tensor of shape `[A1, ..., An, 1]`, where the last dimension
represents a non-negative shininess coefficient.
albedo: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents albedo with values in [0,1].
brdf_normalization: A `bool` indicating whether normalization should be
applied to enforce the energy conservation property of BRDFs. Note that
`brdf_normalization` must be set to False in order to use the original
Blinn specular model.
name: A name for this op. Defaults to "phong_brdf".
Returns:
A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents
the amount of light reflected in the outgoing light direction.
Raises:
ValueError: if the shape of `direction_incoming_light`,
`direction_outgoing_light`, `surface_normal`, `shininess` or `albedo` is not
supported.
InvalidArgumentError: if not all of shininess values are non-negative, or if
at least one element of `albedo` is outside of [0,1].
"""
with tf.name_scope(name):
direction_incoming_light = tf.convert_to_tensor(
value=direction_incoming_light)
direction_outgoing_light = tf.convert_to_tensor(
value=direction_outgoing_light)
surface_normal = tf.convert_to_tensor(value=surface_normal)
shininess = tf.convert_to_tensor(value=shininess)
albedo = tf.convert_to_tensor(value=albedo)
shape.check_static(tensor=direction_incoming_light,
tensor_name="direction_incoming_light",
has_dim_equals=(-1, 3))
shape.check_static(tensor=direction_outgoing_light,
tensor_name="direction_outgoing_light",
has_dim_equals=(-1, 3))
shape.check_static(tensor=surface_normal,
tensor_name="surface_normal",
has_dim_equals=(-1, 3))
shape.check_static(tensor=shininess,
tensor_name="shininess",
has_dim_equals=(-1, 1))
shape.check_static(tensor=albedo,
tensor_name="albedo",
has_dim_equals=(-1, 3))
shape.compare_batch_dimensions(
tensors=(direction_incoming_light, direction_outgoing_light,
surface_normal, shininess, albedo),
tensor_names=("direction_incoming_light",
"direction_outgoing_light", "surface_normal",
"shininess", "albedo"),
last_axes=-2,
broadcast_compatible=True)
direction_incoming_light = asserts.assert_normalized(
direction_incoming_light)
direction_outgoing_light = asserts.assert_normalized(
direction_outgoing_light)
surface_normal = asserts.assert_normalized(surface_normal)
albedo = asserts.assert_all_in_range(albedo,
0.0,
1.0,
open_bounds=False)
shininess = asserts.assert_all_above(shininess, 0.0, open_bound=False)
# Checks whether the incoming or outgoing light point behind the surface.
dot_incoming_light_surface_normal = vector.dot(
-direction_incoming_light, surface_normal)
dot_outgoing_light_surface_normal = vector.dot(
direction_outgoing_light, surface_normal)
min_dot = tf.minimum(dot_incoming_light_surface_normal,
dot_outgoing_light_surface_normal)
perfect_reflection_direction = vector.reflect(direction_incoming_light,
surface_normal)
perfect_reflection_direction = tf.math.l2_normalize(
perfect_reflection_direction, axis=-1)
cos_alpha = vector.dot(perfect_reflection_direction,
direction_outgoing_light,
axis=-1)
cos_alpha = tf.maximum(cos_alpha, tf.zeros_like(cos_alpha))
phong_model = albedo * tf.pow(cos_alpha, shininess)
if brdf_normalization:
phong_model *= _brdf_normalization_factor(shininess)
common_shape = shape.get_broadcasted_shape(min_dot.shape,
phong_model.shape)
d_val = lambda dim: 1 if dim is None else tf.compat.dimension_value(dim
)
common_shape = [d_val(dim) for dim in common_shape]
condition = tf.broadcast_to(tf.greater_equal(min_dot, 0.0),
common_shape)
phong_model = tf.broadcast_to(phong_model, common_shape)
return tf.where(condition, phong_model, tf.zeros_like(phong_model))
def evaluate_spherical_harmonics(degree_l, order_m, theta, phi, name=None):
"""Evaluates a point sample of a Spherical Harmonic basis function.
Note:
This function is implementating the algorithm and variable names described
p. 12 of 'Spherical Harmonic Lighting: The Gritty Details.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
degree_l: An integer tensor of shape `[A1, ..., An, C]`, where the last
dimension represents the band of the spherical harmonics. Note that
`degree_l` must be non-negative.
order_m: An integer tensor of shape `[A1, ..., An, C]`, where the last
dimension represents the index of the spherical harmonics in the band
`degree_l`. Note that `order_m` must satisfy `0 <= order_m <= l`.
theta: A tensor of shape `[A1, ..., An, 1]`. This variable stores the polar
angle of the sameple. Values of theta must be in [0, pi].
phi: A tensor of shape `[A1, ..., An, 1]`. This variable stores the
azimuthal angle of the sameple. Values of phi must be in [0, 2pi].
name: A name for this op. Defaults to
'spherical_harmonics_evaluate_spherical_harmonics'.
Returns:
A tensor of shape `[A1, ..., An, C]` containing the evaluation of each basis
of the spherical harmonics.
Raises:
ValueError: if the shape of `theta` or `phi` is not supported.
InvalidArgumentError: if at least an element of `l`, `m`, `theta` or `phi`
is outside the expected range.
"""
with tf.compat.v1.name_scope(
name, "spherical_harmonics_evaluate_spherical_harmonics",
[degree_l, order_m, theta, phi]):
degree_l = tf.convert_to_tensor(value=degree_l)
order_m = tf.convert_to_tensor(value=order_m)
theta = tf.convert_to_tensor(value=theta)
phi = tf.convert_to_tensor(value=phi)
if not degree_l.dtype.is_integer:
raise ValueError("`degree_l` must be of an integer type.")
if not order_m.dtype.is_integer:
raise ValueError("`order_m` must be of an integer type.")
shape.compare_dimensions(tensors=(degree_l, order_m),
axes=-1,
tensor_names=("degree_l", "order_m"))
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=(degree_l, order_m, theta, phi),
last_axes=-2,
tensor_names=("degree_l", "order_m",
"theta", "phi"),
broadcast_compatible=False)
# Checks that tensors contain appropriate data.
degree_l = asserts.assert_all_above(degree_l, 0)
order_m = asserts.assert_all_in_range(order_m, -degree_l, degree_l)
theta = asserts.assert_all_in_range(theta, 0.0, np.pi)
phi = asserts.assert_all_in_range(phi, 0.0, 2.0 * np.pi)
var_type = theta.dtype
sign_m = tf.math.sign(order_m)
order_m = tf.abs(order_m)
zeros = tf.zeros_like(order_m)
result_m_zero = _spherical_harmonics_normalization(
degree_l, zeros, var_type) * evaluate_legendre_polynomial(
degree_l, zeros, tf.cos(theta))
result_branch = _evaluate_spherical_harmonics_branch(
degree_l, order_m, theta, phi, sign_m, var_type)
return tf.where(tf.equal(order_m, zeros), result_m_zero, result_branch)
def ndc_to_screen(point_ndc_space,
lower_left_corner,
screen_dimensions,
near,
far,
name=None):
"""Transforms points from normalized device coordinates to screen coordinates.
Note:
In the following, A1 to An are optional batch dimensions which must be
broadcast compatible between `point_ndc_space` and the other variables.
Args:
point_ndc_space: A tensor of shape `[A1, ..., An, 3]`, where the last
dimension represents points in normalized device coordinates.
lower_left_corner: A tensor of shape `[A1, ..., An, 2]`, where the last
dimension captures the position (in pixels) of the lower left corner of
the screen.
screen_dimensions: A tensor of shape `[A1, ..., An, 2]`, where the last
dimension is expressed in pixels and captures the width and the height (in
pixels) of the screen.
near: A tensor of shape `[A1, ..., An, 1]`, where the last dimension
captures the distance between the viewer and the near clipping plane. Note
that values for `near` must be non-negative.
far: A tensor of shape `[A1, ..., An, 1]`, where the last dimension
captures the distance between the viewer and the far clipping plane. Note
that values for `far` must be greater than those of `near`.
name: A name for this op. Defaults to 'ndc_to_screen'.
Raises:
InvalidArgumentError: if any input contains data not in the specified range
of valid values.
ValueError: If any input is of an unsupported shape.
Returns:
A tensor of shape `[A1, ..., An, 3]`, containing `point_ndc_space` in
screen coordinates.
"""
with tf.compat.v1.name_scope(
name, "ndc_to_screen",
[point_ndc_space, lower_left_corner, screen_dimensions, near, far]):
point_ndc_space = tf.convert_to_tensor(value=point_ndc_space)
lower_left_corner = tf.convert_to_tensor(value=lower_left_corner)
screen_dimensions = tf.convert_to_tensor(value=screen_dimensions)
near = tf.convert_to_tensor(value=near)
far = tf.convert_to_tensor(value=far)
shape.check_static(tensor=point_ndc_space,
tensor_name="point_ndc_space",
has_dim_equals=(-1, 3))
shape.check_static(tensor=lower_left_corner,
tensor_name="lower_left_corner",
has_dim_equals=(-1, 2))
shape.check_static(tensor=screen_dimensions,
tensor_name="screen_dimensions",
has_dim_equals=(-1, 2))
shape.check_static(tensor=near,
tensor_name="near",
has_dim_equals=(-1, 1))
shape.check_static(tensor=far,
tensor_name="far",
has_dim_equals=(-1, 1))
shape.compare_batch_dimensions(tensors=(lower_left_corner,
screen_dimensions, near, far),
last_axes=-2,
tensor_names=("lower_left_corner",
"screen_dimensions",
"near", "far"),
broadcast_compatible=False)
shape.compare_batch_dimensions(tensors=(point_ndc_space, near),
last_axes=-2,
tensor_names=("point_ndc_space",
"near"),
broadcast_compatible=True)
screen_dimensions = asserts.assert_all_above(screen_dimensions,
0.0,
open_bound=True)
near = asserts.assert_all_above(near, 0.0, open_bound=True)
far = asserts.assert_all_above(far, near, open_bound=True)
ndc_to_screen_factor = tf.concat(
(screen_dimensions, far - near), axis=-1) / 2.0
screen_center = tf.concat(
(lower_left_corner + screen_dimensions / 2.0, (near + far) / 2.0),
axis=-1)
return ndc_to_screen_factor * point_ndc_space + screen_center
def triangulate(startpoints, endpoints, weights, name=None):
"""Triangulates 3d points by miminizing the sum of squared distances to rays.
The rays are defined by their start points and endpoints. At least two rays
are required to triangulate any given point. Contrary to the standard
reprojection-error metric, the sum of squared distances to rays can be
minimized in a closed form.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
startpoints: A tensor of ray start points with shape `[A1, ..., An, V, 3]`,
the number of rays V around which the solution points live should be
greater or equal to 2, otherwise triangulation is impossible.
endpoints: A tensor of ray endpoints with shape `[A1, ..., An, V, 3]`, the
number of rays V around which the solution points live should be greater
or equal to 2, otherwise triangulation is impossible. The `endpoints`
tensor should have the same shape as the `startpoints` tensor.
weights: A tensor of ray weights (certainties) with shape `[A1, ..., An,
V]`. Weights should have all positive entries. Weight should have at least
two non-zero entries for each point (at least two rays should have
certainties > 0).
name: A name for this op. The default value of None means "ray_triangulate".
Returns:
A tensor of triangulated points with shape `[A1, ..., An, 3]`.
Raises:
ValueError: If the shape of the arguments is not supported.
"""
with tf.compat.v1.name_scope(name, "ray_triangulate",
[startpoints, endpoints, weights]):
startpoints = tf.convert_to_tensor(value=startpoints)
endpoints = tf.convert_to_tensor(value=endpoints)
weights = tf.convert_to_tensor(value=weights)
shape.check_static(
tensor=startpoints,
tensor_name="startpoints",
has_rank_greater_than=1,
has_dim_equals=(-1, 3),
has_dim_greater_than=(-2, 1))
shape.check_static(
tensor=endpoints,
tensor_name="endpoints",
has_rank_greater_than=1,
has_dim_equals=(-1, 3),
has_dim_greater_than=(-2, 1))
shape.compare_batch_dimensions(
tensors=(startpoints, endpoints, weights),
last_axes=(-2, -2, -1),
broadcast_compatible=False)
weights = asserts.assert_all_above(weights, 0.0, open_bound=False)
weights = asserts.assert_at_least_k_non_zero_entries(weights, k=2)
left_hand_side_list = []
right_hand_side_list = []
# TODO: Replace the inefficient for loop and add comments here.
for ray_id in range(weights.shape[-1]):
weights_single_ray = weights[..., ray_id]
startpoints_single_ray = startpoints[..., ray_id, :]
endpoints_singleview = endpoints[..., ray_id, :]
ray = endpoints_singleview - startpoints_single_ray
ray = tf.nn.l2_normalize(ray, axis=-1)
ray_x, ray_y, ray_z = tf.unstack(ray, axis=-1)
zeros = tf.zeros_like(ray_x)
cross_product_matrix = tf.stack(
(zeros, -ray_z, ray_y, ray_z, zeros, -ray_x, -ray_y, ray_x, zeros),
axis=-1)
cross_product_matrix_shape = tf.concat(
(tf.shape(input=cross_product_matrix)[:-1], (3, 3)), axis=-1)
cross_product_matrix = tf.reshape(
cross_product_matrix, shape=cross_product_matrix_shape)
weights_single_ray = tf.expand_dims(weights_single_ray, axis=-1)
weights_single_ray = tf.expand_dims(weights_single_ray, axis=-1)
left_hand_side = weights_single_ray * cross_product_matrix
left_hand_side_list.append(left_hand_side)
dot_product = tf.matmul(cross_product_matrix,
tf.expand_dims(startpoints_single_ray, axis=-1))
right_hand_side = weights_single_ray * dot_product
right_hand_side_list.append(right_hand_side)
left_hand_side_multi_rays = tf.concat(left_hand_side_list, axis=-2)
right_hand_side_multi_rays = tf.concat(right_hand_side_list, axis=-2)
points = tf.linalg.lstsq(left_hand_side_multi_rays,
right_hand_side_multi_rays)
points = tf.squeeze(points, axis=-1)
return points
def intersection_ray_sphere(sphere_center,
sphere_radius,
ray,
point_on_ray,
name=None):
"""Finds positions and surface normals where the sphere and the ray intersect.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
sphere_center: A tensor of shape `[3]` representing the 3d sphere center.
sphere_radius: A tensor of shape `[1]` containing a strictly positive value
defining the radius of the sphere.
ray: A tensor of shape `[A1, ..., An, 3]` containing normalized 3D vectors.
point_on_ray: A tensor of shape `[A1, ..., An, 3]`.
name: A name for this op. The default value of None means
"ray_intersection_ray_sphere".
Returns:
A tensor of shape `[2, A1, ..., An, 3]` containing the position of the
intersections, and a tensor of shape `[2, A1, ..., An, 3]` the associated
surface normals at that point. Both tensors contain NaNs when there is no
intersections. The first dimension of the returned tensor provides access to
the first and second intersections of the ray with the sphere.
Raises:
ValueError: if the shape of `sphere_center`, `sphere_radius`, `ray` or
`point_on_ray` is not supported.
tf.errors.InvalidArgumentError: If `ray` is not normalized.
"""
with tf.compat.v1.name_scope(
name, "ray_intersection_ray_sphere",
[sphere_center, sphere_radius, ray, point_on_ray]):
sphere_center = tf.convert_to_tensor(value=sphere_center)
sphere_radius = tf.convert_to_tensor(value=sphere_radius)
ray = tf.convert_to_tensor(value=ray)
point_on_ray = tf.convert_to_tensor(value=point_on_ray)
shape.check_static(
tensor=sphere_center,
tensor_name="sphere_center",
has_rank=1,
has_dim_equals=(0, 3))
shape.check_static(
tensor=sphere_radius,
tensor_name="sphere_radius",
has_rank=1,
has_dim_equals=(0, 1))
shape.check_static(tensor=ray, tensor_name="ray", has_dim_equals=(-1, 3))
shape.check_static(
tensor=point_on_ray, tensor_name="point_on_ray", has_dim_equals=(-1, 3))
shape.compare_batch_dimensions(
tensors=(ray, point_on_ray),
last_axes=(-2, -2),
broadcast_compatible=False)
sphere_radius = asserts.assert_all_above(
sphere_radius, 0.0, open_bound=True)
ray = asserts.assert_normalized(ray)
vector_sphere_center_to_point_on_ray = sphere_center - point_on_ray
distance_sphere_center_to_point_on_ray = tf.norm(
tensor=vector_sphere_center_to_point_on_ray, axis=-1, keepdims=True)
distance_projection_sphere_center_on_ray = vector.dot(
vector_sphere_center_to_point_on_ray, ray)
closest_distance_sphere_center_to_ray = tf.sqrt(
tf.square(distance_sphere_center_to_point_on_ray) -
tf.pow(distance_projection_sphere_center_on_ray, 2))
half_secant_length = tf.sqrt(
tf.square(sphere_radius) -
tf.square(closest_distance_sphere_center_to_ray))
distances = tf.stack(
(distance_projection_sphere_center_on_ray - half_secant_length,
distance_projection_sphere_center_on_ray + half_secant_length),
axis=0)
intersections_points = distances * ray + point_on_ray
normals = tf.math.l2_normalize(
intersections_points - sphere_center, axis=-1)
return intersections_points, normals
def distortion_factor(
squared_radius: type_alias.TensorLike,
distortion_coefficient: type_alias.TensorLike,
name: str = "quadratic_radial_distortion_distortion_factor"
) -> Tuple[tf.Tensor, tf.Tensor]:
"""Calculates a quadratic distortion factor given squared radii.
Given a vector describing a location in camera space in homogeneous
coordinates, `(x/z, y/z, 1)`, squared_radius is `r^2 = (x/z)^2 + (y/z)^2`.
distortion_factor multiplies `x/z` and `y/z` to obtain the distorted
coordinates. In this function, `distortion_factor` is given by
`1.0 + distortion_coefficient * squared_radius`.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
Args:
squared_radius: A tensor of shape `[A1, ..., An, H, W]`, containing the
radii of the image pixels computed as `(x/z)^2 + (y/z)^2`. We use squared
radius rather than the radius itself to avoid an unnecessary `sqrt`, which
may introduce gradient singularities. The non-negativity of squared radius
is only enforced in debug mode.
distortion_coefficient: A `scalar` or a tensor of shape `[A1, ..., An]`,
which contains the distortion coefficients of each image.
name: A name for this op. Defaults to
"quadratic_radial_distortion_distortion_factor".
Returns:
distortion_factor: A tensor of shape `[A1, ..., An, H, W]`, the correction
factor that should multiply the projective coordinates `(x/z)` and `(y/z)`
to apply the distortion.
overflow_mask: A boolean tensor of shape `[A1, ..., An, H, W]`, `True` where
`squared_radius` is beyond the range where the distortion function is
monotonically increasing. Wherever `overflow_mask` is True,
`distortion_factor`'s value is meaningless.
"""
with tf.name_scope(name, ):
squared_radius = tf.convert_to_tensor(value=squared_radius)
distortion_coefficient = tf.convert_to_tensor(
value=distortion_coefficient)
if distortion_coefficient.shape.ndims == 0:
distortion_coefficient = tf.expand_dims(distortion_coefficient,
axis=0)
shape.check_static(tensor=squared_radius,
tensor_name="squared_radius",
has_rank_greater_than=1)
shape.compare_batch_dimensions(tensors=(squared_radius,
distortion_coefficient),
tensor_names=("squared_radius",
"distortion_coefficient"),
last_axes=(-3, -1),
broadcast_compatible=True)
squared_radius = asserts.assert_all_above(squared_radius,
0.0,
open_bound=False)
distortion_coefficient = tf.expand_dims(distortion_coefficient,
axis=-1)
distortion_coefficient = tf.expand_dims(distortion_coefficient,
axis=-1)
distortion_coefficient_times_squared_radius = (distortion_coefficient *
squared_radius)
distortion_factor_ = 1.0 + distortion_coefficient_times_squared_radius
# This condition needs to hold for the distortion to be monotomnically
# increasing, as can be derived by differentiating it.
overflow_mask = tf.less(
1.0 + 3.0 * distortion_coefficient_times_squared_radius, 0.0)
return distortion_factor_, overflow_mask
def undistortion_factor(
distorted_squared_radius: type_alias.TensorLike,
distortion_coefficient: type_alias.TensorLike,
num_iterations: int = 5,
name: str = "quadratic_radial_distortion_undistortion_factor"
) -> Tuple[tf.Tensor, tf.Tensor]:
"""Calculates the inverse quadratic distortion function given squared radii.
Given a vector describing a location in camera space in homogeneous
coordinates `(x/z, y/z, 1)`, after distortion has been applied, these become
`(x'/z, y'/z, 1)`. `distorted_squared_radius` is `(x'/z)^2 + (y'/z)^2`.
`undistortion_factor` multiplies `x'/z` and `y'/z` to obtain the undistorted
projective coordinates `x/z` and `y/z`.
The undustortion factor in this function is derived from a quadratic.
distortion function, where the distortion factor equals
`1.0 + distortion_coefficient * squared_radius`.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
Args:
distorted_squared_radius: A tensor of shape `[A1, ..., An, H, W]` containing
the value of projective coordinates `(x/z)^2 + (y/z)^2`. For each pixel
it contains the squared distance of that pixel to the center of the image
plane. We use `distorted_squared_radius` rather than the distorted radius
itself to avoid an unnecessary `sqrt`, which may introduce gradient
singularities. The non-negativity of `distorted_squared_radius` is only
enforced in debug mode.
distortion_coefficient: A `scalar` or a tensor of shape `[A1, ..., An]`,
which contains the distortion coefficients of each image.
num_iterations: Number of Newton-Raphson iterations to calculate the inverse
distortion function. Defaults to 5, which is on the high-accuracy side.
name: A name for this op. Defaults to
"quadratic_radial_distortion_undistortion_factor".
Returns:
undistortion: A tensor of shape `[A1, ..., An, H, W]` containing the
correction factor that should multiply the distorted projective
coordinates `(x'/z)` and `(y'/z)` to obtain the undistorted ones.
overflow_mask: A `bool` tensor of shape `[A1, ..., An, H, W]`, `True` where
`distorted_squared_radius` is beyond the range where the distortion
function is monotonically increasing. Wherever `overflow_mask` is `True`,
`undistortion_factor`'s value is meaningless.
"""
with tf.name_scope(name):
distorted_squared_radius = tf.convert_to_tensor(
value=distorted_squared_radius)
distortion_coefficient = tf.convert_to_tensor(
value=distortion_coefficient)
if distortion_coefficient.shape.ndims == 0:
distortion_coefficient = tf.expand_dims(distortion_coefficient,
axis=0)
shape.check_static(tensor=distorted_squared_radius,
tensor_name="distorted_squared_radius",
has_rank_greater_than=1)
shape.compare_batch_dimensions(
tensors=(distorted_squared_radius, distortion_coefficient),
tensor_names=("distorted_squared_radius",
"distortion_coefficient"),
last_axes=(-3, -1),
broadcast_compatible=True)
distorted_squared_radius = asserts.assert_all_above(
distorted_squared_radius, 0.0, open_bound=False)
distortion_coefficient = tf.expand_dims(distortion_coefficient,
axis=-1)
distortion_coefficient = tf.expand_dims(distortion_coefficient,
axis=-1)
# For a distortion function of r' = (1 + ar^2)r, with a negative a, the
# maximum r until which r'(r) is monotonically increasing is r^2 = -1/(3a).
# At that value, r'^2 = -4 / (27a). Therefore the overflow condition for r'
# is ar'^2 +(4/27.0) < 0. For a positive a it never holds, as it should,
# because then r' is monotonic in r everywhere and thus never overflows.
distortion_coefficient_times_distorted_squared_radius = (
distortion_coefficient * distorted_squared_radius)
overflow_mask = tf.less(
4.0 / 27.0 + distortion_coefficient_times_distorted_squared_radius,
0.0)
# Newton-raphson iterations. The expression below is obtained from
# algebrically simplifying the Newton-Raphson formula
# (https://en.wikipedia.org/wiki/Newtons_method).
# We initialize with the approximate formula for the undistortion function
# given here https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4934233/.
undistortion = (1.0 -
distortion_coefficient_times_distorted_squared_radius)
for _ in range(num_iterations):
two_thirds_undistortion = 2.0 * undistortion / 3.0
undistortion = (1.0 - two_thirds_undistortion) / (
1.0 +
3.0 * distortion_coefficient_times_distorted_squared_radius *
undistortion * undistortion) + two_thirds_undistortion
return undistortion, overflow_mask
def interpolate(points,
weights,
indices,
normalize=True,
allow_negative_weights=False,
name=None):
"""Weighted interpolation for M-D point sets.
Given an M-D point set, this function can be used to generate a new point set
that is formed by interpolating a subset of points in the set.
Note:
In the following, A1 to An, and B1 to Bk are optional batch dimensions.
Args:
points: A tensor with shape `[B1, ..., Bk, M] and rank R > 1, where M is the
dimensionality of the points.
weights: A tensor with shape `[A1, ..., An, P]`, where P is the number of
points to interpolate for each output point.
indices: A tensor of dtype tf.int32 and shape `[A1, ..., An, P, R-1]`, which
contains the point indices to be used for each output point. The R-1
dimensional axis gives the slice index of a single point in `points`. The
first n+1 dimensions of weights and indices must match, or be broadcast
compatible.
normalize: A `bool` describing whether or not to normalize the weights on
the last axis.
allow_negative_weights: A `bool` describing whether or not negative weights
are allowed.
name: A name for this op. Defaults to "weighted_interpolate".
Returns:
A tensor of shape `[A1, ..., An, M]` storing the interpolated M-D
points. The first n dimensions will be the same as weights and indices.
"""
with tf.compat.v1.name_scope(name, "weighted_interpolate",
[points, weights, indices]):
points = tf.convert_to_tensor(value=points)
weights = tf.convert_to_tensor(value=weights)
indices = tf.convert_to_tensor(value=indices)
shape.check_static(tensor=points,
tensor_name="points",
has_rank_greater_than=1)
shape.check_static(tensor=indices,
tensor_name="indices",
has_rank_greater_than=1,
has_dim_equals=(-1, points.shape.ndims - 1))
shape.compare_dimensions(tensors=(weights, indices),
axes=(-1, -2),
tensor_names=("weights", "indices"))
shape.compare_batch_dimensions(tensors=(weights, indices),
last_axes=(-2, -3),
tensor_names=("weights", "indices"),
broadcast_compatible=True)
if not allow_negative_weights:
weights = asserts.assert_all_above(weights, 0.0, open_bound=False)
if normalize:
sums = tf.reduce_sum(input_tensor=weights, axis=-1, keepdims=True)
sums = asserts.assert_nonzero_norm(sums)
weights = safe_ops.safe_signed_div(weights, sums)
point_lists = tf.gather_nd(points, indices)
return vector.dot(point_lists,
tf.expand_dims(weights, axis=-1),
axis=-2,
keepdims=False)
