Merge pull request #93 from ethanweber/tancik/allow_multiple_nerf_heads

Refactor field heads

mattport/nerf/field_modules/field_heads.py

@@ -26,17 +26,35 @@ class FieldHeadNames(Enum):
 class FieldHead(FieldModule):
     """Base field output"""
 
-    def __init__(self, in_dim: int, out_dim: int, activation: Optional[nn.Module] = None) -> None:
+    def __init__(
+        self,
+        out_dim: int,
+        field_head_name: FieldHeadNames,
+        in_dim: Optional[int] = None,
+        activation: Optional[nn.Module] = None,
+    ) -> None:
         """
         Args:
-            in_dim (int): input dimension
             out_dim (int): output dimension for renderer
+            field_head_name (FieldHeadNames): Field type
+            in_dim (int, optional): input dimension. If not defined in constructor, it must be set later.
             activation (Optional[nn.Module]): output head activation
         """
         super().__init__()
-        self.in_dim = in_dim
         self.out_dim = out_dim
         self.activation = activation
+        self.field_head_name = field_head_name
+        self.net = None
+        if in_dim is not None:
+            self.in_dim = in_dim
+            self._construct_net()
+
+    def set_in_dim(self, in_dim: int) -> None:
+        """Set input dimension of Field Head"""
+        self.in_dim = in_dim
+        self._construct_net()
+
+    def _construct_net(self):
         layers = [nn.Linear(self.in_dim, self.out_dim)]
         if self.activation:
             layers.append(self.activation)
@@ -52,7 +70,7 @@ class FieldHead(FieldModule):
             TensorType[..., "out_dim"]: Render head output
         """
         if not self.net:
-            raise SystemError("Render head network not initialized. build_nn_modules() should be called.")
+            raise SystemError("in_dim not set. Must be provided to construtor, or set_in_dim() should be called.")
         out_tensor = self.net(in_tensor)
         return out_tensor
 
@@ -60,23 +78,24 @@ class FieldHead(FieldModule):
 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, activation=activation)
+    def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None:
+        super().__init__(in_dim=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, activation=activation)
+    def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Sigmoid()) -> None:
+        super().__init__(in_dim=in_dim, out_dim=3, field_head_name=FieldHeadNames.RGB, activation=activation)
 
 
 class SHFieldHead(FieldHead):
     """Spherical harmonics output"""
 
-    def __init__(self, in_dim: int, levels: int = 3, channels: int = 3, activation: Optional[nn.Module] = None) -> None:
-        """_summary_
-
+    def __init__(
+        self, in_dim: Optional[int] = None, levels: int = 3, channels: int = 3, activation: Optional[nn.Module] = None
+    ) -> None:
+        """
         Args:
             in_dim (int): Input dimension
             levels (int, optional): Number of spherical harmonics layers. Defaults to 3.
@@ -84,37 +103,41 @@ 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, activation=activation)
+        super().__init__(in_dim=in_dim, out_dim=out_dim, field_head_name=FieldHeadNames.SH, activation=activation)
 
 
 class UncertaintyFieldHead(FieldHead):
     """Uncertainty output"""
 
-    def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Softplus()) -> None:
-        super().__init__(in_dim, out_dim=1, activation=activation)
+    def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None:
+        super().__init__(in_dim=in_dim, out_dim=1, field_head_name=FieldHeadNames.UNCERTAINTY, activation=activation)
 
 
 class TransientRGBHead(FieldHead):
     """Transient RGB output"""
 
-    def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Sigmoid()) -> None:
-        super().__init__(in_dim, out_dim=3, activation=activation)
+    def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Sigmoid()) -> None:
+        super().__init__(in_dim=in_dim, out_dim=3, field_head_name=FieldHeadNames.TRANSIENT_RGB, activation=activation)
 
 
 class TransientDensityHead(FieldHead):
     """Transient density output"""
 
-    def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Softplus()) -> None:
-        super().__init__(in_dim, out_dim=1, activation=activation)
+    def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None:
+        super().__init__(
+            in_dim=in_dim, out_dim=1, field_head_name=FieldHeadNames.TRANSIENT_DENSITY, activation=activation
+        )
 
 
 class SemanticStuffHead(FieldHead):
     """Semantic stuff output"""
 
-    def __init__(self, in_dim: int, num_classes: int) -> None:
+    def __init__(self, num_classes: int, in_dim: Optional[int] = None) -> None:
         """
         Args:
-            in_dim (int): Input dimension
             num_classes (int): Number of semantic classes
+            in_dim (int): Input dimension
         """
-        super().__init__(in_dim, out_dim=num_classes, activation=None)
+        super().__init__(
+            in_dim=in_dim, out_dim=num_classes, field_head_name=FieldHeadNames.SEMANTICS_STUFF, activation=None
+        )

mattport/nerf/fields/base.py

@@ -7,6 +7,7 @@ from typing import Tuple
 
 from torch import nn
 from torchtyping import TensorType
+from mattport.nerf.field_modules.field_heads import FieldHeadNames
 
 from mattport.structures.rays import PointSamples
 
@@ -39,5 +40,5 @@ class Field(nn.Module):
         """Evaluates the field at points along the ray."""
         density, density_embedding = self.get_density(point_samples)
         field_outputs = self.get_outputs(point_samples, density_embedding=density_embedding)
-        field_outputs["density"] = density
+        field_outputs[FieldHeadNames.DENSITY] = density
         return field_outputs

mattport/nerf/fields/nerf_field.py

@@ -4,8 +4,10 @@
 from typing import Tuple
 
 import torch
+from torch import nn
+
 from mattport.nerf.field_modules.encoding import Encoding, Identity
-from mattport.nerf.field_modules.field_heads import DensityFieldHead, RGBFieldHead
+from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHead, RGBFieldHead
 from mattport.nerf.field_modules.mlp import MLP
 from mattport.nerf.fields.base import Field
 from mattport.structures.rays import PointSamples
@@ -23,7 +25,8 @@ class NeRFField(Field):
         base_mlp_layer_width: int = 256,
         head_mlp_num_layers: int = 2,
         head_mlp_layer_width: int = 128,
-        skip_connections: Tuple = (4,),
+        skip_connections: Tuple[int] = (4,),
+        field_heads: Tuple[FieldHead] = (RGBFieldHead(),),
     ) -> None:
         """Create boilerplate NeRF field.
 
@@ -54,7 +57,9 @@ class NeRFField(Field):
         )
 
         self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim())
-        self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_head.get_out_dim())
+        self.field_heads = nn.ModuleList(field_heads)
+        for field_head in self.field_heads:
+            field_head.set_in_dim(self.mlp_head.get_out_dim())
 
     def get_density(self, point_samples: PointSamples):
         """Computes and returns the densities."""
@@ -79,11 +84,14 @@ class NeRFField(Field):
         directions = point_samples.directions
         if not is_not_none(valid_mask):
             valid_mask = torch.ones_like(directions[..., 0]).bool()
-        # placeholders for values to return
-        rgb = torch.zeros(*valid_mask.shape, 3, dtype=torch.float32, device=directions.device)
-        if not valid_mask.any():  # empty mask
-            return {"rgb": rgb}
-        encoded_dir = self.direction_encoding(directions[valid_mask])
-        rgb_mlp_out = self.mlp_head(torch.cat([encoded_dir, density_embedding[valid_mask]], dim=-1))
-        rgb[valid_mask] = self.field_output_rgb(rgb_mlp_out)
-        return {"rgb": rgb}
+        outputs = {}
+        for field_head in self.field_heads:
+            # placeholders for values to return
+            out = torch.zeros(*valid_mask.shape, field_head.out_dim, dtype=torch.float32, device=directions.device)
+            if not valid_mask.any():  # empty mask
+                return {field_head.field_head_name: out}
+            encoded_dir = self.direction_encoding(directions[valid_mask])
+            mlp_out = self.mlp_head(torch.cat([encoded_dir, density_embedding[valid_mask]], dim=-1))
+            out[valid_mask] = field_head(mlp_out)
+            outputs[field_head.field_head_name] = out
+        return outputs

mattport/nerf/graph/vanilla_nerf.py

@@ -10,6 +10,7 @@ from torch.nn import Parameter
 from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
 from mattport.nerf.field_modules.encoding import NeRFEncoding
+from mattport.nerf.field_modules.field_heads import FieldHeadNames
 
 from mattport.nerf.fields.nerf_field import NeRFField
 from mattport.nerf.graph.base import Graph
@@ -88,9 +89,9 @@ class NeRFGraph(Graph):
 
         # coarse field:
         field_outputs_coarse = self.field_coarse.forward(ray_samples_uniform.to_point_samples())
-        weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse["density"])
+        weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY])
         rgb_coarse = self.renderer_rgb(
-            rgb=field_outputs_coarse["rgb"],
+            rgb=field_outputs_coarse[FieldHeadNames.RGB],
             weights=weights_coarse,
         )
         accumulation_coarse = self.renderer_accumulation(weights_coarse)
@@ -101,9 +102,9 @@ class NeRFGraph(Graph):
 
         # fine field:
         field_outputs_fine = self.field_fine.forward(ray_samples_pdf.to_point_samples())
-        weights_fine = ray_samples_pdf.get_weights(field_outputs_fine["density"])
+        weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY])
         rgb_fine = self.renderer_rgb(
-            rgb=field_outputs_fine["rgb"],
+            rgb=field_outputs_fine[FieldHeadNames.RGB],
             weights=weights_fine,
         )
         accumulation_fine = self.renderer_accumulation(weights_fine)

tests/nerf/field_modules/test_field_outputs.py

@@ -3,10 +3,12 @@ Field output tests
 """
 import torch
 from torch import nn
+import pytest
 
 from mattport.nerf.field_modules.field_heads import (
     DensityFieldHead,
     FieldHead,
+    FieldHeadNames,
     RGBFieldHead,
     SHFieldHead,
 )
@@ -17,12 +19,19 @@ def test_field_output():
     in_dim = 6
     out_dim = 4
     activation = nn.ReLU()
-    render_head = FieldHead(in_dim=in_dim, out_dim=out_dim, activation=activation)
+    render_head = FieldHead(in_dim=in_dim, out_dim=out_dim, field_head_name=FieldHeadNames.RGB, activation=activation)
     assert render_head.get_out_dim() == out_dim
 
     x = torch.ones((9, in_dim))
     render_head(x)
 
+    # Test in_dim not provided at construction
+    render_head = FieldHead(out_dim=out_dim, field_head_name=FieldHeadNames.RGB, activation=activation)
+    with pytest.raises(SystemError):
+        render_head(x)
+    render_head.set_in_dim(in_dim)
+    render_head(x)
+
 
 def test_density_output():
     """Test rgb output"""