Merge pull request #17 from ethanweber/tancik/restructure

Restructure into field modules

mattport/nerf/field_modules/__init__.py

@@ -1,4 +1,4 @@
-"""init modules"""
+"""init field modules"""
 from .encoding import Encoding, ScalingAndOffset
 from .mlp import MLP
-from .module import Module
+from .base import FieldModule

mattport/nerf/field_modules/base.py

@@ -1,5 +1,5 @@
 """
-The module baseclass.
+The field module baseclass.
 """
 from abc import abstractmethod
 
@@ -7,12 +7,8 @@ from torch import nn
 from torchtyping import TensorType
 
 
-class Module(nn.Module):
-    """_summary_
-
-    Args:
-        nn (_type_): _description_
-    """
+class FieldModule(nn.Module):
+    """Field modules that can be combined to store and compute the fields."""
 
     def __init__(self) -> None:
         """Default initialization of module"""

mattport/nerf/field_modules/encoding.py

@@ -3,10 +3,10 @@ Encoding functions
 """
 import torch
 from torchtyping import TensorType
-from mattport.nerf.modules.module import Module
+from mattport.nerf.field_modules.base import FieldModule
 
 
-class Encoding(Module):
+class Encoding(FieldModule):
     """Encode an input tensor. Intended to be subclassed"""
 
     def __init__(self, in_dim: int) -> None:

mattport/nerf/field_modules/field_outputs.py

@@ -6,11 +6,11 @@ from typing import Optional
 from torch import nn
 from torchtyping import TensorType
 
-from mattport.nerf.modules.module import Module
+from mattport.nerf.field_modules.base import FieldModule
 
 
-class RenderHead(Module):
-    """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(Module):
         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/field_modules/mlp.py

@@ -5,10 +5,10 @@ from typing import Optional
 
 from torch import nn
 from torchtyping import TensorType
-from mattport.nerf.modules.module import Module
+from mattport.nerf.field_modules.base import FieldModule
 
 
-class MLP(Module):
+class MLP(FieldModule):
     """Multilayer perceptron"""
 
     def __init__(

mattport/nerf/graph.py

@@ -43,7 +43,7 @@ class Graph(nn.Module):
         self.modules = {}
         # initialize graph with known input dimensions; set default in_dim to 0
         for module_name, module_dict in modules_config.items():
-            module = getattr(importlib.import_module("mattport.nerf.modules"), module_dict.class_name)
+            module = getattr(importlib.import_module("mattport.nerf.field_modules"), module_dict.class_name)
             if not module_dict.meta_data.in_dim:
                 module_dict.meta_data.in_dim = 0
             self.modules[module_name] = module(**module_dict.meta_data)

mattport/nerf/renderers.py

@@ -7,13 +7,13 @@ import torch
 from torch import nn
 from torchtyping import TensorType
 
-from mattport.nerf.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_encodings.py

@@ -3,7 +3,7 @@ Encoding Tests
 """
 import pytest
 import torch
-from mattport.nerf.modules import encoding
+from mattport.nerf.field_modules import encoding
 
 
 def test_scaling_and_offset():

tests/nerf/field_modules/test_field_outputs.py

@@ -5,15 +5,15 @@ import pytest
 import torch
 from torch import nn
 
-from mattport.nerf.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))

tests/nerf/test_graph.py

@@ -1,5 +1,7 @@
+"""
+Tests of graph construction
+"""
 import numpy as np
-import pytest
 from hydra import compose, initialize
 from omegaconf import open_dict
 
@@ -28,7 +30,7 @@ def instantiate_graph(config_name: str) -> Graph:
 
 
 def test_graph_init():
-    """test modules dictionary is instantiated properly"""
+    """test field modules dictionary is instantiated properly"""
     test_graph, _ = instantiate_graph("test_basic_graph")
     assert test_graph.modules is not None
     assert type(test_graph.modules["encoder_0"]).__name__ == "ScalingAndOffset"
@@ -37,7 +39,7 @@ def test_graph_init():
 
 
 def test_input_dimension():
-    """test module input dimension calculation"""
+    """test field module input dimension calculation"""
     test_graph, _ = instantiate_graph("test_repeat_graph")
     assert test_graph.modules["mlp_0"].in_dim == 16
     assert test_graph.modules["mlp_1"].in_dim == 1

tests/nerf/test_ray_generator.py

@@ -4,7 +4,7 @@ Test ray generation.
 import torch
 
 from mattport.nerf.dataset.blender_dataset import load_blender_data
-from mattport.nerf.modules.ray_generator import RayGenerator
+from mattport.nerf.field_modules.ray_generator import RayGenerator
 from mattport.structures.rays import RayBundle
 from mattport.utils.io import get_absolute_path