diff --git a/mattport/nerf/modules/__init__.py b/mattport/nerf/field_modules/__init__.py similarity index 56% rename from mattport/nerf/modules/__init__.py rename to mattport/nerf/field_modules/__init__.py index 2a96e1cc9..80c8bcdb2 100644 --- a/mattport/nerf/modules/__init__.py +++ b/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 diff --git a/mattport/nerf/modules/module.py b/mattport/nerf/field_modules/base.py similarity index 91% rename from mattport/nerf/modules/module.py rename to mattport/nerf/field_modules/base.py index 57dc6b5e7..b3c0a8e52 100644 --- a/mattport/nerf/modules/module.py +++ b/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""" diff --git a/mattport/nerf/modules/encoding.py b/mattport/nerf/field_modules/encoding.py similarity index 98% rename from mattport/nerf/modules/encoding.py rename to mattport/nerf/field_modules/encoding.py index da418f406..751370c6b 100644 --- a/mattport/nerf/modules/encoding.py +++ b/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: diff --git a/tests/nerf/configs/__init__.py b/mattport/nerf/field_modules/field_losses.py similarity index 100% rename from tests/nerf/configs/__init__.py rename to mattport/nerf/field_modules/field_losses.py diff --git a/mattport/nerf/render_heads.py b/mattport/nerf/field_modules/field_outputs.py similarity index 57% rename from mattport/nerf/render_heads.py rename to mattport/nerf/field_modules/field_outputs.py index be509d80b..262b53a00 100644 --- a/mattport/nerf/render_heads.py +++ b/mattport/nerf/field_modules/field_outputs.py @@ -6,22 +6,31 @@ from typing import Optional from torch import nn from torchtyping import TensorType +from mattport.nerf.field_modules.base import FieldModule -class RenderHead(nn.Module): - """Base Render head""" - def __init__(self, in_dim: int, out_dim: int, activation: Optional[nn.Module]) -> None: +class FieldOutput(FieldModule): + """Base field output""" + + def __init__(self, in_dim: int, out_dim: int, activation: Optional[nn.Module] = None) -> 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) + print(layers) self.net = nn.Sequential(*layers) def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType[..., "out_dim"]: @@ -33,19 +42,20 @@ 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) -class DensityHead(RenderHead): - """Density head""" +class DensityFieldOutput(FieldOutput): + """Density output""" - def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Softplus) -> None: + 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: + def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Sigmoid()) -> None: super().__init__(in_dim, 3, activation) diff --git a/mattport/nerf/modules/mlp.py b/mattport/nerf/field_modules/mlp.py similarity index 96% rename from mattport/nerf/modules/mlp.py rename to mattport/nerf/field_modules/mlp.py index 6710742a3..008920d97 100644 --- a/mattport/nerf/modules/mlp.py +++ b/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__( diff --git a/mattport/nerf/modules/ray_generator.py b/mattport/nerf/field_modules/ray_generator.py similarity index 100% rename from mattport/nerf/modules/ray_generator.py rename to mattport/nerf/field_modules/ray_generator.py diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py index cd6108080..5dd87440d 100644 --- a/mattport/nerf/graph.py +++ b/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) diff --git a/mattport/nerf/loss.py b/mattport/nerf/loss.py index e69de29bb..d59f38dcd 100644 --- a/mattport/nerf/loss.py +++ b/mattport/nerf/loss.py @@ -0,0 +1,9 @@ +""" +Collection of Losses +""" +from torch import nn + +L1Loss = nn.L1Loss +MSELoss = nn.MSELoss + +LOSSES = {"L1": L1Loss, "MSE": MSELoss} diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py index e7b3c4a46..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 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/modules/test_encodings.py b/tests/nerf/field_modules/test_encodings.py similarity index 97% rename from tests/nerf/modules/test_encodings.py rename to tests/nerf/field_modules/test_encodings.py index 47bd74668..6798ead49 100644 --- a/tests/nerf/modules/test_encodings.py +++ b/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(): diff --git a/tests/nerf/field_modules/test_field_outputs.py b/tests/nerf/field_modules/test_field_outputs.py new file mode 100644 index 000000000..cec54dd90 --- /dev/null +++ b/tests/nerf/field_modules/test_field_outputs.py @@ -0,0 +1,56 @@ +""" +Render Head Tests +""" +import pytest +import torch +from torch import nn + +from mattport.nerf.field_modules import field_outputs + + +def test_field_output(): + """Test rgb render head""" + in_dim = 6 + out_dim = 4 + activation = nn.ReLU() + 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)) + + with pytest.raises(SystemError): + render_head(x) + render_head.build_nn_modules() + y = render_head(x) + + assert y.shape[-1] == out_dim + + +def test_density_head(): + """Test rgb output""" + in_dim = 6 + density_head = field_outputs.DensityFieldOutput(in_dim) + assert density_head.get_out_dim() == 1 + + x = torch.ones((9, in_dim)) + with pytest.raises(SystemError): + density_head(x) + density_head.build_nn_modules() + y = density_head(x) + + assert y.shape[-1] == 1 + + +def test_rgb_head(): + """Test rgb output""" + in_dim = 6 + rgb_head = field_outputs.RGBFieldOutput(in_dim) + assert rgb_head.get_out_dim() == 3 + + x = torch.ones((9, in_dim)) + with pytest.raises(SystemError): + rgb_head(x) + rgb_head.build_nn_modules() + y = rgb_head(x) + + assert y.shape[-1] == 3 diff --git a/tests/nerf/graph/configs/__init__.py b/tests/nerf/graph/configs/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/nerf/configs/test_basic_graph.yaml b/tests/nerf/graph/configs/test_basic_graph.yaml similarity index 100% rename from tests/nerf/configs/test_basic_graph.yaml rename to tests/nerf/graph/configs/test_basic_graph.yaml diff --git a/tests/nerf/configs/test_complex_graph.yaml b/tests/nerf/graph/configs/test_complex_graph.yaml similarity index 100% rename from tests/nerf/configs/test_complex_graph.yaml rename to tests/nerf/graph/configs/test_complex_graph.yaml diff --git a/tests/nerf/configs/test_repeat_graph.yaml b/tests/nerf/graph/configs/test_repeat_graph.yaml similarity index 100% rename from tests/nerf/configs/test_repeat_graph.yaml rename to tests/nerf/graph/configs/test_repeat_graph.yaml diff --git a/tests/nerf/test_graph.py b/tests/nerf/graph/test_graph.py similarity index 96% rename from tests/nerf/test_graph.py rename to tests/nerf/graph/test_graph.py index 73d8f0a25..e45665341 100644 --- a/tests/nerf/test_graph.py +++ b/tests/nerf/graph/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 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" diff --git a/tests/nerf/test_ray_generator.py b/tests/nerf/test_ray_generator.py index 22ba267cc..e9d3bb36b 100644 --- a/tests/nerf/test_ray_generator.py +++ b/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