Cleanup nerfacc integration (#463)

* cleanup nerfacc implementation

* typing

nerfactory/fields/instant_ngp_field.py

@@ -151,26 +151,6 @@ class TCNNInstantNGPField(Field):
         rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions)
         return {FieldHeadNames.RGB: rgb}
 
-    def get_outputs_from_positions_and_direction(self, positions: TensorType["bs", 3], directions: TensorType["bs", 3]):
-        """Given positions and directions, return the outputs of the field.
-
-        Args:
-            positions: positions of the points
-            directions: directions at the point locations
-        """
-        # assume flat positions and directions
-        # TODO: normalize positions
-        positions = SceneBounds.get_normalized_positions(positions, self.aabb)
-        directions = get_normalized_directions(directions)
-        h = self.mlp_base(positions)
-        density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1)
-        density = trunc_exp(density_before_activation.to(positions))
-        # TODO: normalize directions
-        d = self.direction_encoding(directions)
-        h = torch.cat([d, base_mlp_out], dim=-1)
-        rgb = self.mlp_head(h).to(directions)
-        return {FieldHeadNames.RGB: rgb, FieldHeadNames.DENSITY: density}
-
 
 class TorchInstantNGPField(NeRFField):
     """

nerfactory/models/instant_ngp.py

@@ -27,7 +27,6 @@ from torch.nn import Parameter
 from torchmetrics import PeakSignalNoiseRatio
 from torchmetrics.functional import structural_similarity_index_measure
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
-from torchtyping import TensorType
 
 from nerfactory.cameras.rays import RayBundle
 from nerfactory.configs import base as cfg
@@ -135,15 +134,6 @@ class NGPModel(Model):
     def get_outputs(self, ray_bundle: RayBundle):
         assert self.field is not None
 
-        def query_fn(positions: TensorType["bs", 3], directions: TensorType["bs", 3], only_density=False):
-            assert self.field is not None
-            if only_density:
-                return self.field.density_fn(positions)
-            field_outputs = self.field.get_outputs_from_positions_and_direction(positions, directions)
-            rgbs = field_outputs[FieldHeadNames.RGB]
-            sigmas = field_outputs[FieldHeadNames.DENSITY]
-            return rgbs, sigmas
-
         n_rays = ray_bundle.origins.shape[0]
 
         with torch.no_grad():
@@ -152,48 +142,23 @@ class NGPModel(Model):
                 near_plane=self.config.near_plane,
             )
 
-            frustum_starts = ray_samples.frustums.starts
-            frustum_ends = ray_samples.frustums.ends
-            frustum_dirs = ray_samples.frustums.directions
-            frustum_positions = ray_samples.frustums.get_positions()
-
-        # compat the samples thru volumetric rendering
-        with torch.no_grad():
-            densities = self.field.density_fn(frustum_positions)
-            (
-                compact_packed_info,
-                compact_frustum_starts,
-                compact_frustum_ends,
-                compact_frustum_positions,
-                compact_frustum_dirs,
-            ) = nerfacc.volumetric_rendering_steps(
-                packed_info,
-                densities,
-                frustum_starts,
-                frustum_ends,
-                frustum_positions,
-                frustum_dirs,
-            )
-
-        # network
-        compact_query_results = query_fn(compact_frustum_positions, compact_frustum_dirs)
-        compact_rgbs, compact_densities = compact_query_results[0], compact_query_results[1]
+        field_outputs = self.field(ray_samples)
+        rgbs = field_outputs[FieldHeadNames.RGB]
+        densities = field_outputs[FieldHeadNames.DENSITY]
 
         # accumulation
         compact_weights, compact_ray_indices = nerfacc.volumetric_rendering_weights(
-            compact_packed_info,
-            compact_densities,
-            compact_frustum_starts,
-            compact_frustum_ends,
-        )
-        accumulated_color = nerfacc.volumetric_rendering_accumulate(
-            compact_weights, compact_ray_indices, compact_rgbs, n_rays
+            packed_info,
+            densities,
+            ray_samples.frustums.starts,
+            ray_samples.frustums.ends,
         )
+        accumulated_color = nerfacc.volumetric_rendering_accumulate(compact_weights, compact_ray_indices, rgbs, n_rays)
         accumulated_weight = nerfacc.volumetric_rendering_accumulate(compact_weights, compact_ray_indices, None, n_rays)
         accumulated_depth = nerfacc.volumetric_rendering_accumulate(
             compact_weights,
             compact_ray_indices,
-            (compact_frustum_starts + compact_frustum_ends) / 2.0,
+            (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2.0,
             n_rays,
         )
 

nerfactory/models/modules/ray_sampler.py

@@ -469,6 +469,7 @@ class VolumetricSampler(Sampler):
         """
         super().__init__(num_samples=num_samples)
         self.aabb = aabb
+        self.density_fn = density_fn
 
         self.render_step_size = ((self.aabb[3:] - self.aabb[:3]).max() * math.sqrt(3) / num_samples).item()
 
@@ -500,6 +501,7 @@ class VolumetricSampler(Sampler):
         ray_bundle: RayBundle,
         num_samples: Optional[int] = None,
         near_plane: float = 0.0,
+        remove_occluded_samples: bool = True,
     ) -> Tuple[RaySamples, TensorType["total_samples", 3]]:
         """Generate ray samples in a bounding box.
 
@@ -507,6 +509,7 @@ class VolumetricSampler(Sampler):
             ray_bundle: Rays to generate samples for
             num_samples: Number of samples per ray
             near_plane: Near plane for raymarching
+            remove_occluded_samples: Whether to remove occluded samples. (requires evaluating the density function)
 
         Returns:
             a tuple of (ray_samples, packed_info)
@@ -537,8 +540,27 @@ class VolumetricSampler(Sampler):
             stratified=self.training,
         )
 
+        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,
+                )
+
         zeros = torch.zeros_like(origins[:, :1])
         ray_samples = RaySamples(
-            frustums=Frustums(origins=origins, directions=dirs, starts=starts, ends=ends, pixel_area=zeros),
+            frustums=Frustums(
+                origins=origins,
+                directions=dirs,
+                starts=starts,
+                ends=ends,
+                pixel_area=zeros,
+            ),
         )
         return ray_samples, packed_info