Switch to dictionary outputs

mattport/nerf/field/nerf.py

@@ -6,7 +6,7 @@ import torch
 from torch import nn
 
 from mattport.nerf.field_modules.encoding import NeRFEncoding
-from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHeadOutputs, RGBFieldHead
+from mattport.nerf.field_modules.field_heads import DensityFieldHead, RGBFieldHead
 from mattport.nerf.field_modules.mlp import MLP
 from mattport.structures.rays import RaySamples
 
@@ -52,6 +52,5 @@ class NeRFField(nn.Module):
         field_rgb_output = self.field_output_rgb(rgb_mlp_out)
         field_density_out = self.field_output_density(base_mlp_out)
 
-        # TODO(ethan): implement a way to combine the field outputs...
-        field_outputs = FieldHeadOutputs(rgb=field_rgb_output.rgb, density=field_density_out.density)
+        field_outputs = field_rgb_output | field_density_out
         return field_outputs

mattport/nerf/field_modules/field_heads.py

@@ -1,8 +1,8 @@
 """
 Collection of render heads
 """
-from dataclasses import dataclass
-from typing import Optional
+from enum import Enum
+from typing import Dict, Optional
 
 from torch import nn
 from torchtyping import TensorType
@@ -10,31 +10,32 @@ from torchtyping import TensorType
 from mattport.nerf.field_modules.base import FieldModule
 
 
-@dataclass
-class FieldHeadOutputs:
-    """_summary_"""
+class FieldHeadNames(Enum):
+    """Possible field outputs"""
 
-    rgb: TensorType["...", 3] = None
-    density: TensorType["...", 1] = None
+    RGB = "rgb"
+    DENSITY = "density"
 
 
 class FieldHead(FieldModule):
     """Base field output"""
 
     def __init__(
-        self, in_dim: int, out_dim: int, field_quantity_name: str, activation: Optional[nn.Module] = None
+        self, in_dim: int, out_dim: int, field_head_name: FieldHeadNames, activation: Optional[nn.Module] = None
     ) -> None:
         """
         Args:
             in_dim (int): input dimension
             out_dim (int): output dimension for renderer
-            field_quantity_name (srt): name of field head
+            field_head_name (FieldHeadNames): name of field head
             activation (Optional[nn.Module]): output head activation
         """
         super().__init__()
         self.in_dim = in_dim
         self.out_dim = out_dim
-        self.field_quantity_name = field_quantity_name
+        # if field_head_name not in FieldHeadNames:
+        #     raise ValueError("Incorrect field output name. Should be one of FieldHeadNames.")
+        self.field_head_name = field_head_name
         self.activation = activation
         self.net = None
         self.build_nn_modules()
@@ -45,7 +46,7 @@ class FieldHead(FieldModule):
             layers.append(self.activation)
         self.net = nn.Sequential(*layers)
 
-    def forward(self, in_tensor: TensorType[..., "in_dim"]) -> FieldHeadOutputs:
+    def forward(self, in_tensor: TensorType[..., "in_dim"]) -> Dict[FieldHeadNames, TensorType]:
         """Process network output for renderer
 
         Args:
@@ -54,25 +55,23 @@ class FieldHead(FieldModule):
         Returns:
             TensorType[..., "out_dim"]: Render head output
         """
-        if not self.field_quantity_name:
+        if not self.field_head_name:
             raise ValueError("field_quantity_name should be set in the child class. E.g., as 'rgb' or 'density'.")
         if not self.net:
             raise SystemError("Render head network not initialized. build_nn_modules() should be called.")
         out_tensor = self.net(in_tensor)
-        field_head_outputs = FieldHeadOutputs()
-        setattr(field_head_outputs, self.field_quantity_name, out_tensor)
-        return field_head_outputs
+        return {self.field_head_name: out_tensor}
 
 
 class DensityFieldHead(FieldHead):
     """Density output"""
 
     def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Softplus()) -> None:
-        super().__init__(in_dim, out_dim=1, field_quantity_name="density", activation=activation)
+        super().__init__(in_dim, out_dim=1, field_head_name=FieldHeadNames.DENSITY, activation=activation)
 
 
 class RGBFieldHead(FieldHead):
     """RGB output"""
 
     def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Sigmoid()) -> None:
-        super().__init__(in_dim, out_dim=3, field_quantity_name="rgb", activation=activation)
+        super().__init__(in_dim, out_dim=3, field_head_name=FieldHeadNames.RGB, activation=activation)

mattport/nerf/graph/nerf.py

@@ -49,13 +49,13 @@ class NeRFGraph(Graph):
         uniform_ray_samples = self.sampler_uniform(ray_bundle)  # RaySamples
         coarse_field_outputs = self.field_coarse(uniform_ray_samples)  # FieldOutputs
         coarse_renderer_outputs = self.renderer_rgb(
-            rgb=coarse_field_outputs.rgb, density=coarse_field_outputs.density, deltas=uniform_ray_samples.deltas
+            field_outputs=coarse_field_outputs.rgb, deltas=uniform_ray_samples.deltas
         )  # RendererOutputs
         # fine network:
-        pdf_ray_samples = self.sampler_pdf(ray_bundle, uniform_ray_samples, coarse_field_outputs.density)  # RaySamples
+        pdf_ray_samples = self.sampler_pdf(ray_bundle, uniform_ray_samples, coarse_field_outputs)  # RaySamples
         fine_field_outputs = self.field_fine(pdf_ray_samples)  # FieldOutputs
         fine_renderer_outputs = self.renderer_rgb(
-            rgb=fine_field_outputs.rgb, density=fine_field_outputs.density, deltas=pdf_ray_samples.deltas
+            field_outputs=fine_field_outputs, deltas=pdf_ray_samples.deltas
         )  # RendererOutputs
         # outputs:
         outputs = {"rgb_coarse": coarse_renderer_outputs.rgb, "rgb_fine": fine_renderer_outputs.rgb}

mattport/nerf/sampler.py

@@ -2,10 +2,11 @@
 Collection of sampling strategies
 """
 
-from typing import Optional
+from typing import Dict, Optional
 import torch
 from torch import nn
 from torchtyping import TensorType
+from mattport.nerf.field_modules.field_heads import FieldHeadNames
 
 from mattport.structures.rays import RayBundle, RaySamples
 
@@ -78,7 +79,7 @@ class PDFSampler(nn.Module):
         self,
         ray_bundle: RayBundle,
         coarse_ray_samples: RaySamples,
-        density: TensorType[..., "num_samples", 1],
+        field_outputs: Dict[FieldHeadNames, TensorType],
         num_samples: Optional[int] = None,
         randomized: bool = True,
         eps: float = 1e-5,
@@ -102,6 +103,7 @@ class PDFSampler(nn.Module):
 
         # Calculate weight contributions along ray
         # Todo(matt): This computation is duplicated
+        density = field_outputs[FieldHeadNames.DENSITY]
         delta_density = coarse_ray_samples.deltas * density[..., 0]
         alphas = 1 - torch.exp(-delta_density)
 

tests/nerf/field_modules/test_field_outputs.py

@@ -4,17 +4,16 @@ Field output tests
 import torch
 from torch import nn
 
-from mattport.nerf.field_modules import field_heads
+from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHead, FieldHeadNames, RGBFieldHead
 
 
 def test_field_output():
     """Test render output"""
     in_dim = 6
     out_dim = 4
+    field_head_name = FieldHeadNames.DENSITY
     activation = nn.ReLU()
-    render_head = field_heads.FieldHead(
-        in_dim=in_dim, out_dim=out_dim, field_quantity_name="rgb", activation=activation
-    )
+    render_head = FieldHead(in_dim=in_dim, out_dim=out_dim, field_head_name=field_head_name, activation=activation)
     assert render_head.get_out_dim() == out_dim
 
     x = torch.ones((9, in_dim))
@@ -24,7 +23,7 @@ def test_field_output():
 def test_density_output():
     """Test rgb output"""
     in_dim = 6
-    density_head = field_heads.DensityFieldHead(in_dim)
+    density_head = DensityFieldHead(in_dim)
     assert density_head.get_out_dim() == 1
 
     x = torch.ones((9, in_dim))
@@ -34,8 +33,14 @@ def test_density_output():
 def test_rgb_output():
     """Test rgb output"""
     in_dim = 6
-    rgb_head = field_heads.RGBFieldHead(in_dim)
+    rgb_head = RGBFieldHead(in_dim)
     assert rgb_head.get_out_dim() == 3
 
     x = torch.ones((9, in_dim))
     rgb_head(x)
+
+
+if __name__ == "__main__":
+    test_field_output()
+    test_density_output()
+    test_rgb_output()

tests/nerf/test_sampler.py

@@ -2,6 +2,7 @@
 Test samplers
 """
 import torch
+from mattport.nerf.field_modules.field_heads import FieldHeadNames
 from mattport.nerf.sampler import PDFSampler, UniformSampler
 
 from mattport.structures.rays import RayBundle
@@ -35,14 +36,14 @@ def test_pdf_sampler():
     directions = torch.ones_like(origins)
     ray_bundle = RayBundle(origins=origins, directions=directions)
 
-    density = torch.ones((10, num_samples, 1))
+    field_outputs = {FieldHeadNames.DENSITY: torch.ones((10, num_samples, 1))}
 
     uniform_sampler = UniformSampler(near_plane=near_plane, far_plane=far_plane, num_samples=num_samples)
     coarse_ray_samples = uniform_sampler(ray_bundle)
 
     # Just check that it doesn't crash. Need to add some actual tests.
     pdf_sampler = PDFSampler(num_samples)
-    pdf_sampler(ray_bundle, coarse_ray_samples, density, num_samples)
+    pdf_sampler(ray_bundle, coarse_ray_samples, field_outputs, num_samples)
 
     # TODO Tancik: Add more precise tests