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/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/field_modules/render_heads.py

@@ -6,10 +6,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 RenderHead(Module):
+class RenderHead(FieldModule):
     """Base Render head"""
 
     def __init__(self, in_dim: int, out_dim: int, activation: Optional[nn.Module] = None) -> None:

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,7 +7,7 @@ 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 render_heads as rh
 
 
 class Renderer(nn.Module):

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_render_heads.py

@@ -5,7 +5,7 @@ import pytest
 import torch
 from torch import nn
 
-from mattport.nerf.modules import render_heads
+from mattport.nerf.field_modules import render_heads
 
 
 def test_render_head():

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