Remove density field

mattport/nerf/fields/base.py

@@ -3,21 +3,17 @@ Base class for the graphs.
 """
 
 from abc import abstractmethod
+from typing import Tuple
 
 from torch import nn
+from torchtyping import TensorType
 
-from mattport.structures.rays import PointSamples, RaySamples
+from mattport.structures.rays import PointSamples
 
 
 class Field(nn.Module):
     """Base class for fields."""
 
-    def __init__(self) -> None:
-        super().__init__()
-        # minimum weight for it to be valid. used to skip querying the field
-        # TODO(ethan): avoid the hardcoding
-        self.weights_threshold = 1e-4
-
     def density_fn(self, positions):
         """Returns only the density. Used primarily with the occupancy grid."""
         point_samples = PointSamples(positions=positions)
@@ -25,8 +21,15 @@ class Field(nn.Module):
         return density
 
     @abstractmethod
-    def get_density(self, point_samples: PointSamples):
-        """Computes and returns the densities."""
+    def get_density(self, point_samples: PointSamples) -> Tuple[TensorType[..., 1], TensorType[..., "num_features"]]:
+        """Computes and returns the densities.
+
+        Args:
+            point_samples (PointSamples): Samples locations to compute density.
+
+        Returns:
+            Tuple[TensorType[...,1], TensorType[...,"num_features"]]: A tensor of densities and a tensor of features.
+        """
 
     @abstractmethod
     def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None):
@@ -38,14 +41,3 @@ class Field(nn.Module):
         field_outputs = self.get_outputs(point_samples, density_embedding=density_embedding)
         field_outputs["density"] = density
         return field_outputs
-
-    def forward_with_weight_pruning(self, ray_samples: RaySamples):
-        """Forward but with pruning based on weights, from the densities."""
-        density, density_embedding = self.get_density(ray_samples)
-        weights = ray_samples.get_weights(density)
-        valid_mask_outputs = weights >= self.weights_threshold
-        field_outputs = self.get_outputs(
-            ray_samples, density_embedding=density_embedding, valid_mask=valid_mask_outputs
-        )
-        field_outputs["density"] = density
-        return field_outputs

mattport/nerf/fields/density_field.py

@@ -1,40 +0,0 @@
-"""
-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/vanilla_nerf.py

@@ -15,7 +15,6 @@ 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
@@ -43,9 +42,6 @@ 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(
@@ -95,7 +91,9 @@ class NeRFField(Field):
         if not valid_mask.any():  # empty mask
             return density, base_mlp_out
 
-        density[valid_mask], base_mlp_out[valid_mask] = self.density_field(positions[valid_mask])
+        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])
         return density, base_mlp_out
 
     def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None):

mattport/structures/rays.py

@@ -15,22 +15,22 @@ from mattport.utils.misc import is_not_none
 class PointSamples:
     """Samples in space."""
 
-    positions: TensorType[..., 3] = None
-    directions: TensorType[..., 3] = None
-    camera_indices: TensorType[..., 1] = None
-    valid_mask: TensorType[..., 1] = None
+    positions: TensorType[..., 3] = None  # XYZ locations
+    directions: TensorType[..., 3] = None  # Unit direction vector
+    camera_indices: TensorType[..., 1] = None  # Camera index
+    valid_mask: TensorType[..., 1] = None  # Rays that are valid
 
 
 @dataclass
 class RaySamples:
     """Samples along a ray"""
 
-    positions: TensorType[..., 3] = None
-    directions: TensorType[..., 3] = None
-    camera_indices: TensorType[..., 1] = None
-    valid_mask: TensorType[..., 1] = None
-    ts: TensorType[..., 1] = None
-    deltas: TensorType[..., 1] = None
+    positions: TensorType[..., 3] = None  # XYZ locations
+    directions: TensorType[..., 3] = None  # Unit direction vector
+    camera_indices: TensorType[..., 1] = None  # Camera index
+    valid_mask: TensorType[..., 1] = None  # Rays that are valid
+    ts: TensorType[..., 1] = None  # "time steps", distances along ray
+    deltas: TensorType[..., 1] = None  # "width" of each sample
 
     def to_point_samples(self) -> PointSamples:
         """Convert to PointSamples instance and return."""
@@ -67,14 +67,10 @@ class RaySamples:
 
 @dataclass
 class RayBundle:
-    """_summary_
+    """A bundle of ray parameters."""
 
-    Returns:
-        _type_: _description_
-    """
-
-    origins: TensorType["num_rays", 3]
-    directions: TensorType["num_rays", 3]
+    origins: TensorType["num_rays", 3]  # Ray origins
+    directions: TensorType["num_rays", 3]  #
     camera_indices: Optional[TensorType["num_rays"]] = None
     nears: Optional[TensorType["num_rays"]] = None
     fars: Optional[TensorType["num_rays"]] = None