Rename field heads

mattport/nerf/field_modules/field_outputs.py

@@ -9,8 +9,8 @@ from torchtyping import TensorType
 from mattport.nerf.field_modules.base import FieldModule
 
 
-class RenderHead(FieldModule):
-    """Base Render head"""
+class FieldOutput(FieldModule):
+    """Base field output"""
 
     def __init__(self, in_dim: int, out_dim: int, activation: Optional[nn.Module] = None) -> None:
         """
@@ -47,15 +47,15 @@ class RenderHead(FieldModule):
         return self.net(in_tensor)
 
 
-class DensityHead(RenderHead):
-    """Density head"""
+class DensityFieldOutput(FieldOutput):
+    """Density output"""
 
     def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Softplus()) -> None:
         super().__init__(in_dim, 1, activation)
 
 
-class RGBHead(RenderHead):
-    """RGB head"""
+class RGBFieldOutput(FieldOutput):
+    """RGB output"""
 
     def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Sigmoid()) -> None:
         super().__init__(in_dim, 3, activation)

mattport/nerf/renderers.py

@@ -7,13 +7,13 @@ import torch
 from torch import nn
 from torchtyping import TensorType
 
-from mattport.nerf.field_modules import render_heads as rh
+from mattport.nerf.field_modules import field_outputs
 
 
 class Renderer(nn.Module):
     """Base Renderer. Intended to be subclassed"""
 
-    def required_render_heads(self) -> Set[rh.RenderHead]:
+    def required_render_heads(self) -> Set[field_outputs.FieldOutput]:
         """
         Returns:
             Set[RenderHead]: RenderHeads required for this renderer.
@@ -45,8 +45,8 @@ class RGB(Renderer):
         self.rgb_name = rgb_name
         self.density_name = density_name
 
-    def required_render_heads(self) -> Set[rh.RenderHead]:
-        return set(rh.RGBHead, rh.DensityHead)
+    def required_render_heads(self) -> Set[field_outputs.FieldOutput]:
+        return set(field_outputs.RGBFieldOutput, field_outputs.DensityFieldOutput)
 
     def forward(
         self,
@@ -83,8 +83,8 @@ class Depth(Renderer):
         Renderer (_type_): _description_
     """
 
-    def required_render_heads(self) -> Set[rh.RenderHead]:
-        return set(rh.DensityHead)
+    def required_render_heads(self) -> Set[field_outputs.FieldOutput]:
+        return set(field_outputs.DensityFieldOutput)
 
     def forward(
         self,

tests/nerf/field_modules/test_field_outputs.py

@@ -5,15 +5,15 @@ import pytest
 import torch
 from torch import nn
 
-from mattport.nerf.field_modules import render_heads
+from mattport.nerf.field_modules import field_outputs
 
 
-def test_render_head():
+def test_field_output():
     """Test rgb render head"""
     in_dim = 6
     out_dim = 4
     activation = nn.ReLU()
-    render_head = render_heads.RenderHead(in_dim=in_dim, out_dim=out_dim, activation=activation)
+    render_head = field_outputs.FieldOutput(in_dim=in_dim, out_dim=out_dim, activation=activation)
     assert render_head.get_out_dim() == out_dim
 
     x = torch.ones((9, in_dim))
@@ -27,9 +27,9 @@ def test_render_head():
 
 
 def test_density_head():
-    """Test rgb render head"""
+    """Test rgb output"""
     in_dim = 6
-    density_head = render_heads.DensityHead(in_dim)
+    density_head = field_outputs.DensityFieldOutput(in_dim)
     assert density_head.get_out_dim() == 1
 
     x = torch.ones((9, in_dim))
@@ -42,9 +42,9 @@ def test_density_head():
 
 
 def test_rgb_head():
-    """Test rgb render head"""
+    """Test rgb output"""
     in_dim = 6
-    rgb_head = render_heads.RGBHead(in_dim)
+    rgb_head = field_outputs.RGBFieldOutput(in_dim)
     assert rgb_head.get_out_dim() == 3
 
     x = torch.ones((9, in_dim))