Implementation of alpha rendering for samples on a ray using only the density information.

PiperOrigin-RevId: 371591566

Author: krematas

tensorflow_graphics/rendering/volumetric/ray_density.py

@@ -0,0 +1,79 @@
+# Copyright 2020 The TensorFlow Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""This module implements the radiance-based ray rendering."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+from tensorflow_graphics.util import export_api
+from tensorflow_graphics.util import shape
+
+
+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
+
+# API contains all public functions and classes.
+__all__ = export_api.get_functions_and_classes()

tensorflow_graphics/rendering/volumetric/tests/ray_density_test.py

@@ -0,0 +1,84 @@
+# Copyright 2020 The TensorFlow Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for radiance ray rendering."""
+
+from absl.testing import flagsaver
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow_graphics.rendering.volumetric import ray_density
+from tensorflow_graphics.util import test_case
+
+
+def generate_random_test_ray_render():
+  """Generates random test for the voxels rendering functions."""
+  batch_shape = np.random.randint(1, 3)
+  n_rays = np.random.randint(1, 512)
+  random_ray_values = np.random.uniform(size=[batch_shape] + [n_rays, 1])
+  random_ray_dists = random_ray_values[..., -1]
+  return random_ray_values, random_ray_dists
+
+
+class RayDensityTest(test_case.TestCase):
+
+  @parameterized.parameters(
+      ((6, 1), (6,)),
+      ((8, 16, 44, 1), (8, 16, 44)),
+      ((12, 8, 16, 22, 1), (12, 8, 16, 22)),
+      ((32, 32, 256, 1), (32, 32, 256)),
+      ((32, 32, 256, 1), (1, 1, 256)),
+      ((32, 32, 256, 1), (256,)),
+  )
+  def test_render_shape_exception_not_raised(self, *shapes):
+    """Tests that the shape exceptions are not raised."""
+    self.assert_exception_is_not_raised(ray_density.compute_density, shapes)
+
+  @parameterized.parameters(
+      ("must have a rank greater than 1", ((1,), (3,))),
+      ("must have exactly 1 dimensions in axis -1", ((44, 4), (44, 1))),
+      ("Not all batch dimensions are broadcast-compatible.",
+       ((32, 32, 256, 1,), (32, 16, 256,))),
+      ("must have the same number of dimensions",
+       ((32, 32, 128, 1,), (32, 32, 555,))),
+  )
+  def test_render_shape_exception_raised(self, error_msg, shape):
+    """Tests that the shape exception is raised."""
+    self.assert_exception_is_raised(ray_density.compute_density,
+                                    error_msg, shape)
+
+  @flagsaver.flagsaver(tfg_add_asserts_to_graph=False)
+  def test_render_jacobian_random(self):
+    """Tests the Jacobian of render."""
+    point_values, point_distance = generate_random_test_ray_render()
+    self.assert_jacobian_is_correct_fn(
+        lambda x: ray_density.compute_density(x, point_distance)[0],
+        [point_values])
+    self.assert_jacobian_is_correct_fn(
+        lambda x: ray_density.compute_density(point_values, x)[0],
+        [point_distance])
+
+  def test_render_preset(self):
+    """Checks that render returns the expected value."""
+
+    image_rays = np.zeros((128, 128, 64, 1))
+    image_rays[32:96, 32:96, 16:32, :] = 1
+    distances = np.zeros((128, 128, 64)) + 1.5
+    target_image = np.zeros((128, 128, 1))
+    target_image[32:96, 32:96, :] = 1
+    rendered_image, *_ = ray_density.compute_density(image_rays, distances)
+    self.assertAllClose(rendered_image, target_image)
+
+
+if __name__ == "__main__":
+  test_case.main()