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 interpolate_with_weights(knots, weights, name=None):
"""Interpolates knots using knot weights.
Note:
In the following, A1 to An, and B1 to Bk are optional batch dimensions.
Args:
knots: A tensor with shape `[B1, ..., Bk, C]` containing knot values, where
`C` is the number of knots.
weights: A tensor with shape `[A1, ..., An, C]` containing dense weights for
the knots, where `C` is the number of knots.
name: A name for this op. Defaults to "bspline_interpolate_with_weights".
Returns:
A tensor with shape `[A1, ..., An, B1, ..., Bk]`, which is the result of
spline interpolation.
Raises:
ValueError: If the last dimension of knots and weights is not equal.
"""
with tf.compat.v1.name_scope(name, "bspline_interpolate_with_weights",
[knots, weights]):
knots = tf.convert_to_tensor(value=knots)
weights = tf.convert_to_tensor(value=weights)
shape.compare_dimensions(
tensors=(knots, weights), axes=-1, tensor_names=("knots", "weights"))
return tf.tensordot(weights, knots, (-1, -1))
def adjacency_from_edges(edges, weights, num_edges, num_vertices):
"""Returns a batched sparse 1-ring adj tensor from edge list tensor.
Args:
edges: [B, E, 2] `int32` tensor of edges, possibly 0 padded.
weights: [B, E] `float32` tensor of edge weights, possibly 0 padded.
num_edges: [B] `int32` tensor of number of valid edges per batch sample.
num_vertices: [B] `int32` tensor of number of valid vertices per batch
sample.
Returns:
adj: A batched SparseTensor of weighted adjacency graph, of
dense_shape [B, V, V] where V is max(num_vertices)
"""
edges = tf.convert_to_tensor(value=edges)
weights = tf.convert_to_tensor(value=weights)
num_edges = tf.convert_to_tensor(value=num_edges)
num_vertices = tf.convert_to_tensor(value=num_vertices)
if not edges.dtype.is_integer:
raise TypeError("'edges' must have an integer type.")
if not num_edges.dtype.is_integer:
raise TypeError("'num_edges' must have an integer type.")
if not num_vertices.dtype.is_integer:
raise TypeError("'num_vertices' must have an integer type.")
if not weights.dtype.is_floating:
raise TypeError("'weights' must have a floating type.")
shape.check_static(tensor=edges, tensor_name='edges', has_rank=3)
shape.check_static(tensor=weights, tensor_name='weights', has_rank=2)
shape.check_static(tensor=num_edges, tensor_name='num_edges', has_rank=1)
shape.check_static(tensor=num_vertices,
tensor_name='num_vertices',
has_rank=1)
shape.compare_dimensions(tensors=(edges, weights, num_edges, num_vertices),
tensor_names=('edges', 'weights', 'num_edges',
'num_vertices'),
axes=(-3, -2, -1, -1))
shape.compare_dimensions(tensors=(edges, weights),
tensor_names=('edges', 'weights'),
axes=(-2, -1))
batch_size = tf.shape(input=edges)[0]
max_num_vertices = tf.reduce_max(input_tensor=num_vertices)
max_num_edges = tf.shape(input=edges)[1]
batch_col = tf.reshape(tf.range(batch_size, dtype=edges.dtype), [-1, 1, 1])
batch_col = tf.tile(batch_col, [1, max_num_edges, 1])
batch_edges = tf.concat([batch_col, edges], axis=-1)
indices, _ = conv_utils.flatten_batch_to_2d(batch_edges, sizes=num_edges)
values, _ = conv_utils.flatten_batch_to_2d(tf.expand_dims(weights, -1),
sizes=num_edges)
values = tf.squeeze(values)
adjacency = tf.SparseTensor(
indices=tf.cast(indices, tf.int64),
values=values,
dense_shape=[batch_size, max_num_vertices, max_num_vertices])
adjacency = tf.sparse.reorder(adjacency)
return adjacency
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 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 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 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 check_valid_graph_convolution_input(data: type_alias.TensorLike,
neighbors: tf.sparse.SparseTensor,
sizes: type_alias.TensorLike):
"""Checks that the inputs are valid for graph convolution ops.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
data: A `float` tensor with shape `[A1, ..., An, V1, V2]`.
neighbors: A SparseTensor with the same type as `data` and with shape `[A1,
..., An, V1, V1]`.
sizes: An `int` tensor of shape `[A1, ..., An]`. Optional, can be `None`.
Raises:
TypeError: if the input types are invalid.
ValueError: if the input dimensions are invalid.
"""
if not data.dtype.is_floating:
raise TypeError("'data' must have a float type.")
if neighbors.dtype != data.dtype:
raise TypeError("'neighbors' and 'data' must have the same type.")
if sizes is not None and not sizes.dtype.is_integer:
raise TypeError("'sizes' must have an integer type.")
if not isinstance(neighbors, tf.sparse.SparseTensor):
raise ValueError("'neighbors' must be a SparseTensor.")
data_ndims = data.shape.ndims
shape.check_static(tensor=data,
tensor_name="data",
has_rank_greater_than=1)
shape.check_static(tensor=neighbors,
tensor_name="neighbors",
has_rank=data_ndims)
if not _is_dynamic_shape(tensors=(data, neighbors)):
shape.compare_dimensions(tensors=(data, neighbors, neighbors),
tensor_names=("data", "neighbors",
"neighbors"),
axes=(-2, -2, -1))
if sizes is None:
shape.compare_batch_dimensions(tensors=(data, neighbors),
tensor_names=("data", "neighbors"),
last_axes=-3,
broadcast_compatible=False)
else:
shape.check_static(tensor=sizes,
tensor_name="sizes",
has_rank=data_ndims - 2)
shape.compare_batch_dimensions(tensors=(data, neighbors, sizes),
tensor_names=("data", "neighbors",
"sizes"),
last_axes=(-3, -3, -1),
broadcast_compatible=False)
def partition_sums_2d(data, group_ids, row_weights=None, name=None):
"""Sum over subsets of rows in a 2-D tensor.
Args:
data: 2-D tensor with shape `[D1, D2]`.
group_ids: 1-D `int` tensor with shape `[D1]`.
row_weights: 1-D tensor with shape `[D1]`. Can be `None`.
name: A name for this op. Defaults to 'utils_partition_sums_2d'.
Returns:
A 2-D tensor with shape `[max(group_ids) + 1, D2]` where
`output[i, :] = sum(data[j, :] * weight[j] * 1(group_ids[j] == i)), 1(.)`
is the indicator function.
Raises:
ValueError: if the inputs have invalid dimensions or types.
"""
with tf.compat.v1.name_scope(name, "utils_partition_sums_2d",
[data, group_ids, row_weights]):
data = tf.convert_to_tensor(value=data)
group_ids = tf.convert_to_tensor(value=group_ids)
if not group_ids.dtype.is_integer:
raise TypeError("'group_ids' must be an integer tensor.")
elif group_ids.dtype != tf.int64:
group_ids = tf.cast(group_ids, dtype=tf.int64)
if row_weights is None:
row_weights = tf.ones_like(group_ids, dtype=data.dtype)
else:
row_weights = tf.convert_to_tensor(value=row_weights)
if row_weights.dtype != data.dtype:
raise TypeError(
"'data' and 'row_weights' must have the same type.")
shape.check_static(tensor=data, tensor_name="data", has_rank=2)
shape.check_static(tensor=group_ids,
tensor_name="group_ids",
has_rank=1)
shape.check_static(tensor=row_weights,
tensor_name="row_weights",
has_rank=1)
shape.compare_dimensions(tensors=(data, group_ids, row_weights),
tensor_names=("data", "group_ids",
"row_weights"),
axes=0)
num_rows = tf.size(input=group_ids, out_type=tf.int64)
sparse_indices = tf.stack((group_ids, tf.range(num_rows)), axis=1)
out_shape = (tf.reduce_max(input_tensor=group_ids) + 1, num_rows)
sparse = tf.SparseTensor(sparse_indices,
row_weights,
dense_shape=out_shape)
return tf.sparse.sparse_dense_matmul(sparse, data)
def metis_parent_to_perm(metis_parents, num_metis_parents, num_vertices):
shape.compare_dimensions(
tensors=(metis_parents, num_metis_parents, num_vertices),
tensor_names=('metis_parents', 'num_metis_parents', 'num_vertices'),
axes=(-2, -2, -1)
)
batch_size = tf.shape(input=num_metis_parents)[0]
layer_size = tf.shape(input=num_metis_parents)[1]
max_num_pool = tf.reduce_max(metis_parents) + 1
max_num_vertices = tf.reduce_max(num_vertices)
num_parents_per_batch = tf.reduce_sum(num_metis_parents, axis=1)
metis_parents, _ = conv_utils.flatten_batch_to_2d(metis_parents, sizes=num_parents_per_batch)
num_metis_parents, _ = conv_utils.flatten_batch_to_2d(num_metis_parents)
# Directly use metis correspondence as row indices
row_indices = metis_parents
"""
We first allocate the whole range from 0 to length of flattened metis parents(thus, row indices
have the same size with col indices. Then we want to minus a offset to obtain the actual col
indices in each layer(i.e. axis 0). We find the offset by cumulatively sum up each length and
repeat them for the corresponding length, so that we can perform subtraction with range and
offsets.
"""
slice_lengths = tf.cumsum(num_metis_parents)
# Note that, after cumsum, ith offset is actually the (i+1)th offset, so we apply
# the following line to re-align the offset.
slice_lengths -= num_metis_parents
slice_lengths = tf.repeat(slice_lengths, repeats=num_metis_parents)
col_indices = tf.range(tf.shape(metis_parents)[0])
col_indices -= slice_lengths
layer_indices = tf.range(layer_size)
layer_indices = tf.tile(layer_indices, repeats=batch_size)
layer_indices = tf.repeat(layer_indices, repeats=num_metis_parents)
batch_indices = tf.range(batch_size)
batch_indices = tf.repeat(batch_indices, repeats=num_parents_per_batch)
indices = tf.stack((batch_indices, layer_indices, row_indices, col_indices), axis=1)
values = tf.ones(tf.shape(indices)[0], dtype=tf.float32)
metis_pooling = tf.SparseTensor(
indices=tf.cast(indices, tf.int64),
values=values,
dense_shape=[batch_size, layer_size, max_num_pool, max_num_vertices]
)
metis_pooling = tf.sparse.reorder(metis_pooling)
return metis_pooling
def compute_density(density_values, distances, name=None):
"""Renders the density values (alpha) for points along a ray, as described in ["NeRF Representing Scenes as Neural Radiance Fields for View Synthesis"](https://github.com/bmild/nerf).
Note:
In the following, A1 to An are optional batch dimensions.
Args:
density_values: A tensor of shape `[A1, ..., An, N, 1]`,
where N are the samples on the ray.
distances: A tensor of shape `[A1, ..., An, N]` containing the distances
between the samples, where N are the samples on the ray.
name: A name for this op. Defaults to "ray_radiance".
Returns:
A tensor of shape `[A1, ..., An, 1]` for the estimated density values,
and a tensor of shape `[A1, ..., An, N]` for the sample weights.
"""
with tf.compat.v1.name_scope(name, "ray_density",
[density_values, distances]):
density_values = tf.convert_to_tensor(value=density_values)
distances = tf.convert_to_tensor(value=distances)
distances = tf.expand_dims(distances, -1)
shape.check_static(tensor=density_values,
tensor_name="density_values",
has_dim_equals=(-1, 1))
shape.check_static(tensor=density_values,
tensor_name="density_values",
has_rank_greater_than=1)
shape.check_static(tensor=distances,
tensor_name="distances",
has_rank_greater_than=1)
shape.compare_batch_dimensions(tensors=(density_values, distances),
tensor_names=("density_values",
"dists"),
last_axes=-3,
broadcast_compatible=True)
shape.compare_dimensions(tensors=(density_values, distances),
tensor_names=("density_values", "dists"),
axes=-2)
alpha = 1. - tf.exp(-density_values * distances)
alpha = tf.squeeze(alpha, -1)
ray_sample_weights = alpha * tf.math.cumprod(
1. - alpha + 1e-10, -1, exclusive=True)
ray_alpha = tf.expand_dims(tf.reduce_sum(ray_sample_weights, -1),
axis=-1)
return ray_alpha, ray_sample_weights
def 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 is_valid(matrix, atol=1e-3, name=None):
r"""Determines if a matrix in K-dimensions is a valid rotation matrix.
Determines if a matrix $$\mathbf{R}$$ is a valid rotation matrix by checking
that $$\mathbf{R}^T\mathbf{R} = \mathbf{I}$$ and $$\det(\mathbf{R}) = 1$$.
Note: In the following, A1 to An are optional batch dimensions.
Args:
matrix: A tensor of shape `[A1, ..., An, K, K]`, where the last two
dimensions represent a rotation matrix in K-dimensions.
atol: The absolute tolerance parameter.
name: A name for this op that defaults to "rotation_matrix_common_is_valid".
Returns:
A tensor of type `bool` and shape `[A1, ..., An, 1]` where False indicates
that the input is not a valid rotation matrix.
"""
with tf.compat.v1.name_scope(name, "rotation_matrix_common_is_valid",
[matrix]):
matrix = tf.convert_to_tensor(value=matrix)
shape.check_static(tensor=matrix,
tensor_name="matrix",
has_rank_greater_than=1)
shape.compare_dimensions(tensors=(matrix, matrix),
tensor_names=("matrix", "matrix"),
axes=(-1, -2))
distance_to_unit_determinant = tf.abs(tf.linalg.det(matrix) - 1.)
# Computes how far the product of the transposed rotation matrix with itself
# is from the identity matrix.
ndims = matrix.shape.ndims
permutation = list(range(ndims - 2)) + [ndims - 1, ndims - 2]
identity = tf.eye(tf.compat.v1.dimension_value(matrix.shape[-1]),
dtype=matrix.dtype)
difference_to_identity = tf.matmul(
tf.transpose(a=matrix, perm=permutation), matrix) - identity
norm_diff = tf.norm(tensor=difference_to_identity, axis=(-2, -1))
# Computes the mask of entries that satisfies all conditions.
mask = tf.logical_and(distance_to_unit_determinant < atol,
norm_diff < atol)
output = tf.compat.v1.where(
mask, tf.ones_like(distance_to_unit_determinant, dtype=bool),
tf.zeros_like(distance_to_unit_determinant, dtype=bool))
return tf.expand_dims(output, axis=-1)
def integration_product(harmonics1,
harmonics2,
keepdims=True,
name="spherical_harmonics_convolution"):
"""Computes the integral of harmonics1.harmonics2 over the sphere.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
harmonics1: A tensor of shape `[A1, ..., An, C]`, where the last dimension
represents spherical harmonics coefficients.
harmonics2: A tensor of shape `[A1, ..., An, C]`, where the last dimension
represents spherical harmonics coefficients.
keepdims: If True, retains reduced dimensions with length 1.
name: A name for this op. Defaults to "spherical_harmonics_convolution".
Returns:
A tensor of shape `[A1, ..., An]` containing scalar values resulting from
integrating the product of the spherical harmonics `harmonics1` and
`harmonics2`.
Raises:
ValueError: if the last dimension of `harmonics1` is different from the last
dimension of `harmonics2`.
"""
with tf.name_scope(name):
harmonics1 = tf.convert_to_tensor(value=harmonics1)
harmonics2 = tf.convert_to_tensor(value=harmonics2)
shape.compare_dimensions(
tensors=(harmonics1, harmonics2),
axes=-1,
tensor_names=("harmonics1", "harmonics2"))
shape.compare_batch_dimensions(
tensors=(harmonics1, harmonics2),
last_axes=-2,
tensor_names=("harmonics1", "harmonics2"),
broadcast_compatible=True)
return vector.dot(harmonics1, harmonics2, keepdims=keepdims)
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 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 match_intermediate_batch_dimensions(tensor1, tensor2):
"""Match the batch dimensions.
Args:
tensor1: A tensor of shape `[A1, M]`.
tensor2: A tensor of shape `[A1, ..., An, N]`.
Returns:
A tensor of shape `[A1, ..., An, M]`.
"""
shape.check_static(tensor=tensor1, tensor_name="tensor1", has_rank=2)
shape.check_static(tensor=tensor2,
tensor_name="tensor2",
has_rank_greater_than=1)
shape.compare_dimensions(tensors=(tensor1, tensor2),
tensor_names=("tensor1", "tensor2"),
axes=0)
shape1 = tf.shape(tensor1)
shape2 = tf.shape(tensor2)
shape_diff = len(shape2) - len(shape1)
new_shape = tf.concat([[shape1[0]], [1] * shape_diff, [shape1[-1]]],
axis=-1)
target_shape = tf.concat([shape2[:-1], [shape1[-1]]], axis=-1)
return tf.broadcast_to(tf.reshape(tensor1, new_shape), target_shape)
def 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)
def feature_steered_convolution(
data: type_alias.TensorLike,
neighbors: tf.sparse.SparseTensor,
sizes: type_alias.TensorLike,
var_u: type_alias.TensorLike,
var_v: type_alias.TensorLike,
var_c: type_alias.TensorLike,
var_w: type_alias.TensorLike,
var_b: type_alias.TensorLike,
name="graph_convolution_feature_steered_convolution") -> tf.Tensor:
# pyformat: disable
"""Implements the Feature Steered graph convolution.
FeaStNet: Feature-Steered Graph Convolutions for 3D Shape Analysis
Nitika Verma, Edmond Boyer, Jakob Verbeek
CVPR 2018
https://arxiv.org/abs/1706.05206
The shorthands used below are
`V`: The number of vertices.
`C`: The number of channels in the input data.
`D`: The number of channels in the output after convolution.
`W`: The number of weight matrices used in the convolution.
The input variables (`var_u`, `var_v`, `var_c`, `var_w`, `var_b`) correspond
to the variables with the same names in the paper cited above.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
data: A `float` tensor with shape `[A1, ..., An, V, C]`.
neighbors: A `SparseTensor` with the same type as `data` and with shape
`[A1, ..., An, V, V]` representing vertex neighborhoods. The neighborhood
of a vertex defines the support region for convolution. For a mesh, a
common choice for the neighborhood of vertex i would be the vertices in
the K-ring of i (including i itself). Each vertex must have at least one
neighbor. For a faithful implementation of the FeaStNet convolution,
neighbors should be a row-normalized weight matrix corresponding to the
graph adjacency matrix with self-edges: `neighbors[A1, ..., An, i, j] > 0`
if vertex j is a neighbor of i, and `neighbors[A1, ..., An, i, i] > 0` for
all i, and `sum(neighbors, axis=-1)[A1, ..., An, i] == 1.0 for all i`.
These requirements are relaxed in this implementation.
sizes: An `int` tensor of shape `[A1, ..., An]` indicating the true input
sizes in case of padding (`sizes=None` indicates no padding).Note that
`sizes[A1, ..., An] <= V`. If `data` and `neighbors` are 2-D, `sizes` will
be ignored. An example usage of `sizes`: consider an input consisting of
three graphs G0, G1, and G2 with V0, V1, and V2 vertices respectively. The
padded input would have the following shapes: `data.shape = [3, V, C]` and
`neighbors.shape = [3, V, V]`, where `V = max([V0, V1, V2])`. The true
sizes of each graph will be specified by `sizes=[V0, V1, V2]`,
`data[i, :Vi, :]` and `neighbors[i, :Vi, :Vi]` will be the vertex and
neighborhood data of graph Gi. The `SparseTensor` `neighbors` should have
no nonzero entries in the padded regions.
var_u: A 2-D tensor with shape `[C, W]`.
var_v: A 2-D tensor with shape `[C, W]`.
var_c: A 1-D tensor with shape `[W]`.
var_w: A 3-D tensor with shape `[W, C, D]`.
var_b: A 1-D tensor with shape `[D]`.
name: A name for this op. Defaults to
`graph_convolution_feature_steered_convolution`.
Returns:
Tensor with shape `[A1, ..., An, V, D]`.
Raises:
TypeError: if the input types are invalid.
ValueError: if the input dimensions are invalid.
"""
# pyformat: enable
with tf.name_scope(name):
data = tf.convert_to_tensor(value=data)
neighbors = tf.compat.v1.convert_to_tensor_or_sparse_tensor(
value=neighbors)
if sizes is not None:
sizes = tf.convert_to_tensor(value=sizes)
var_u = tf.convert_to_tensor(value=var_u)
var_v = tf.convert_to_tensor(value=var_v)
var_c = tf.convert_to_tensor(value=var_c)
var_w = tf.convert_to_tensor(value=var_w)
var_b = tf.convert_to_tensor(value=var_b)
data_ndims = data.shape.ndims
utils.check_valid_graph_convolution_input(data, neighbors, sizes)
shape.compare_dimensions(tensors=(data, var_u, var_v, var_w),
tensor_names=("data", "var_u", "var_v",
"var_w"),
axes=(-1, 0, 0, 1))
shape.compare_dimensions(tensors=(var_u, var_v, var_c, var_w),
tensor_names=("var_u", "var_v", "var_c",
"var_w"),
axes=(1, 1, 0, 0))
shape.compare_dimensions(tensors=(var_w, var_b),
tensor_names=("var_w", "var_b"),
axes=-1)
# Flatten the batch dimensions and remove any vertex padding.
if data_ndims > 2:
if sizes is not None:
sizes_square = tf.stack((sizes, sizes), axis=-1)
else:
sizes_square = None
x_flat, unflatten = utils.flatten_batch_to_2d(data, sizes)
adjacency = utils.convert_to_block_diag_2d(neighbors, sizes_square)
else:
x_flat = data
adjacency = neighbors
x_u = tf.matmul(x_flat, var_u)
x_v = tf.matmul(x_flat, var_v)
adjacency_ind_0 = adjacency.indices[:, 0]
adjacency_ind_1 = adjacency.indices[:, 1]
x_u_rep = tf.gather(x_u, adjacency_ind_0)
x_v_sep = tf.gather(x_v, adjacency_ind_1)
weights_q = tf.exp(x_u_rep + x_v_sep + tf.reshape(var_c, (1, -1)))
weights_q_sum = tf.reduce_sum(input_tensor=weights_q,
axis=-1,
keepdims=True)
weights_q = weights_q / weights_q_sum
y_i_m = []
x_sep = tf.gather(x_flat, adjacency_ind_1)
q_m_list = tf.unstack(weights_q, axis=-1)
w_m_list = tf.unstack(var_w, axis=0)
x_flat_shape = tf.shape(input=x_flat)
for q_m, w_m in zip(q_m_list, w_m_list):
# Compute `y_i_m = sum_{j in neighborhood(i)} q_m(x_i, x_j) * w_m * x_j`.
q_m = tf.expand_dims(q_m, axis=-1)
p_sum = tf.math.unsorted_segment_sum(
data=(q_m * x_sep) * tf.expand_dims(adjacency.values, -1),
segment_ids=adjacency_ind_0,
num_segments=x_flat_shape[0])
y_i_m.append(tf.matmul(p_sum, w_m))
y_out = tf.add_n(inputs=y_i_m) + tf.reshape(var_b, [1, -1])
if data_ndims > 2:
y_out = unflatten(y_out)
return y_out
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 sample_inverse_transform_stratified_1d(
ray_org: TensorLike,
ray_dir: TensorLike,
z_values_init: TensorLike,
weights_init: TensorLike,
n_samples: int,
combine_z_values=True,
name: str = "sample_inverse_transform_stratified_1d"):
"""Sample points on a ray using inverse transform stratified sampling.
The rays are defined by their origin and direction. Along each ray, there are
M samples (provided as 1D distances from the ray origin) and the corresponding
weights (probabilities) that facilitate the inverse transform sampling.
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_init: A tensor of shape `[A1, ..., An, M]`,
where the last dimension is the location of M points along the ray.
weights_init: A tensor of shape `[A1, ..., An, M]`,
where the last dimension is the density of M points along the ray.
n_samples: A number M to sample on the ray.
combine_z_values: Wether to combine the new 1D samples with
the initial points.
name: A name for this op that defaults to "stratified_sampling".
Returns:
A tensor of shape `[A1, ..., An, M, 3]` indicating the M points on the ray
and a tensor of shape `[A1, ..., An, M]` for the Z values on the points.
"""
with tf.name_scope(name):
shape.check_static(tensor=ray_org,
tensor_name="ray_org",
has_dim_equals=(-1, 3))
shape.check_static(tensor=ray_dir,
tensor_name="ray_dir",
has_dim_equals=(-1, 3))
shape.compare_batch_dimensions(tensors=(ray_org, ray_dir,
z_values_init, weights_init),
tensor_names=("ray_org", "ray_dir",
"z_values_init",
"weights_init"),
last_axes=-2,
broadcast_compatible=False)
shape.compare_dimensions(tensors=(z_values_init, weights_init),
tensor_names=("z_values_init",
"weights_init"),
axes=-1)
bin_start = z_values_init[..., :-1]
bin_width = z_values_init[..., 1:] - z_values_init[..., :-1]
bin_weights = .5 * (weights_init[..., 1:] + weights_init[..., :-1])
random_z_values = sampling.inverse_transform_stratified_1d(
bin_start, bin_width, bin_weights, n_samples)
random_z_values = tf.stop_gradient(random_z_values)
if combine_z_values:
z_values_final = tf.sort(
tf.concat([z_values_init, random_z_values], -1), -1)
else:
z_values_final = tf.sort(random_z_values, -1)
points3d = _points_from_z_values(ray_org, ray_dir, z_values_final)
return points3d, z_values_final
def inverse_transform_stratified_1d(bin_start: TensorLike,
bin_width: TensorLike,
pdf: TensorLike,
num_samples: int,
name="inverse_transform_stratified_1d"):
"""Stratified 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 (start and width of each bin).
The new samples can be sampled from anywhere inside the bin, unlike
inverse_transform_sampling_1d that returns the selected bin location.
Args:
bin_start: A tensor of shape `[A1, ..., An, M]` containing starting position
of M bins.
bin_width: A tensor of shape `[A1, ..., An, M]` containing the width of
M bins.
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_stratified".
Returns:
A tensor of shape `[A1, ..., An, N]` indicating the N points on the ray
"""
with tf.name_scope(name):
bin_start = tf.convert_to_tensor(value=bin_start)
bin_width = tf.convert_to_tensor(value=bin_width)
pdf = tf.convert_to_tensor(value=pdf)
shape.check_static(tensor=bin_start,
tensor_name="bin_start",
has_rank_greater_than=0)
shape.check_static(tensor=bin_width,
tensor_name="bin_width",
has_rank_greater_than=0)
shape.check_static(tensor=pdf,
tensor_name="pdf",
has_rank_greater_than=0)
shape.compare_batch_dimensions(tensors=(bin_start, pdf, bin_width),
tensor_names=("bins", "pdf",
"bin_width"),
last_axes=-2,
broadcast_compatible=True)
shape.compare_dimensions(tensors=(bin_start, pdf, bin_width),
tensor_names=("bins", "pdf", "bin_width"),
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]
batch_dims = batch_shape.get_shape().as_list()[0]
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")
below_bin_id = tf.maximum(0, bin_indices - 1)
above_bin_id = tf.minimum(cdf.shape[-1] - 1, bin_indices)
below_bin_cdf = tf.gather(cdf,
below_bin_id,
axis=-1,
batch_dims=batch_dims)
above_bin_cdf = tf.gather(cdf,
above_bin_id,
axis=-1,
batch_dims=batch_dims)
bin_prob = above_bin_cdf - below_bin_cdf
bin_prob = tf.where(bin_prob < 1e-5, tf.ones_like(bin_prob), bin_prob)
below_bin = tf.gather(bin_start,
below_bin_id,
axis=-1,
batch_dims=batch_dims)
bin_width = tf.gather(bin_width,
below_bin_id,
axis=-1,
batch_dims=batch_dims)
return below_bin + (uniform_samples -
below_bin_cdf) / bin_prob * bin_width
def blend(points,
skinning_weights,
bone_rotations,
bone_translations,
name=None):
"""Transforms the points using Linear Blend Skinning.
Note:
In the following, A1 to An are optional batch dimensions, which must be
broadcast compatible and allow transforming full 3D shapes at once.
In the following, B1 to Bm are optional batch dimensions, which allow
transforming multiple poses at once.
Args:
points: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
represents a 3d point.
skinning_weights: A tensor of shape `[A1, ..., An, W]`, where the last
dimension represents the skinning weights of each bone.
bone_rotations: A tensor of shape `[B1, ..., Bm, W, 3, 3]`, which represents
the 3d rotations applied to each bone.
bone_translations: A tensor of shape `[B1, ..., Bm, W, 3]`, which represents
the 3d translation vectors applied to each bone.
name: A name for this op that defaults to "linear_blend_skinning_blend".
Returns:
A tensor of shape `[B1, ..., Bm, A1, ..., An, 3]`, where the last dimension
represents a 3d point.
Raises:
ValueError: If the shape of the input tensors are not supported.
"""
with tf.compat.v1.name_scope(
name, "linear_blend_skinning_blend",
[points, skinning_weights, bone_rotations, bone_translations]):
points = tf.convert_to_tensor(value=points)
skinning_weights = tf.convert_to_tensor(value=skinning_weights)
bone_rotations = tf.convert_to_tensor(value=bone_rotations)
bone_translations = tf.convert_to_tensor(value=bone_translations)
shape.check_static(
tensor=points, tensor_name="points", has_dim_equals=(-1, 3))
shape.check_static(
tensor=bone_rotations,
tensor_name="bone_rotations",
has_rank_greater_than=2,
has_dim_equals=((-2, 3), (-1, 3)))
shape.check_static(
tensor=bone_translations,
tensor_name="bone_translations",
has_rank_greater_than=1,
has_dim_equals=(-1, 3))
shape.compare_dimensions(
tensors=(skinning_weights, bone_rotations),
tensor_names=("skinning_weights", "bone_rotations"),
axes=(-1, -3))
shape.compare_dimensions(
tensors=(skinning_weights, bone_translations),
tensor_names=("skinning_weights", "bone_translations"),
axes=(-1, -2))
shape.compare_batch_dimensions(
tensors=(points, skinning_weights),
tensor_names=("points", "skinning_weights"),
last_axes=(-2, -2),
broadcast_compatible=True)
shape.compare_batch_dimensions(
tensors=(bone_rotations, bone_translations),
tensor_names=("bone_rotations", "bone_translations"),
last_axes=(-3, -2),
broadcast_compatible=True)
num_bones = skinning_weights.shape[-1]
def dim_value(dim):
return 1 if dim is None else tf.compat.v1.dimension_value(dim)
# TODO(b/148362025): factorize this block out
points_batch_shape = shape.get_broadcasted_shape(
points.shape[:-1], skinning_weights.shape[:-1])
points_batch_shape = [dim_value(dim) for dim in points_batch_shape]
points = tf.broadcast_to(points, points_batch_shape + [3])
skinning_weights = tf.broadcast_to(skinning_weights,
points_batch_shape + [num_bones])
bones_batch_shape = shape.get_broadcasted_shape(
bone_rotations.shape[:-3], bone_translations.shape[:-2])
bones_batch_shape = [dim_value(dim) for dim in bones_batch_shape]
bone_rotations = tf.broadcast_to(bone_rotations,
bones_batch_shape + [num_bones, 3, 3])
bone_translations = tf.broadcast_to(bone_translations,
bones_batch_shape + [num_bones, 3])
points_batch_dims = points.shape.ndims - 1
bones_batch_dims = bone_rotations.shape.ndims - 3
points = tf.reshape(points,
[1] * bones_batch_dims + points_batch_shape + [1, 3])
skinning_weights = tf.reshape(skinning_weights, [1] * bones_batch_dims +
points_batch_shape + [num_bones, 1])
bone_rotations = tf.reshape(
bone_rotations,
bones_batch_shape + [1] * points_batch_dims + [num_bones, 3, 3])
bone_translations = tf.reshape(
bone_translations,
bones_batch_shape + [1] * points_batch_dims + [num_bones, 3])
transformed_points = rotation_matrix_3d.rotate(
points, bone_rotations) + bone_translations
weighted_points = tf.multiply(skinning_weights, transformed_points)
blended_points = tf.reduce_sum(input_tensor=weighted_points, axis=-2)
return blended_points
def energy(vertices_rest_pose,
vertices_deformed_pose,
quaternions,
edges,
vertex_weight=None,
edge_weight=None,
conformal_energy=True,
aggregate_loss=True,
name=None):
"""Estimates an As Conformal As Possible (ACAP) fitting energy.
For a given mesh in rest pose, this function evaluates a variant of the ACAP
[1] fitting energy for a batch of deformed meshes. The vertex weights and edge
weights are defined on the rest pose.
The method implemented here is similar to [2], but with an added free variable
capturing a scale factor per vertex.
[1]: Yusuke Yoshiyasu, Wan-Chun Ma, Eiichi Yoshida, and Fumio Kanehiro.
"As-Conformal-As-Possible Surface Registration." Computer Graphics Forum. Vol.
33. No. 5. 2014.</br>
[2]: Olga Sorkine, and Marc Alexa.
"As-rigid-as-possible surface modeling". Symposium on Geometry Processing.
Vol. 4. 2007.
Note:
In the description of the arguments, V corresponds to
the number of vertices in the mesh, and E to the number of edges in this
mesh.
Note:
In the following, A1 to An are optional batch dimensions.
Args:
vertices_rest_pose: A tensor of shape `[V, 3]` containing the position of
all the vertices of the mesh in rest pose.
vertices_deformed_pose: A tensor of shape `[A1, ..., An, V, 3]` containing
the position of all the vertices of the mesh in deformed pose.
quaternions: A tensor of shape `[A1, ..., An, V, 4]` defining a rigid
transformation to apply to each vertex of the rest pose. See Section 2
from [1] for further details.
edges: A tensor of shape `[E, 2]` defining indices of vertices that are
connected by an edge.
vertex_weight: An optional tensor of shape `[V]` defining the weight
associated with each vertex. Defaults to a tensor of ones.
edge_weight: A tensor of shape `[E]` defining the weight of edges. Common
choices for these weights include uniform weighting, and cotangent
weights. Defaults to a tensor of ones.
conformal_energy: A `bool` indicating whether each vertex is associated with
a scale factor or not. If this parameter is True, scaling information must
be encoded in the norm of `quaternions`. If this parameter is False, this
function implements the energy described in [2].
aggregate_loss: A `bool` defining whether the returned loss should be an
aggregate measure. When True, the mean squared error is returned. When
False, returns two losses for every edge of the mesh.
name: A name for this op. Defaults to "as_conformal_as_possible_energy".
Returns:
When aggregate_loss is `True`, returns a tensor of shape `[A1, ..., An]`
containing the ACAP energies. When aggregate_loss is `False`, returns a
tensor of shape `[A1, ..., An, 2*E]` containing each term of the summation
described in the equation 7 of [2].
Raises:
ValueError: if the shape of `vertices_rest_pose`, `vertices_deformed_pose`,
`quaternions`, `edges`, `vertex_weight`, or `edge_weight` is not supported.
"""
with tf.compat.v1.name_scope(name, "as_conformal_as_possible_energy", [
vertices_rest_pose, vertices_deformed_pose, quaternions, edges,
conformal_energy, vertex_weight, edge_weight
]):
vertices_rest_pose = tf.convert_to_tensor(value=vertices_rest_pose)
vertices_deformed_pose = tf.convert_to_tensor(
value=vertices_deformed_pose)
quaternions = tf.convert_to_tensor(value=quaternions)
edges = tf.convert_to_tensor(value=edges)
if vertex_weight is not None:
vertex_weight = tf.convert_to_tensor(value=vertex_weight)
if edge_weight is not None:
edge_weight = tf.convert_to_tensor(value=edge_weight)
shape.check_static(tensor=vertices_rest_pose,
tensor_name="vertices_rest_pose",
has_rank=2,
has_dim_equals=(-1, 3))
shape.check_static(tensor=vertices_deformed_pose,
tensor_name="vertices_deformed_pose",
has_rank_greater_than=1,
has_dim_equals=(-1, 3))
shape.check_static(tensor=quaternions,
tensor_name="quaternions",
has_rank_greater_than=1,
has_dim_equals=(-1, 4))
shape.compare_batch_dimensions(tensors=(vertices_deformed_pose,
quaternions),
last_axes=(-3, -3),
broadcast_compatible=False)
shape.check_static(tensor=edges,
tensor_name="edges",
has_rank=2,
has_dim_equals=(-1, 2))
tensors_with_vertices = [
vertices_rest_pose, vertices_deformed_pose, quaternions
]
names_with_vertices = [
"vertices_rest_pose", "vertices_deformed_pose", "quaternions"
]
axes_with_vertices = [-2, -2, -2]
if vertex_weight is not None:
shape.check_static(tensor=vertex_weight,
tensor_name="vertex_weight",
has_rank=1)
tensors_with_vertices.append(vertex_weight)
names_with_vertices.append("vertex_weight")
axes_with_vertices.append(0)
shape.compare_dimensions(tensors=tensors_with_vertices,
axes=axes_with_vertices,
tensor_names=names_with_vertices)
if edge_weight is not None:
shape.check_static(tensor=edge_weight,
tensor_name="edge_weight",
has_rank=1)
shape.compare_dimensions(tensors=(edges, edge_weight),
axes=(0, 0),
tensor_names=("edges", "edge_weight"))
if not conformal_energy:
quaternions = quaternion.normalize(quaternions)
# Extracts the indices of vertices.
indices_i, indices_j = tf.unstack(edges, axis=-1)
# Extracts the vertices we need per term.
vertices_i_rest = tf.gather(vertices_rest_pose, indices_i, axis=-2)
vertices_j_rest = tf.gather(vertices_rest_pose, indices_j, axis=-2)
vertices_i_deformed = tf.gather(vertices_deformed_pose,
indices_i,
axis=-2)
vertices_j_deformed = tf.gather(vertices_deformed_pose,
indices_j,
axis=-2)
# Extracts the weights we need per term.
weights_shape = vertices_i_rest.shape.as_list()[-2]
if vertex_weight is not None:
weight_i = tf.gather(vertex_weight, indices_i)
weight_j = tf.gather(vertex_weight, indices_j)
else:
weight_i = weight_j = tf.ones(weights_shape,
dtype=vertices_rest_pose.dtype)
weight_i = tf.expand_dims(weight_i, axis=-1)
weight_j = tf.expand_dims(weight_j, axis=-1)
if edge_weight is not None:
weight_ij = edge_weight
else:
weight_ij = tf.ones(weights_shape, dtype=vertices_rest_pose.dtype)
weight_ij = tf.expand_dims(weight_ij, axis=-1)
# Extracts the rotation we need per term.
quaternion_i = tf.gather(quaternions, indices_i, axis=-2)
quaternion_j = tf.gather(quaternions, indices_j, axis=-2)
# Computes the energy.
deformed_ij = vertices_i_deformed - vertices_j_deformed
rotated_rest_ij = quaternion.rotate(
(vertices_i_rest - vertices_j_rest), quaternion_i)
energy_ij = weight_i * weight_ij * (deformed_ij - rotated_rest_ij)
deformed_ji = vertices_j_deformed - vertices_i_deformed
rotated_rest_ji = quaternion.rotate(
(vertices_j_rest - vertices_i_rest), quaternion_j)
energy_ji = weight_j * weight_ij * (deformed_ji - rotated_rest_ji)
energy_ij_squared = vector.dot(energy_ij, energy_ij, keepdims=False)
energy_ji_squared = vector.dot(energy_ji, energy_ji, keepdims=False)
if aggregate_loss:
average_energy_ij = tf.reduce_mean(input_tensor=energy_ij_squared,
axis=-1)
average_energy_ji = tf.reduce_mean(input_tensor=energy_ji_squared,
axis=-1)
return (average_energy_ij + average_energy_ji) / 2.0
return tf.concat((energy_ij_squared, energy_ji_squared), axis=-1)
def generate(starts, stops, nums, name=None):
r"""Generates a M-D uniform axis-aligned grid.
Warning:
This op is not differentiable. Indeed, the gradient of tf.linspace and
tf.meshgrid are currently not defined.
Note:
In the following, `B` is an optional batch dimension.
Args:
starts: A tensor of shape `[M]` or `[B, M]`, where the last dimension
represents a M-D start point.
stops: A tensor of shape `[M]` or `[B, M]`, where the last dimension
represents a M-D end point.
nums: A tensor of shape `[M]` representing the number of subdivisions for
each dimension.
name: A name for this op. Defaults to "grid_generate".
Returns:
A tensor of shape `[nums[0], ..., nums[M-1], M]` containing an M-D uniform
grid or a tensor of shape [B, nums[0], ..., nums[M-1], M]` containing B
M-D uniform grids. Please refer to the example below for more details.
Raises:
ValueError: If the shape of `starts`, `stops`, or 'nums' is not supported.
Examples:
```python
print(generate((-1.0, -2.0), (1.0, 2.0), (3, 5)))
>>> [[[-1. -2.]
[-1. -1.]
[-1. 0.]
[-1. 1.]
[-1. 2.]]
[[ 0. -2.]
[ 0. -1.]
[ 0. 0.]
[ 0. 1.]
[ 0. 2.]]
[[ 1. -2.]
[ 1. -1.]
[ 1. 0.]
[ 1. 1.]
[ 1. 2.]]]
```
Generates a 3x5 2d grid from -1.0 to 1.0 with 3 subdivisions for the x
axis and from -2.0 to 2.0 with 5 subdivisions for the y axis. This lead to a
tensor of shape (3, 5, 2).
"""
with tf.compat.v1.name_scope(name, "grid_generate", [starts, stops, nums]):
starts = tf.convert_to_tensor(value=starts)
stops = tf.convert_to_tensor(value=stops)
nums = tf.convert_to_tensor(value=nums)
shape.check_static(tensor=starts,
tensor_name="starts",
has_rank_greater_than=0,
has_rank_less_than=3)
shape.check_static(tensor=stops,
tensor_name="stops",
has_rank_greater_than=0,
has_rank_less_than=3)
shape.check_static(tensor=nums, tensor_name="nums", has_rank=1)
shape.compare_batch_dimensions(tensors=(starts, stops),
last_axes=(-1, -1),
broadcast_compatible=False)
shape.compare_dimensions((starts, stops, nums), -1,
("starts", "stops", "nums"))
if starts.shape.ndims == 1:
return _grid(starts, stops, nums)
else:
return tf.stack([
_grid(starts, stops, nums)
for starts, stops in zip(tf.unstack(starts), tf.unstack(stops))
])
