Merge pull request #81 from ethanweber/tancik/refactor_occupancy

Refactor occupancy. Work in progress, still a bit gross.

TODOs:
+ Remove Field.py and refactor VanillaNeRF
+ Move masking to field_module
+ Decorator or some other fancy python thing for occupancy
+ Clean up documentation

mattport/nerf/fields/density_field.py

@@ -0,0 +1,40 @@
+"""
+Density Field.
+"""
+
+from torch import nn
+
+from torchtyping import TensorType
+from mattport.nerf.field_modules.field_heads import DensityFieldHead
+from mattport.structures.rays import RaySamples
+
+
+class DensityField(nn.Module):
+    """Base class for fields."""
+
+    def __init__(self, model: nn.Sequential, density_activation: nn.Module = nn.Softplus()) -> None:
+        super().__init__()
+        self.model = model
+        self.density_head = DensityFieldHead(in_dim=self.model[-1].get_out_dim(), activation=density_activation)
+
+    def process_point(self, in_tensor: TensorType):
+        """_summary_"""
+        features = self.model(in_tensor)
+        density = self.density_head(features)
+        return density, features
+
+    def proccess_ray(self, rays: RaySamples):
+        """_summary_"""
+        densities, features = self.process_point(in_tensor=rays.positions)
+
+        weights = rays.get_weights(densities)
+        return densities, weights, features
+
+    def forward(self, in_tensor):
+        """summary"""
+        return self.process_point(in_tensor)
+
+    def density_fn(self, positions):
+        """Returns only the density."""
+        density, _ = self.process_point(in_tensor=positions)
+        return density

mattport/nerf/graph/base.py

@@ -1,7 +1,6 @@
 """
 The Graph module contains all trainable parameters.
 """
-import logging
 from abc import abstractmethod
 from collections import defaultdict
 from typing import Dict, List, Union
@@ -13,7 +12,6 @@ from torch.nn import Parameter
 from torchtyping import TensorType
 
 from mattport.nerf.dataset.structs import SceneBounds
-from mattport.nerf.occupancy_grid import OccupancyGrid
 from mattport.nerf.ray_generator import RayGenerator
 from mattport.nerf.colliders import AABBBoxCollider
 from mattport.structures.cameras import get_camera_model
@@ -65,7 +63,6 @@ class Graph(AbstractGraph):
         self.steps_per_occupancy_grid_update = steps_per_occupancy_grid_update
         self.kwargs = kwargs
         self.collider = None
-        self.occupancy_grid = OccupancyGrid(aabb=self.scene_bounds.aabb)
         self.ray_generator = RayGenerator(self.intrinsics, self.camera_to_world)
         self.populate_collider()
         self.populate_fields()
@@ -97,10 +94,6 @@ class Graph(AbstractGraph):
         batch: Union[str, Dict[str, torch.tensor]] = None,
         step: int = None,
     ):
-        # update the occupancy grid
-        if step and step % self.steps_per_occupancy_grid_update == 0:
-            logging.info("Updating occupancy grid.")
-            self.occupancy_grid.update_occupancy_grid(density_fn=self.field_fine.density_fn)
 
         # get the rays
         original_ray_bundle = self.ray_generator.forward(ray_indices)  # RayBundle

mattport/nerf/graph/vanilla_nerf.py

@@ -15,6 +15,7 @@ from mattport.nerf.field_modules.encoding import NeRFEncoding
 from mattport.nerf.field_modules.field_heads import DensityFieldHead, RGBFieldHead
 from mattport.nerf.field_modules.mlp import MLP
 from mattport.nerf.fields.base import Field
+from mattport.nerf.fields.density_field import DensityField
 from mattport.nerf.graph.base import Graph
 from mattport.nerf.loss import MSELoss
 from mattport.nerf.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer
@@ -42,6 +43,9 @@ class NeRFField(Field):
         self.build_mlp_rgb()
         self.build_heads()
 
+        sequence = nn.Sequential(self.encoding_xyz, self.mlp_base)
+        self.density_field = DensityField(sequence)
+
     def build_encodings(self):
         """Build the encodings."""
         self.encoding_xyz = NeRFEncoding(
@@ -89,9 +93,8 @@ class NeRFField(Field):
         )
         if not valid_mask.any():  # empty mask
             return density, base_mlp_out
-        encoded_xyz = self.encoding_xyz(positions[valid_mask])
-        base_mlp_out[valid_mask] = self.mlp_base(encoded_xyz)
-        density[valid_mask] = self.field_output_density(base_mlp_out[valid_mask])
+
+        density[valid_mask], base_mlp_out[valid_mask] = self.density_field(positions[valid_mask])
         return density, base_mlp_out
 
     def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None):
@@ -138,7 +141,6 @@ class NeRFGraph(Graph):
         self.field_fine = NeRFField()
 
     def populate_misc_modules(self):
-
         # samplers
         self.sampler_uniform = UniformSampler(num_samples=self.num_coarse_samples)
         self.sampler_pdf = PDFSampler(num_samples=self.num_importance_samples)
@@ -162,7 +164,6 @@ class NeRFGraph(Graph):
         return param_groups
 
     def get_outputs(self, ray_bundle: RayBundle):
-
         # uniform sampling
         ray_samples_uniform = self.sampler_uniform(ray_bundle)
 
@@ -189,12 +190,6 @@ class NeRFGraph(Graph):
         accumulation_fine = self.renderer_accumulation(weights_fine)
         depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf.ts)
 
-        # --- occupancy grid ---
-        # TODO(ethan): move this elsewhere. currently here for debugging purposes
-        density_occ = self.occupancy_grid.get_densities(ray_samples_pdf.positions)
-        weights_occ = ray_samples_pdf.get_weights(density_occ)
-        depth_occ = self.renderer_depth(weights_occ, ray_samples_pdf.ts)
-
         outputs = {
             "rgb_coarse": rgb_coarse,
             "rgb_fine": rgb_fine,
@@ -202,7 +197,6 @@ class NeRFGraph(Graph):
             "accumulation_fine": accumulation_fine,
             "depth_coarse": depth_coarse,
             "depth_fine": depth_fine,
-            "depth_occ": depth_occ,
         }
         return outputs
 
@@ -241,9 +235,6 @@ class NeRFGraph(Graph):
         writer.write_image(name=f"image_idx_{image_idx}", image=combined_acc, step=step, group="accumulation")
         writer.write_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth")
 
-        depth_occ = visualization.apply_depth_colormap(outputs["depth_occ"])
-        writer.write_image(name=f"image_idx_{image_idx}", image=depth_occ, step=step, group="depth_occupancy_grid")
-
         # Switch images from [H, W, C] to [1, C, H, W] for metrics computations
         image = torch.moveaxis(image, -1, 0)[None, ...]
         rgb_coarse = torch.moveaxis(rgb_coarse, -1, 0)[None, ...]

mattport/nerf/occupancy_grid.py

@@ -13,7 +13,9 @@ from mattport.utils.misc import is_not_none
 class OccupancyGrid(nn.Module):
     """Module to keep track of the density and occupancy."""
 
-    def __init__(self, num_cascades=1, resolution=128, aabb=None) -> None:
+    def __init__(
+        self, num_cascades=1, resolution=128, aabb=None, density_fn=None, update_every_num_iters: int = 16
+    ) -> None:
         super().__init__()
         assert is_not_none(aabb), "The axis-aligned bounding box aabb is not defined!"
         self.num_cascades = num_cascades  # the number of levels (i.e, cascades)
@@ -26,6 +28,10 @@ class OccupancyGrid(nn.Module):
         self.register_buffer("occupancy_grid", occupancy_grid)
         self.mean_density = 0.0
 
+        self.density_fn = density_fn
+        self.update_every_num_iters = update_every_num_iters
+        self.iteration_count = 0
+
     def reset(self):
         """Zero out the occupancy grid."""
         self.occupancy_grid.zero_()
@@ -96,8 +102,18 @@ class OccupancyGrid(nn.Module):
         )
         self.mean_density = torch.mean(self.occupancy_grid[valid_mask]).item()
 
-    def get_densities(self, xyzs: TensorType[..., 3]):
-        """Trilinear interpolation to get the density values."""
+    @torch.no_grad()
+    def get_densities(self, xyzs: TensorType[..., 3], update_iter_count: bool = False) -> TensorType[..., 1]:
+        """Trilinear interpolation to get the density values.
+
+        Args:
+            xyzs (TensorType[..., 3]): 3D querry coordinate
+            update_iter_count (bool, optional): Update current iteration. Used for knowing when to update grid.
+                Defaults to False.
+
+        Returns:
+            TensorType[..., 1]: Density values
+        """
         occupancy_grid = self.occupancy_grid[None, ...]  # shape (1, num_cascades, X_res, Y_res, Z_res)
         xyzs_shape = xyzs.shape
         xyzs_reshaped = xyzs.view(1, -1, 1, 1, 3)
@@ -112,6 +128,14 @@ class OccupancyGrid(nn.Module):
             padding_mode="zeros",
         )
         densities = densities.view(*xyzs_shape[:-1], 1)
+
+        # TODO we should be tracking the iteration another way
+        if self.training and update_iter_count:
+            self.iteration_count += 1
+            if self.iteration_count % self.update_every_num_iters == 0:
+                if self.density_fn is not None:
+                    self.update_occupancy_grid(self.density_fn)
+
         return densities
 
     def forward(self, x):

mattport/nerf/sampler.py

@@ -6,10 +6,75 @@ from typing import Optional
 import torch
 from torch import nn
 from torchtyping import TensorType
+from mattport.nerf.occupancy_grid import OccupancyGrid
 
 from mattport.structures.rays import RayBundle, RaySamples
 
 
+class OccupancyGridSampler(nn.Module):
+    """Sample based on occupancy grid"""
+
+    def __init__(
+        self,
+        occupancy_grid: OccupancyGrid,
+        num_samples: int = 256,
+        weight_threshold: float = 1e-4,
+        train_stratified=True,
+    ) -> None:
+        """
+        Args:
+        """
+        super().__init__()
+        self.num_samples = num_samples
+        self.train_stratified = train_stratified
+        self.weight_threshold = weight_threshold
+        self.occupancy_grid = occupancy_grid
+
+    @torch.no_grad()
+    def forward(
+        self,
+        ray_bundle: RayBundle,
+        num_samples: Optional[int] = None,
+    ) -> RaySamples:
+        """Generates position samples uniformly.
+
+        Args:
+            ray_bundle (RayBundle): Rays to generate samples for
+            num_samples (Optional[int]): Number of samples per ray
+
+        Returns:
+            RaySamples: Positions and deltas for samples along a ray
+        """
+        assert ray_bundle.nears is not None
+        assert ray_bundle.fars is not None
+
+        num_samples = num_samples or self.num_samples
+        num_rays = ray_bundle.origins.shape[0]
+
+        bins = torch.linspace(0.0, 1.0, num_samples + 1).to(ray_bundle.origins.device)  # shape (num_samples+1,)
+        bins = ray_bundle.nears[:, None] + bins[None, :] * (
+            ray_bundle.fars[:, None] - ray_bundle.nears[:, None]
+        )  # shape (num_rays, num_samples+1)
+
+        if self.train_stratified and self.training:
+            t_rand = torch.rand((num_rays, num_samples), dtype=bins.dtype, device=bins.device)
+            ts = bins[:, :-1] + t_rand * (bins[:, 1:] - bins[:, :-1])  # shape (num_rays, num_samples)
+        else:
+            ts = (bins[:, 1:] + bins[:, :-1]) / 2
+
+        ray_samples = RaySamples(
+            ts=ts,
+            ray_bundle=ray_bundle,
+        )
+
+        densities = self.occupancy_grid.get_densities(ray_samples.positions, update_iter_count=True)
+        weights = ray_samples.get_weights(densities)
+
+        valid_mask = weights >= self.weight_threshold
+
+        return ray_samples, weights, valid_mask
+
+
 class UniformSampler(nn.Module):
     """Sample uniformly along a ray"""