From 96e1f2a698d8277e9667c2e15fcbd9e9eabb64ed Mon Sep 17 00:00:00 2001
From: Matthew Tancik <tancik@berkeley.edu>
Date: Mon, 19 Sep 2022 17:38:47 -0700
Subject: [PATCH] Disable occlusion sampling and speed up eval (#488)

disable occlusion sampling and speed up eval
---
 nerfactory/cameras/cameras.py            |  6 +++---
 nerfactory/datamanagers/base.py          |  2 +-
 nerfactory/datamanagers/dataloaders.py   |  1 -
 nerfactory/models/modules/ray_sampler.py | 15 +++++++--------
 4 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/nerfactory/cameras/cameras.py b/nerfactory/cameras/cameras.py
index 66cfc2ad8..217994fb9 100644
--- a/nerfactory/cameras/cameras.py
+++ b/nerfactory/cameras/cameras.py
@@ -232,17 +232,17 @@ class Cameras:
 
         dx = torch.sqrt(torch.sum((directions - directions_stack[1]) ** 2, dim=-1))
         dy = torch.sqrt(torch.sum((directions - directions_stack[2]) ** 2, dim=-1))
-        pixel_area = dx * dy
+        pixel_area = (dx * dy)[..., None]
 
         if not isinstance(camera_indices, torch.Tensor):
-            ray_bundle_camera_indices = torch.Tensor([camera_indices]).broadcast_to((self._num_cameras)).to(self.device)
+            ray_bundle_camera_indices = torch.Tensor([camera_indices]).broadcast_to(pixel_area.shape).to(self.device)
         else:
             ray_bundle_camera_indices = camera_indices
 
         return RayBundle(
             origins=origins,
             directions=directions,
-            pixel_area=pixel_area[..., None],
+            pixel_area=pixel_area,
             camera_indices=ray_bundle_camera_indices,
         )
 
diff --git a/nerfactory/datamanagers/base.py b/nerfactory/datamanagers/base.py
index dfd82a59d..84d7f7bf6 100644
--- a/nerfactory/datamanagers/base.py
+++ b/nerfactory/datamanagers/base.py
@@ -332,6 +332,6 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
     def next_eval_image(self, step: int) -> Tuple[int, RayBundle, Dict]:
         for camera_ray_bundle, batch in self.eval_dataloader:
             assert camera_ray_bundle.camera_indices is not None
-            image_idx = int(camera_ray_bundle.camera_indices[0, 0])
+            image_idx = int(camera_ray_bundle.camera_indices[0, 0, 0])
             return image_idx, camera_ray_bundle, batch
         raise ValueError("No more eval images")
diff --git a/nerfactory/datamanagers/dataloaders.py b/nerfactory/datamanagers/dataloaders.py
index b70e0625f..2520beef3 100644
--- a/nerfactory/datamanagers/dataloaders.py
+++ b/nerfactory/datamanagers/dataloaders.py
@@ -156,7 +156,6 @@ class EvalDataloader(DataLoader):
             image_idx: Camera image index
         """
         ray_bundle = self.cameras.generate_rays(camera_indices=image_idx)
-        ray_bundle.camera_indices = torch.Tensor([image_idx])[..., None].int()
         batch = self.input_dataset[image_idx]
         batch = get_dict_to_torch(batch, device=self.device, exclude=["image"])
         return ray_bundle, batch
diff --git a/nerfactory/models/modules/ray_sampler.py b/nerfactory/models/modules/ray_sampler.py
index 929d06b40..9116d9b8e 100644
--- a/nerfactory/models/modules/ray_sampler.py
+++ b/nerfactory/models/modules/ray_sampler.py
@@ -347,6 +347,12 @@ class VolumetricSampler(Sampler):
     """Sampler inspired by the one proposed in the Instant-NGP paper.
 
     This sampler does ray-box AABB test, ray samples generation and density check, all together.
+
+        Args:
+        aabb: Bounding box of the scene.
+        density_fn: Function that takes a tensor of points and returns a tensor of densities.
+        grid_resolution: Resolution of the density grid.
+        num_samples: Number of samples per ray.
     """
 
     def __init__(
@@ -356,14 +362,7 @@ class VolumetricSampler(Sampler):
         grid_resolution: int = 128,
         num_samples: int = 1024,
     ) -> None:
-        """Init.
 
-        Args:
-            aabb: Bounding box of the scene.
-            density_fn: Function that takes a tensor of points and returns a tensor of densities.
-            grid_resolution: Resolution of the density grid.
-            num_samples: Number of samples per ray.
-        """
         super().__init__(num_samples=num_samples)
         self.aabb = aabb
         self.density_fn = density_fn
@@ -398,7 +397,7 @@ class VolumetricSampler(Sampler):
         ray_bundle: RayBundle,
         num_samples: Optional[int] = None,
         near_plane: float = 0.0,
-        remove_occluded_samples: bool = True,
+        remove_occluded_samples: bool = False,
     ) -> Tuple[RaySamples, TensorType["total_samples", 3], TensorType["total_samples", 2]]:
         """Generate ray samples in a bounding box.