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)) ])