Update to nerfacc 0.0.8 (#471)

nerfactory/models/instant_ngp.py

@@ -62,9 +62,6 @@ class InstantNGPModelConfig(cfg.ModelConfig):
     """Instant NGP doesn't use a collider."""
     field_implementation: Literal["torch", "tcnn"] = "tcnn"  # torch, tcnn, ...
     """one of "torch" or "tcnn", or other fields in 'field_implementation_to_class'"""
-    # TODO(ethan): remove the density field specified here
-    enable_density_field: bool = False
-    """Whether to create a density field to filter samples."""
     num_samples: int = 1024
     """Number of samples in field evaluation."""
     grid_resolution: int = 128
@@ -129,7 +126,7 @@ class NGPModel(Model):
             TrainingCallback(
                 where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],
                 update_every_num_iters=1,
-                func=self.sampler.occ_field.every_n_step,  # will take in step
+                func=self.sampler.occ_field.every_n_step,
             ),
         ]
 
@@ -142,11 +139,10 @@ class NGPModel(Model):
 
     def get_outputs(self, ray_bundle: RayBundle):
         assert self.field is not None
-
-        n_rays = ray_bundle.origins.shape[0]
+        num_rays = len(ray_bundle)
 
         with torch.no_grad():
-            ray_samples, packed_info = self.sampler(
+            ray_samples, packed_info, ray_indices = self.sampler(
                 ray_bundle=ray_bundle,
                 near_plane=self.config.near_plane,
             )
@@ -154,7 +150,7 @@ class NGPModel(Model):
         field_outputs = self.field(ray_samples)
 
         # accumulation
-        weights, ray_indices = nerfacc.volumetric_rendering_weights(
+        weights = nerfacc.volumetric_rendering_weights(
             packed_info=packed_info,
             sigmas=field_outputs[FieldHeadNames.DENSITY],
             frustum_starts=ray_samples.frustums.starts,
@@ -165,10 +161,12 @@ class NGPModel(Model):
             rgb=field_outputs[FieldHeadNames.RGB],
             weights=weights,
             ray_indices=ray_indices,
-            num_rays=n_rays,
+            num_rays=num_rays,
         )
-        depth = self.renderer_depth(weights=weights, ray_samples=ray_samples, ray_indices=ray_indices, num_rays=n_rays)
-        accumulation = self.renderer_accumulation(weights=weights, ray_indices=ray_indices, num_rays=n_rays)
+        depth = self.renderer_depth(
+            weights=weights, ray_samples=ray_samples, ray_indices=ray_indices, num_rays=num_rays
+        )
+        accumulation = self.renderer_accumulation(weights=weights, ray_indices=ray_indices, num_rays=num_rays)
         alive_ray_mask = accumulation.squeeze(-1) > 0
 
         outputs = {

nerfactory/models/modules/ray_sampler.py

@@ -502,7 +502,7 @@ class VolumetricSampler(Sampler):
         num_samples: Optional[int] = None,
         near_plane: float = 0.0,
         remove_occluded_samples: bool = True,
-    ) -> Tuple[RaySamples, TensorType["total_samples", 3]]:
+    ) -> Tuple[RaySamples, TensorType["total_samples", 3], TensorType["total_samples", 2]]:
         """Generate ray samples in a bounding box.
 
         Args:
@@ -512,9 +512,10 @@ class VolumetricSampler(Sampler):
             remove_occluded_samples: Whether to remove occluded samples. (requires evaluating the density function)
 
         Returns:
-            a tuple of (ray_samples, packed_info)
+            a tuple of (ray_samples, packed_info, ray_indices)
             The ray_samples are packed, only storing the valid samples.
             The packed_info contains all the information to recover packed samples into unpacked mode for rendering.
+            The ray_indices contains the indices of the rays that each sample belongs to.
         """
 
         if self.density_field is not None:
@@ -528,8 +529,12 @@ class VolumetricSampler(Sampler):
 
         rays_o = ray_bundle.origins.contiguous()
         rays_d = ray_bundle.directions.contiguous()
+        if ray_bundle.camera_indices is not None:
+            camera_indices = ray_bundle.camera_indices.contiguous()
+        else:
+            camera_indices = None
 
-        packed_info, origins, dirs, starts, ends = nerfacc.volumetric_marching(
+        packed_info, starts, ends = nerfacc.volumetric_marching(
             rays_o=rays_o,
             rays_d=rays_d,
             aabb=self.aabb,
@@ -540,18 +545,25 @@ class VolumetricSampler(Sampler):
             stratified=self.training,
         )
 
+        ray_indices = nerfacc.unpack_to_ray_indices(packed_info).long()
+        origins = rays_o[ray_indices]
+        dirs = rays_d[ray_indices]
+        if camera_indices is not None:
+            camera_indices = camera_indices[ray_indices]
+
         if remove_occluded_samples:
             with torch.no_grad():
                 positions = origins + dirs * (starts + ends) / 2.0
                 densities = self.density_fn(positions)
-                packed_info, starts, ends, origins, dirs = nerfacc.volumetric_rendering_steps(
-                    packed_info,
-                    densities,
-                    starts,
-                    ends,
-                    origins,
-                    dirs,
-                )
+                if camera_indices is not None:
+                    packed_info, starts, ends, origins, dirs, camera_indices = nerfacc.volumetric_rendering_steps(
+                        packed_info, densities, starts, ends, origins, dirs, camera_indices
+                    )
+                else:
+                    packed_info, starts, ends, origins, dirs = nerfacc.volumetric_rendering_steps(
+                        packed_info, densities, starts, ends, origins, dirs
+                    )
+            ray_indices = nerfacc.unpack_to_ray_indices(packed_info).long()
 
         zeros = torch.zeros_like(origins[:, :1])
         ray_samples = RaySamples(
@@ -562,5 +574,6 @@ class VolumetricSampler(Sampler):
                 ends=ends,
                 pixel_area=zeros,
             ),
+            camera_indices=camera_indices,
         )
-        return ray_samples, packed_info
+        return ray_samples, packed_info, ray_indices

pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
     "mediapy==1.1.0",
     "msgpack==1.0.4",
     "msgpack_numpy==0.4.8",
-    "nerfacc==0.0.6",
+    "nerfacc==0.0.8",
     "numpy-quaternion",
     "opencv-python==4.6.0.66",
     "plotly==5.7.0",