refactor field heads

mattport/nerf/field_modules/base.py

@@ -44,15 +44,11 @@ class FieldModule(nn.Module):
 
     @abstractmethod
     def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
-        """_summary_
-
+        """
         Args:
-            x (TensorType[..., "input_dim"]): Input tensor to be encoded
-
-        Raises:
-            NotImplementedError: _description_
+            in_tensor (TensorType[..., "input_dim"]): Input tensor to process
 
         Returns:
-            TensorType[..., "output_dim"]: A encoded input tensor
+            TensorType[..., "output_dim"]: Processed tensor
         """
         raise NotImplementedError

mattport/nerf/field_modules/field_heads.py

@@ -14,6 +14,7 @@ class FieldHeadNames(Enum):
     """Possible field outputs"""
 
     RGB = "rgb"
+    SH = "sh"
     DENSITY = "density"
     UNCERTAINTY = "uncertainty"
     TRANSIENT_RGB = "transient_rgb"
@@ -38,14 +39,8 @@ class FieldHead(FieldModule):
         super().__init__()
         self.in_dim = in_dim
         self.out_dim = out_dim
-        # 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()
-
-    def build_nn_modules(self) -> None:
         layers = [nn.Linear(self.in_dim, self.out_dim)]
         if self.activation:
             layers.append(self.activation)
@@ -60,10 +55,6 @@ class FieldHead(FieldModule):
         Returns:
             TensorType[..., "out_dim"]: Render head output
         """
-        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)
         return {self.field_head_name: out_tensor}
 
@@ -95,7 +86,7 @@ class SHFieldHead(FieldHead):
             activation (Optional[nn.Module], optional): Output activation. Defaults to None.
         """
         out_dim = channels * levels**2
-        super().__init__(in_dim, out_dim=out_dim, field_head_name=FieldHeadNames.RGB, activation=activation)
+        super().__init__(in_dim, out_dim=out_dim, field_head_name=FieldHeadNames.SH, activation=activation)
 
 
 class UncertaintyFieldHead(FieldHead):
@@ -123,4 +114,9 @@ class SemanticStuffHead(FieldHead):
     """Semantic stuff output"""
 
     def __init__(self, in_dim: int, num_classes: int) -> None:
+        """
+        Args:
+            in_dim (int): Input dimension
+            num_classes (int): Number of semantic classes
+        """
         super().__init__(in_dim, out_dim=num_classes, field_head_name=FieldHeadNames.SEMANTICS_STUFF, activation=None)

mattport/nerf/renderers.py

@@ -35,6 +35,29 @@ class RGBRenderer(nn.Module):
         super().__init__()
         self.background_color = background_color
 
+    @classmethod
+    def combine_rgb(
+        cls,
+        rgb: TensorType[..., "num_samples", 3],
+        weights: TensorType[..., "num_samples"],
+        background_color: Optional[TensorType[3]] = None,
+    ) -> TensorType[..., 3]:
+        """Composite samples along ray and render color image
+
+        Args:
+            rgb (TensorType[..., "num_samples", -1]): RGB for each sample
+            weights (TensorType[..., "num_samples"]): Weights for each sample
+
+        Returns:
+            TensorType[..., 3]: Outputs rgb values.
+        """
+        rgb = torch.sum(weights[..., None] * rgb, dim=-2)
+
+        if background_color is not None:
+            rgb = rgb + background_color.to(weights.device)[None, ...] * (1.0 - torch.sum(weights, dim=-1)[..., None])
+
+        return rgb
+
     def forward(
         self,
         rgb: TensorType[..., "num_samples", 3],
@@ -50,17 +73,9 @@ class RGBRenderer(nn.Module):
             RendererOutputs: Outputs with rgb values.
         """
 
-        rgb = torch.sum(weights[..., None] * rgb, dim=-2)
-
-        if self.background_color is not None:
-            rgb = rgb + self.background_color.to(weights.device)[None, ...] * (
-                1.0 - torch.sum(weights, dim=-1)[..., None]
-            )
-        torch.clamp_(rgb, min=0.0, max=1.0)
-
-        assert torch.max(rgb) <= 1.0
-        assert torch.min(rgb) >= 0.0
-
+        rgb = self.combine_rgb(rgb, weights)
+        if not self.training:
+            torch.clamp_(rgb, min=0.0, max=1.0)
         renderer_outputs = RendererOutputs(rgb=rgb)
         return renderer_outputs
 
@@ -69,7 +84,7 @@ class SHRenderer(nn.Module):
     """Render RGB value from spherical harmonics."""
 
     def __init__(
-        self, activation: Optional[nn.Module] = nn.Sigmoid(), background_color: Optional[TensorType[3]] = None
+        self, background_color: Optional[TensorType[3]] = None, activation: Optional[nn.Module] = nn.Sigmoid()
     ) -> None:
         """
         Args:
@@ -103,21 +118,12 @@ class SHRenderer(nn.Module):
         components = components_from_spherical_harmonics(levels=levels, directions=directions)
 
         rgb = sh * components[..., None, :]  # [..., num_samples, 3, sh_components]
-        rgb = torch.sum(sh, dim=-1)  # [..., num_samples, 3]
-
-        rgb = torch.sum(weights[..., None] * rgb, dim=-2) + 0.5
+        rgb = torch.sum(sh, dim=-1) + 0.5  # [..., num_samples, 3]
 
         if self.activation is not None:
             self.activation(rgb)
 
-        if self.background_color is not None:
-            rgb = rgb + self.background_color.to(weights.device)[None, ...] * (
-                1.0 - torch.sum(weights, dim=-1)[..., None]
-            )
-        torch.clamp_(rgb, min=0.0, max=1.0)
-
-        assert torch.max(rgb) <= 1.0
-        assert torch.min(rgb) >= 0.0
+        rgb = RGBRenderer.combine_rgb(rgb, weights, background_color=self.background_color)
 
         renderer_outputs = RendererOutputs(rgb=rgb)
         return renderer_outputs

tests/nerf/field_modules/test_field_outputs.py

@@ -4,7 +4,13 @@ Field output tests
 import torch
 from torch import nn
 
-from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHead, FieldHeadNames, RGBFieldHead
+from mattport.nerf.field_modules.field_heads import (
+    DensityFieldHead,
+    FieldHead,
+    FieldHeadNames,
+    RGBFieldHead,
+    SHFieldHead,
+)
 
 
 def test_field_output():
@@ -40,7 +46,20 @@ def test_rgb_output():
     rgb_head(x)
 
 
+def test_sh_output():
+    """Test sh output"""
+    in_dim = 6
+    levels = 4
+    channels = 3
+    rgb_head = SHFieldHead(in_dim, levels=levels, channels=channels)
+    assert rgb_head.get_out_dim() == channels * levels**2
+
+    x = torch.ones((9, in_dim))
+    rgb_head(x)
+
+
 if __name__ == "__main__":
     test_field_output()
     test_density_output()
     test_rgb_output()
+    test_sh_output()