From 03738b1066a34bdbe294f88fad0d3b106d84b150 Mon Sep 17 00:00:00 2001
From: Matt Tancik <tancik@berkeley.edu>
Date: Thu, 14 Apr 2022 17:31:10 -0700
Subject: [PATCH] Update render_heads to use Module class

---
 mattport/nerf/{ => modules}/render_heads.py | 21 +++++++++++++++------
 tests/nerf/modules/test_render_heads.py     | 18 ++++++++++++++++++
 tests/nerf/test_encoding.py                 |  9 ---------
 3 files changed, 33 insertions(+), 15 deletions(-)
 rename mattport/nerf/{ => modules}/render_heads.py (71%)
 create mode 100644 tests/nerf/modules/test_render_heads.py
 delete mode 100644 tests/nerf/test_encoding.py

diff --git a/mattport/nerf/render_heads.py b/mattport/nerf/modules/render_heads.py
similarity index 71%
rename from mattport/nerf/render_heads.py
rename to mattport/nerf/modules/render_heads.py
index be509d80b..5efdceec1 100644
--- a/mattport/nerf/render_heads.py
+++ b/mattport/nerf/modules/render_heads.py
@@ -6,22 +6,30 @@ from typing import Optional
 from torch import nn
 from torchtyping import TensorType
 
+from mattport.nerf.modules.module import Module
 
-class RenderHead(nn.Module):
+
+class RenderHead(Module):
     """Base Render head"""
 
     def __init__(self, in_dim: int, out_dim: int, activation: Optional[nn.Module]) -> None:
         """
-
         Args:
             in_dim (int): input dimension
             out_dim (int): output dimension for renderer
             activation (Optional[nn.Module]): output head activation
         """
         super().__init__()
-        layers = [nn.Linear(in_dim, out_dim)]
-        if activation:
-            layers.append(activation)
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.activation = activation
+
+        self.net = None
+
+    def build_nn_modules(self) -> None:
+        layers = [nn.Linear(self.in_dim, self.out_dim)]
+        if self.activation:
+            layers.append(self.activation)
         self.net = nn.Sequential(*layers)
 
     def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType[..., "out_dim"]:
@@ -33,7 +41,8 @@ class RenderHead(nn.Module):
         Returns:
             TensorType[..., "out_dim"]: Render head output
         """
-
+        if not self.net:
+            raise SystemError("Render head network not initialized. build_nn_modules() should be called.")
         return self.net(in_tensor)
 
 
diff --git a/tests/nerf/modules/test_render_heads.py b/tests/nerf/modules/test_render_heads.py
new file mode 100644
index 000000000..0f5f4b7dd
--- /dev/null
+++ b/tests/nerf/modules/test_render_heads.py
@@ -0,0 +1,18 @@
+"""
+Render Head Tests
+"""
+from mattport.nerf.modules import render_heads
+
+
+def test_density_head():
+    """Test rgb render head"""
+    in_dim = 6
+    density_head = render_heads.DensityHead(in_dim)
+    assert density_head.get_out_dim() == 1
+
+
+def test_rgb_head():
+    """Test rgb render head"""
+    in_dim = 6
+    rgb_head = render_heads.RGBHead(in_dim)
+    assert rgb_head.get_out_dim() == 3
diff --git a/tests/nerf/test_encoding.py b/tests/nerf/test_encoding.py
deleted file mode 100644
index 99f26096c..000000000
--- a/tests/nerf/test_encoding.py
+++ /dev/null
@@ -1,9 +0,0 @@
-import pytest
-
-
-def test_positional_encoding():
-    assert True
-
-
-def test_lower_case():
-    assert "MATTPORT".lower() == "mattport"