From 8e852c577161e2caf6962bd1e53b406c0f0b1525 Mon Sep 17 00:00:00 2001
From: Matt Tancik <tancik@berkeley.edu>
Date: Fri, 15 Apr 2022 14:38:26 -0700
Subject: [PATCH] Rename field heads

---
 .../{render_heads.py => field_outputs.py}          | 12 ++++++------
 mattport/nerf/renderers.py                         | 12 ++++++------
 ...{test_render_heads.py => test_field_outputs.py} | 14 +++++++-------
 3 files changed, 19 insertions(+), 19 deletions(-)
 rename mattport/nerf/field_modules/{render_heads.py => field_outputs.py} (90%)
 rename tests/nerf/field_modules/{test_render_heads.py => test_field_outputs.py} (72%)

diff --git a/mattport/nerf/field_modules/render_heads.py b/mattport/nerf/field_modules/field_outputs.py
similarity index 90%
rename from mattport/nerf/field_modules/render_heads.py
rename to mattport/nerf/field_modules/field_outputs.py
index 4504cad72..262b53a00 100644
--- a/mattport/nerf/field_modules/render_heads.py
+++ b/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)
diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py
index bb8b7aed6..d319ea74e 100644
--- a/mattport/nerf/renderers.py
+++ b/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,
diff --git a/tests/nerf/field_modules/test_render_heads.py b/tests/nerf/field_modules/test_field_outputs.py
similarity index 72%
rename from tests/nerf/field_modules/test_render_heads.py
rename to tests/nerf/field_modules/test_field_outputs.py
index d05caeb7e..cec54dd90 100644
--- a/tests/nerf/field_modules/test_render_heads.py
+++ b/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))