Raymarching

Summary: Implements two basic raymarchers.

Reviewed By: gkioxari

Differential Revision: D24064250

fbshipit-source-id: 18071bd039995336b7410caa403ea29fafb5c66f

tests/test_raymarching.py

@@ -0,0 +1,196 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+
+import unittest
+
+import torch
+from common_testing import TestCaseMixin
+from pytorch3d.renderer import AbsorptionOnlyRaymarcher, EmissionAbsorptionRaymarcher
+
+
+class TestRaymarching(TestCaseMixin, unittest.TestCase):
+    def setUp(self) -> None:
+        torch.manual_seed(42)
+
+    @staticmethod
+    def _init_random_rays(
+        n_rays=10, n_pts_per_ray=9, device="cuda", dtype=torch.float32
+    ):
+        """
+        Generate a batch of ray points with features, densities, and z-coodinates
+        such that their EmissionAbsorption renderring results in
+        feature renders `features_gt`, depth renders `depths_gt`,
+        and opacity renders `opacities_gt`.
+        """
+
+        # generate trivial ray z-coordinates of sampled points coinciding with
+        # each point's order along a ray.
+        rays_z = torch.arange(n_pts_per_ray, dtype=dtype, device=device)[None].repeat(
+            n_rays, 1
+        )
+
+        # generate ground truth depth values of the underlying surface.
+        depths_gt = torch.randint(
+            low=1, high=n_pts_per_ray + 2, size=(n_rays,)
+        ).type_as(rays_z)
+
+        # compute ideal densities that are 0 before the surface and 1 after
+        # the corresponding ground truth depth value
+        rays_densities = (rays_z >= depths_gt[..., None]).type_as(rays_z)[..., None]
+        opacities_gt = (depths_gt < n_pts_per_ray).type_as(rays_z)
+
+        # generate random per-ray features
+        rays_features = torch.rand(
+            (n_rays, n_pts_per_ray, 3), device=rays_z.device, dtype=rays_z.dtype
+        )
+
+        # infer the expected feature render "features_gt"
+        gt_surface = ((rays_z - depths_gt[..., None]).abs() <= 1e-4).type_as(rays_z)
+        features_gt = (rays_features * gt_surface[..., None]).sum(dim=-2)
+
+        return (
+            rays_z,
+            rays_densities,
+            rays_features,
+            depths_gt,
+            features_gt,
+            opacities_gt,
+        )
+
+    @staticmethod
+    def raymarcher(
+        raymarcher_type=EmissionAbsorptionRaymarcher, n_rays=10, n_pts_per_ray=10
+    ):
+        (
+            rays_z,
+            rays_densities,
+            rays_features,
+            depths_gt,
+            features_gt,
+            opacities_gt,
+        ) = TestRaymarching._init_random_rays(
+            n_rays=n_rays, n_pts_per_ray=n_pts_per_ray
+        )
+
+        raymarcher = raymarcher_type()
+
+        def run_raymarcher():
+            raymarcher(
+                rays_densities=rays_densities,
+                rays_features=rays_features,
+                rays_z=rays_z,
+            )
+            torch.cuda.synchronize()
+
+        return run_raymarcher
+
+    def test_emission_absorption_inputs(self):
+        """
+        Test the checks of validity of the inputs to `EmissionAbsorptionRaymarcher`.
+        """
+
+        # init the EA raymarcher
+        raymarcher_ea = EmissionAbsorptionRaymarcher()
+
+        # bad ways of passing densities and features
+        # [rays_densities, rays_features, rays_z]
+        bad_inputs = [
+            [torch.rand(10, 5, 4), None],
+            [torch.Tensor(3)[0], torch.rand(10, 5, 4)],
+            [1.0, torch.rand(10, 5, 4)],
+            [torch.rand(10, 5, 4), 1.0],
+            [torch.rand(10, 5, 4), None],
+            [torch.rand(10, 5, 4), torch.rand(10, 5, 4)],
+            [torch.rand(10, 5, 4), torch.rand(10, 5, 4, 3)],
+            [torch.rand(10, 5, 4, 3), torch.rand(10, 5, 4, 3)],
+        ]
+
+        for bad_input in bad_inputs:
+            with self.assertRaises(ValueError):
+                raymarcher_ea(*bad_input)
+
+    def test_absorption_only_inputs(self):
+        """
+        Test the checks of validity of the inputs to `AbsorptionOnlyRaymarcher`.
+        """
+
+        # init the AO raymarcher
+        raymarcher_ao = AbsorptionOnlyRaymarcher()
+
+        # bad ways of passing densities and features
+        # [rays_densities, rays_features, rays_z]
+        bad_inputs = [[torch.Tensor(3)[0]]]
+
+        for bad_input in bad_inputs:
+            with self.assertRaises(ValueError):
+                raymarcher_ao(*bad_input)
+
+    def test_emission_absorption(self):
+        """
+        Test the EA raymarching algorithm.
+        """
+        (
+            rays_z,
+            rays_densities,
+            rays_features,
+            depths_gt,
+            features_gt,
+            opacities_gt,
+        ) = TestRaymarching._init_random_rays(
+            n_rays=1000, n_pts_per_ray=9, device=None, dtype=torch.float32
+        )
+
+        # init the EA raymarcher
+        raymarcher_ea = EmissionAbsorptionRaymarcher()
+
+        # allow gradients for a differentiability check
+        rays_densities.requires_grad = True
+        rays_features.requires_grad = True
+
+        # render the features first and check with gt
+        data_render = raymarcher_ea(rays_densities, rays_features)
+        features_render, opacities_render = data_render[..., :-1], data_render[..., -1]
+        self.assertClose(opacities_render, opacities_gt)
+        self.assertClose(
+            features_render * opacities_render[..., None],
+            features_gt * opacities_gt[..., None],
+        )
+
+        # get the depth map by rendering the ray z components and check with gt
+        depths_render = raymarcher_ea(rays_densities, rays_z[..., None])[..., 0]
+        self.assertClose(depths_render * opacities_render, depths_gt * opacities_gt)
+
+        # check differentiability
+        loss = features_render.mean()
+        loss.backward()
+        for field in (rays_densities, rays_features):
+            self.assertTrue(field.grad.data.isfinite().all())
+
+    def test_absorption_only(self):
+        """
+        Test the AO raymarching algorithm.
+        """
+        (
+            rays_z,
+            rays_densities,
+            rays_features,
+            depths_gt,
+            features_gt,
+            opacities_gt,
+        ) = TestRaymarching._init_random_rays(
+            n_rays=1000, n_pts_per_ray=9, device=None, dtype=torch.float32
+        )
+
+        # init the AO raymarcher
+        raymarcher_ao = AbsorptionOnlyRaymarcher()
+
+        # allow gradients for a differentiability check
+        rays_densities.requires_grad = True
+
+        # render opacities, check with gt and check that returned features are None
+        opacities_render = raymarcher_ao(rays_densities)[..., 0]
+        self.assertClose(opacities_render, opacities_gt)
+
+        # check differentiability
+        loss = opacities_render.mean()
+        loss.backward()
+        self.assertTrue(rays_densities.grad.data.isfinite().all())