def conjugate(dual_quaternion, name=None):
"""Computes the conjugate of a dual quaternion.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
dual_quaternion: A tensor of shape `[A1, ..., An, 8]`, where the last
dimension represents a normalized dual quaternion.
name: A name for this op that defaults to "dual_quaternion_conjugate".
Returns:
A tensor of shape `[A1, ..., An, 8]`, where the last dimension represents
a normalized dual quaternion.
Raises:
ValueError: If the shape of `dual_quaternion` is not supported.
"""
with tf.compat.v1.name_scope(name, "dual_quaternion_conjugate",
[dual_quaternion]):
dual_quaternion = tf.convert_to_tensor(value=dual_quaternion)
shape.check_static(tensor=dual_quaternion,
tensor_name="dual_quaternion",
has_dim_equals=(-1, 8))
quaternion_real, quaternion_dual = tf.split(dual_quaternion, (4, 4),
axis=-1)
quaternion_real = asserts.assert_normalized(quaternion_real)
return tf.concat((quaternion.conjugate(quaternion_real),
quaternion.conjugate(quaternion_dual)),
axis=-1)
def from_quaternion(quaternion, name=None):
"""Converts a quaternion to an axis-angle representation.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
quaternion: A tensor of shape `[A1, ..., An, 4]`, where the last dimension
represents a normalized quaternion.
name: A name for this op that defaults to "axis_angle_from_quaternion".
Returns:
Tuple of two tensors 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 `quaternion` is not supported.
"""
with tf.compat.v1.name_scope(name, "axis_angle_from_quaternion",
[quaternion]):
quaternion = tf.convert_to_tensor(value=quaternion)
shape.check_static(tensor=quaternion,
tensor_name="quaternion",
has_dim_equals=(-1, 4))
quaternion = asserts.assert_normalized(quaternion)
# This prevents zero norm xyz and zero w, and is differentiable.
quaternion += asserts.select_eps_for_addition(quaternion.dtype)
xyz, w = tf.split(quaternion, (3, 1), axis=-1)
norm = tf.norm(tensor=xyz, axis=-1, keepdims=True)
angle = 2.0 * tf.atan2(norm, tf.abs(w))
axis = safe_ops.safe_unsigned_div(safe_ops.nonzero_sign(w) * xyz, norm)
return axis, angle
def relative_angle(quaternion1, quaternion2, name=None):
r"""Computes the unsigned relative rotation angle between 2 unit quaternions.
Given two normalized quanternions $$\mathbf{q}_1$$ and $$\mathbf{q}_2$$, the
relative angle is computed as
$$\theta = 2\arccos(\mathbf{q}_1^T\mathbf{q}_2)$$.
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 normalized quaternion.
quaternion2: A tensor of shape `[A1, ..., An, 4]`, where the last dimension
represents a normalized quaternion.
name: A name for this op that defaults to "quaternion_relative_angle".
Returns:
A tensor of shape `[A1, ..., An, 1]` where the last dimension represents
rotation angles in the range [0.0, pi].
Raises:
ValueError: If the shape of `quaternion1` or `quaternion2` is not supported.
"""
with (tf.compat.v1.name_scope(name, "quaternion_relative_angle",
[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))
quaternion1 = asserts.assert_normalized(quaternion1)
quaternion2 = asserts.assert_normalized(quaternion2)
dot_product = vector.dot(quaternion1, quaternion2, keepdims=False)
# Ensure dot product is in range [-1. 1].
eps_dot_prod = 4.0 * asserts.select_eps_for_addition(dot_product.dtype)
dot_product = safe_ops.safe_shrink(dot_product,
-1.0,
1.0,
False,
eps=eps_dot_prod)
return 2.0 * tf.acos(tf.abs(dot_product))
def from_axis_angle(axis, angle, name="rotation_matrix_3d_from_axis_angle"):
"""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.name_scope(name):
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 test_assert_normalized_exception_raised(self, dtype):
"""Checks that assert_normalized raises exceptions for invalid input."""
vector = _pick_random_vector() + 10.0
vector = tf.convert_to_tensor(value=vector, dtype=dtype)
vector = tf.abs(vector)
with self.assertRaises(tf.errors.InvalidArgumentError):
self.evaluate(asserts.assert_normalized(vector))
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 from_quaternion(
quaternion: type_alias.TensorLike,
name: str = "rotation_matrix_3d_from_quaternion") -> tf.Tensor:
"""Convert a quaternion to a rotation matrix.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
quaternion: A tensor of shape `[A1, ..., An, 4]`, where the last dimension
represents a normalized quaternion.
name: A name for this op that defaults to
"rotation_matrix_3d_from_quaternion".
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 `quaternion` is not supported.
"""
with tf.name_scope(name):
quaternion = tf.convert_to_tensor(value=quaternion)
shape.check_static(tensor=quaternion,
tensor_name="quaternion",
has_dim_equals=(-1, 4))
quaternion = asserts.assert_normalized(quaternion)
x, y, z, w = tf.unstack(quaternion, axis=-1)
tx = 2.0 * x
ty = 2.0 * y
tz = 2.0 * z
twx = tx * w
twy = ty * w
twz = tz * w
txx = tx * x
txy = ty * x
txz = tz * x
tyy = ty * y
tyz = tz * y
tzz = tz * z
matrix = tf.stack(
(1.0 - (tyy + tzz), txy - twz, txz + twy, txy + twz, 1.0 -
(txx + tzz), tyz - twx, txz - twy, tyz + twx, 1.0 - (txx + tyy)),
axis=-1) # pyformat: disable
output_shape = tf.concat((tf.shape(input=quaternion)[:-1], (3, 3)),
axis=-1)
return tf.reshape(matrix, shape=output_shape)
def distance_to_ray(point: type_alias.TensorLike,
origin: type_alias.TensorLike,
direction: type_alias.TensorLike,
keepdims: bool = True,
name: str = "point_distance_to_ray"
) -> tf.Tensor:
"""Computes the distance from a M-d point to a M-d ray.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
Args:
point: A tensor of shape `[A1, ..., An, M]`.
origin: A tensor of shape `[A1, ..., An, M]`.
direction: A tensor of shape `[A1, ..., An, M]`. The last dimension must be
normalized.
keepdims: A `bool`, whether to keep the last dimension with length 1 or to
remove it.
name: A name for this op. Defaults to "point_distance_to_ray".
Returns:
A tensor of shape `[A1, ..., An, 1]` containing the distance from each point
to the corresponding ray.
Raises:
ValueError: If the shape of `point`, `origin`, or 'direction' is not
supported.
"""
with tf.name_scope(name):
point = tf.convert_to_tensor(value=point)
origin = tf.convert_to_tensor(value=origin)
direction = tf.convert_to_tensor(value=direction)
shape.compare_dimensions((point, origin, direction), -1,
("point", "origin", "direction"))
shape.compare_batch_dimensions(
tensors=(point, origin, direction),
last_axes=-2,
broadcast_compatible=True)
direction = asserts.assert_normalized(direction)
vec = point - origin
dot = vector.dot(vec, direction)
vec -= dot * direction
return tf.norm(tensor=vec, axis=-1, keepdims=keepdims)
def rotate(point: type_alias.TensorLike,
quaternion: type_alias.TensorLike,
name: str = "quaternion_rotate") -> tf.Tensor:
"""Rotates a point using a quaternion.
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.
quaternion: A tensor of shape `[A1, ..., An, 4]`, where the last dimension
represents a normalized quaternion.
name: A name for this op that defaults to "quaternion_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 `quaternion` is not supported.
"""
with tf.name_scope(name):
point = tf.convert_to_tensor(value=point)
quaternion = tf.convert_to_tensor(value=quaternion)
shape.check_static(tensor=point,
tensor_name="point",
has_dim_equals=(-1, 3))
shape.check_static(tensor=quaternion,
tensor_name="quaternion",
has_dim_equals=(-1, 4))
shape.compare_batch_dimensions(tensors=(point, quaternion),
last_axes=-2,
broadcast_compatible=True)
quaternion = asserts.assert_normalized(quaternion)
padding = [[0, 0] for _ in range(point.shape.ndims)]
padding[-1][-1] = 1
point = tf.pad(tensor=point, paddings=padding, mode="CONSTANT")
point = multiply(quaternion, point)
point = multiply(point, conjugate(quaternion))
xyz, _ = tf.split(point, (3, 1), axis=-1)
return xyz
def from_axis_angle(axis: type_alias.TensorLike,
angle: type_alias.TensorLike,
name: str = "quaternion_from_axis_angle") -> tf.Tensor:
"""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.name_scope(name):
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 reflect(vector: TensorLike,
normal: TensorLike,
axis: int = -1,
name: str = "vector_reflect") -> TensorLike:
r"""Computes the reflection direction for an incident vector.
For an incident vector \\(\mathbf{v}\\) and normal $$\mathbf{n}$$ this
function computes the reflected vector as
\\(\mathbf{r} = \mathbf{v} - 2(\mathbf{n}^T\mathbf{v})\mathbf{n}\\).
Note:
In the following, A1 to An are optional batch dimensions, which should be
broadcast compatible.
Args:
vector: A tensor of shape `[A1, ..., Ai, ..., An]`, where the dimension i =
axis represents a vector.
normal: A tensor of shape `[A1, ..., Ai, ..., An]`, where the dimension i =
axis represents a normal around which the vector needs to be reflected.
The normal vector needs to be normalized.
axis: The dimension along which to compute the reflection.
name: A name for this op which defaults to "vector_reflect".
Returns:
A tensor of shape `[A1, ..., Ai, ..., An]`, where the dimension i = axis
represents a reflected vector.
"""
with tf.name_scope(name):
vector = tf.convert_to_tensor(value=vector)
normal = tf.convert_to_tensor(value=normal)
shape.compare_dimensions(tensors=(vector, normal),
axes=axis,
tensor_names=("vector", "normal"))
shape.compare_batch_dimensions(tensors=(vector, normal),
last_axes=-1,
broadcast_compatible=True)
normal = asserts.assert_normalized(normal, axis=axis)
dot_product = dot(vector, normal, axis=axis)
return vector - 2.0 * dot_product * normal
def project_to_ray(point: type_alias.TensorLike,
origin: type_alias.TensorLike,
direction: type_alias.TensorLike,
name: str = "point_project_to_ray"
) -> tf.Tensor:
"""Computes the projection of a M-d point on a M-d ray.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
Args:
point: A tensor of shape `[A1, ..., An, M]`.
origin: A tensor of shape `[A1, ..., An, M]`.
direction: A tensor of shape `[A1, ..., An, M]`. The last dimension must be
normalized.
name: A name for this op. Defaults to "point_project_to_ray".
Returns:
A tensor of shape `[A1, ..., An, M]` containing the projected point.
Raises:
ValueError: If the shape of `point`, `origin`, or 'direction' is not
supported.
"""
with tf.name_scope(name):
point = tf.convert_to_tensor(value=point)
origin = tf.convert_to_tensor(value=origin)
direction = tf.convert_to_tensor(value=direction)
shape.compare_dimensions((point, origin, direction), -1,
("point", "origin", "direction"))
shape.compare_batch_dimensions(
tensors=(point, origin, direction),
last_axes=-2,
broadcast_compatible=True)
direction = asserts.assert_normalized(direction)
vec = point - origin
dot = vector.dot(vec, direction)
return origin + dot * direction
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 interpolate_attributes(attribute, barycentric, name=None):
"""Interpolates attributes using barycentric weights.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
attribute: A tensor of shape `[A1, ..., An, 3, B]`, where the last dimension
stores a per-vertex `B`-dimensional attribute.
barycentric: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
contains barycentric coordinates.
name: A name for this op. Defaults to 'interpolate_attributes'.
Returns:
A tensor of shape `[A1, ..., An, B]`, containing interpolated attributes.
"""
with tf.compat.v1.name_scope(name, "interpolate_attributes",
(attribute, barycentric)):
attribute = tf.convert_to_tensor(value=attribute)
barycentric = tf.convert_to_tensor(value=barycentric)
shape.check_static(tensor=attribute,
tensor_name="attribute",
has_dim_equals=(-2, 3))
shape.check_static(tensor=barycentric,
tensor_name="barycentric",
has_dim_equals=(-1, 3))
shape.compare_batch_dimensions(tensors=(attribute, barycentric),
last_axes=(-2, -1),
tensor_names=("attribute",
"barycentric"),
broadcast_compatible=True)
barycentric = asserts.assert_normalized(barycentric, order=1)
return tf.reduce_sum(
input_tensor=tf.expand_dims(barycentric, axis=-1) * attribute,
axis=-2)
def brdf(direction_incoming_light: type_alias.TensorLike,
direction_outgoing_light: type_alias.TensorLike,
surface_normal: type_alias.TensorLike,
albedo: type_alias.TensorLike,
name: str = "lambertian_brdf") -> tf.Tensor:
"""Evaluates the brdf of a Lambertian surface.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
Note:
The gradient of this function is not smooth when the dot product of the
normal with any light is 0.0.
Args:
direction_incoming_light: A tensor of shape `[A1, ..., An, 3]`, where the
last dimension represents a normalized incoming light vector.
direction_outgoing_light: A tensor of shape `[A1, ..., An, 3]`, where the
last dimension represents a normalized outgoing light vector.
surface_normal: A tensor of shape `[A1, ..., An, 3]`, where the last
dimension represents a normalized surface normal.
albedo: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents albedo with values in [0,1].
name: A name for this op. Defaults to "lambertian_brdf".
Returns:
A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents
the amount of reflected light in any outgoing direction.
Raises:
ValueError: if the shape of `direction_incoming_light`,
`direction_outgoing_light`, `surface_normal`, `shininess` or `albedo` is not
supported.
InvalidArgumentError: if at least one element of `albedo` is outside of
[0,1].
"""
with tf.name_scope(name):
direction_incoming_light = tf.convert_to_tensor(
value=direction_incoming_light)
direction_outgoing_light = tf.convert_to_tensor(
value=direction_outgoing_light)
surface_normal = tf.convert_to_tensor(value=surface_normal)
albedo = tf.convert_to_tensor(value=albedo)
shape.check_static(tensor=direction_incoming_light,
tensor_name="direction_incoming_light",
has_dim_equals=(-1, 3))
shape.check_static(tensor=direction_outgoing_light,
tensor_name="direction_outgoing_light",
has_dim_equals=(-1, 3))
shape.check_static(tensor=surface_normal,
tensor_name="surface_normal",
has_dim_equals=(-1, 3))
shape.check_static(tensor=albedo,
tensor_name="albedo",
has_dim_equals=(-1, 3))
shape.compare_batch_dimensions(
tensors=(direction_incoming_light, direction_outgoing_light,
surface_normal, albedo),
tensor_names=("direction_incoming_light",
"direction_outgoing_light", "surface_normal",
"albedo"),
last_axes=-2,
broadcast_compatible=True)
direction_incoming_light = asserts.assert_normalized(
direction_incoming_light)
direction_outgoing_light = asserts.assert_normalized(
direction_outgoing_light)
surface_normal = asserts.assert_normalized(surface_normal)
albedo = asserts.assert_all_in_range(albedo,
0.0,
1.0,
open_bounds=False)
# Checks whether the incoming or outgoing light point behind the surface.
dot_incoming_light_surface_normal = vector.dot(
-direction_incoming_light, surface_normal)
dot_outgoing_light_surface_normal = vector.dot(
direction_outgoing_light, surface_normal)
min_dot = tf.minimum(dot_incoming_light_surface_normal,
dot_outgoing_light_surface_normal)
common_shape = shape.get_broadcasted_shape(min_dot.shape, albedo.shape)
d_val = lambda dim: 1 if dim is None else tf.compat.dimension_value(dim
)
common_shape = [d_val(dim) for dim in common_shape]
condition = tf.broadcast_to(tf.greater_equal(min_dot, 0.0),
common_shape)
albedo = tf.broadcast_to(albedo, common_shape)
return tf.where(condition, albedo / math.pi, tf.zeros_like(albedo))
def quaternion_weights(
quaternion1: type_alias.TensorLike,
quaternion2: type_alias.TensorLike,
percent: Union[type_alias.Float, type_alias.TensorLike],
eps: Optional[type_alias.Float] = None,
name: str = "quaternion_weights") -> Tuple[tf.Tensor, tf.Tensor]:
"""Calculates slerp weights for two normalized quaternions.
Given a percent and two normalized quaternions, this function returns the
slerp weights. It can also produce extrapolation weights when percent is
outside of the [0, 1] range. It reduces to lerp when input quaternions are
almost parallel or anti-parallel. Input quaternions are assumed to be
normalized. The tf.graphics debug flag TFG_ADD_ASSERTS_TO_GRAPH defined
in tfg_flags.py can be set to add assertions to the graph that check whether
the inputs are normalized, and whether Inf or Nan values are produced.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
quaternion1: A tensor of shape `[A1, ... , An, 4]` storing normalized
quaternions in its last dimension.
quaternion2: A tensor of shape `[A1, ... , An, 4]` storing normalized
quaternions in its last dimension.
percent: A `float` or a tensor with a shape broadcastable to the shape `[A1,
... , An]`.
eps: A `float` used to make operations safe. When left as None, the function
automatically picks the best epsilon based on the dtype and the operation.
name: A name for this op. Defaults to "quaternion_weights".
Raises:
ValueError: If the shapes of quaternions do not match, if the last
dimensions of quaternions are not 4, or if percent is neither a float, nor
a tensor with last dimension 1.
Returns:
Two tensors of shape `[A1, ... , An, 1]` each, which are the two slerp
weights for each quaternion.
"""
with tf.name_scope(name):
quaternion1 = tf.convert_to_tensor(value=quaternion1)
quaternion2 = tf.convert_to_tensor(value=quaternion2)
percent = tf.convert_to_tensor(value=percent, dtype=quaternion1.dtype)
if percent.shape.ndims == 0:
percent = tf.expand_dims(percent, axis=0)
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))
shape.compare_batch_dimensions(
tensors=(quaternion1, quaternion2, percent),
last_axes=(-2, -2, -1),
broadcast_compatible=True,
tensor_names=("quaternion1", "quaternion2", "percent"))
quaternion1 = asserts.assert_normalized(quaternion1)
quaternion2 = asserts.assert_normalized(quaternion2)
dot_product = _safe_dot(quaternion1, quaternion2, eps)
# Take the shorter path
theta = tf.acos(tf.abs(dot_product))
# safe_sinpx_div_sinx returns p for very small x, which means slerp reduces
# to lerp automatically.
scale1 = safe_ops.safe_sinpx_div_sinx(theta, 1.0 - percent, eps)
scale2 = safe_ops.safe_sinpx_div_sinx(theta, percent, eps)
# Flip the sign of scale1 if quaternions are in different hemispheres.
# tf.sign can make scale1 zero if quaternions are orthogonal.
scale1 *= safe_ops.nonzero_sign(dot_product)
return scale1, scale2
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 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 estimate_radiance(point_light_radiance,
point_light_position,
surface_point_position,
surface_point_normal,
observation_point,
brdf,
name="estimate_radiance",
reflected_light_fall_off=False):
"""Estimates the spectral radiance of a point light reflected from the surface point towards the observation point.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible.
B1 to Bm are optional batch dimensions for the lights, which must be
broadcast compatible.
Note:
In case the light or the observation point are located behind the surface
the function will return 0.
Note:
The gradient of this function is not smooth when the dot product of the
normal with the light-to-surface or surface-to-observation vectors is 0.
Args:
point_light_radiance: A tensor of shape '[B1, ..., Bm, K]', where the last
axis represents the radiance of the point light at a specific wave length.
point_light_position: A tensor of shape `[B1, ..., Bm, 3]`, where the last
axis represents the position of the point light.
surface_point_position: A tensor of shape `[A1, ..., An, 3]`, where the last
axis represents the position of the surface point.
surface_point_normal: A tensor of shape `[A1, ..., An, 3]`, where the last
axis represents the normalized surface normal at the given surface point.
observation_point: A tensor of shape `[A1, ..., An, 3]`, where the last axis
represents the observation point.
brdf: The BRDF of the surface as a function of: incoming_light_direction -
The incoming light direction as the last axis of a tensor with shape `[A1,
..., An, 3]`. outgoing_light_direction - The outgoing light direction as
the last axis of a tensor with shape `[A1, ..., An, 3]`.
surface_point_normal - The surface normal as the last axis of a tensor
with shape `[A1, ..., An, 3]`. Note - The BRDF should return a tensor of
size '[A1, ..., An, K]' where the last axis represents the amount of
reflected light in each wave length.
name: A name for this op. Defaults to "estimate_radiance".
reflected_light_fall_off: A boolean specifying whether or not to include the
fall off of the light reflected from the surface towards the observation
point in the calculation. Defaults to False.
Returns:
A tensor of shape `[A1, ..., An, B1, ..., Bm, K]`, where the last
axis represents the amount of light received at the observation point
after being reflected from the given surface point.
Raises:
ValueError: if the shape of `point_light_position`,
`surface_point_position`, `surface_point_normal`, or `observation_point` is
not supported.
InvalidArgumentError: if 'surface_point_normal' is not normalized.
"""
with tf.name_scope(name):
point_light_radiance = tf.convert_to_tensor(value=point_light_radiance)
point_light_position = tf.convert_to_tensor(value=point_light_position)
surface_point_position = tf.convert_to_tensor(value=surface_point_position)
surface_point_normal = tf.convert_to_tensor(value=surface_point_normal)
observation_point = tf.convert_to_tensor(value=observation_point)
shape.check_static(
tensor=point_light_position,
tensor_name="point_light_position",
has_dim_equals=(-1, 3))
shape.check_static(
tensor=surface_point_position,
tensor_name="surface_point_position",
has_dim_equals=(-1, 3))
shape.check_static(
tensor=surface_point_normal,
tensor_name="surface_point_normal",
has_dim_equals=(-1, 3))
shape.check_static(
tensor=observation_point,
tensor_name="observation_point",
has_dim_equals=(-1, 3))
shape.compare_batch_dimensions(
tensors=(surface_point_position, surface_point_normal,
observation_point),
tensor_names=("surface_point_position", "surface_point_normal",
"observation_point"),
last_axes=-2,
broadcast_compatible=True)
shape.compare_batch_dimensions(
tensors=(point_light_radiance, point_light_position),
tensor_names=("point_light_radiance", "point_light_position"),
last_axes=-2,
broadcast_compatible=True)
surface_point_normal = asserts.assert_normalized(surface_point_normal)
# Get the number of lights dimensions (B1,...,Bm).
lights_num_dimensions = max(
len(point_light_radiance.shape), len(point_light_position.shape)) - 1
# Reshape the other parameters so they can be broadcasted to the output of
# shape [A1,...,An, B1,...,Bm, K].
surface_point_position = tf.reshape(
surface_point_position,
surface_point_position.shape[:-1] + (1,) * lights_num_dimensions + (3,))
surface_point_normal = tf.reshape(
surface_point_normal,
surface_point_normal.shape[:-1] + (1,) * lights_num_dimensions + (3,))
observation_point = tf.reshape(
observation_point,
observation_point.shape[:-1] + (1,) * lights_num_dimensions + (3,))
light_to_surface_point = surface_point_position - point_light_position
distance_light_surface_point = tf.norm(
tensor=light_to_surface_point, axis=-1, keepdims=True)
incoming_light_direction = tf.math.l2_normalize(
light_to_surface_point, axis=-1)
surface_to_observation_point = observation_point - surface_point_position
outgoing_light_direction = tf.math.l2_normalize(
surface_to_observation_point, axis=-1)
brdf_value = brdf(incoming_light_direction, outgoing_light_direction,
surface_point_normal)
incoming_light_dot_surface_normal = vector.dot(-incoming_light_direction,
surface_point_normal)
outgoing_light_dot_surface_normal = vector.dot(outgoing_light_direction,
surface_point_normal)
estimated_radiance = (point_light_radiance * \
brdf_value * incoming_light_dot_surface_normal) / \
(4. * math.pi * tf.math.square(distance_light_surface_point))
if reflected_light_fall_off:
distance_surface_observation_point = tf.norm(
tensor=surface_to_observation_point, axis=-1, keepdims=True)
estimated_radiance = estimated_radiance / \
tf.math.square(distance_surface_observation_point)
# Create a condition for checking whether the light or observation point are
# behind the surface.
min_dot = tf.minimum(incoming_light_dot_surface_normal,
outgoing_light_dot_surface_normal)
common_shape = shape.get_broadcasted_shape(min_dot.shape,
estimated_radiance.shape)
d_val = lambda dim: 1 if dim is None else tf.compat.dimension_value(dim)
common_shape = [d_val(dim) for dim in common_shape]
condition = tf.broadcast_to(tf.greater_equal(min_dot, 0.0), common_shape)
return tf.where(condition, estimated_radiance,
tf.zeros_like(estimated_radiance))
