Update render_heads to use Module class

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)
 
 

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

tests/nerf/test_encoding.py

@@ -1,9 +0,0 @@
-import pytest
-
-
-def test_positional_encoding():
-    assert True
-
-
-def test_lower_case():
-    assert "MATTPORT".lower() == "mattport"