Merge branch 'master' of github.com:ethanweber/mattport

.github/workflows/code_checks.yml

@@ -36,7 +36,8 @@ jobs:
             --good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W,ts,t,u,k,v \
             --disable=too-many-arguments \
             --disable=too-many-locals \
-            --disable=fixme
+            --disable=fixme \
+            --disable=too-many-instance-attributes
       - name: Test with pytest
         run: |
           pytest tests --ignore tests/nerf/dataloader/test_ray_generator.py --typeguard-packages=mattport --torchtyping-patch-typeguard

.vscode/settings.json

@@ -32,6 +32,7 @@
         "--good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W,ts,t,u,k,v",
         "--disable=too-many-arguments",
         "--disable=fixme",
-        "--disable=too-many-locals"
+        "--disable=too-many-locals",
+        "--disable=too-many-instance-attributes"
     ],
 }

environment/requirements.txt

@@ -8,6 +8,7 @@ mediapy
 plotly==5.7.0
 pylint==2.13.4
 pytest
+tensorboard==2.8.0
 torch==1.11.0
 torchtyping
 torchvision==0.12.0

mattport/nerf/field/nerf.py

@@ -1,56 +0,0 @@
-"""
-Multi Layer Perceptron
-"""
-
-import torch
-from torch import nn
-
-from mattport.nerf.field_modules.encoding import NeRFEncoding
-from mattport.nerf.field_modules.field_heads import DensityFieldHead, RGBFieldHead
-from mattport.nerf.field_modules.mlp import MLP
-from mattport.structures.rays import RaySamples
-
-
-class NeRFField(nn.Module):
-    """Multilayer perceptron"""
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.build_nn_modules()
-
-    def build_nn_modules(self) -> None:
-        """Initialize modules."""
-        self.encoding_xyz = NeRFEncoding(in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0)
-        self.encoding_dir = NeRFEncoding(in_dim=3, num_frequencies=6, min_freq_exp=0.0, max_freq_exp=4.0)
-        self.mlp_base = MLP(
-            in_dim=self.encoding_xyz.get_out_dim(), out_dim=64, num_layers=8, layer_width=64, activation=nn.ReLU()
-        )
-        self.mlp_rgb = MLP(
-            in_dim=self.mlp_base.get_out_dim() + self.encoding_dir.get_out_dim(),
-            out_dim=64,
-            num_layers=2,
-            layer_width=64,
-            activation=nn.ReLU(),
-        )
-        self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_rgb.get_out_dim())
-        self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim())
-
-    def forward(self, ray_samples: RaySamples):
-        """Evaluates the field at points along the ray
-        Args:
-            xyz: ()
-        # TODO(ethan): change the input to be something more abstracted
-        e.g., a FieldInput structure
-        """
-        positions = ray_samples.positions
-        directions = ray_samples.directions
-        encoded_xyz = self.encoding_xyz(positions)
-        encoded_dir = self.encoding_dir(directions)
-        base_mlp_out = self.mlp_base(encoded_xyz)
-        rgb_mlp_out = self.mlp_rgb(torch.cat([encoded_dir, base_mlp_out], dim=-1))
-
-        field_rgb_output = self.field_output_rgb(rgb_mlp_out)
-        field_density_out = self.field_output_density(base_mlp_out)
-
-        field_outputs = field_rgb_output | field_density_out
-        return field_outputs

mattport/nerf/graph/vanilla_nerf.py

@@ -4,23 +4,66 @@ Implementation of vanilla nerf.
 
 
 from typing import Dict, List
+
+import torch
+from torch import nn
 from torch.nn import Parameter
 from torchtyping import TensorType
 
-from mattport.nerf.field.nerf import NeRFField
+from mattport.nerf.field_modules.encoding import NeRFEncoding
+from mattport.nerf.field_modules.field_heads import DensityFieldHead, RGBFieldHead
+from mattport.nerf.field_modules.mlp import MLP
 from mattport.nerf.field_modules.ray_generator import RayGenerator
 from mattport.nerf.graph.base import Graph
 from mattport.nerf.loss import MSELoss
 from mattport.nerf.renderers import RGBRenderer
 from mattport.nerf.sampler import PDFSampler, UniformSampler
+from mattport.structures.rays import RaySamples
+
+
+class NeRFField(nn.Module):
+    """NeRF module"""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.encoding_xyz = NeRFEncoding(in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0)
+        self.encoding_dir = NeRFEncoding(in_dim=3, num_frequencies=6, min_freq_exp=0.0, max_freq_exp=4.0)
+        self.mlp_base = MLP(
+            in_dim=self.encoding_xyz.get_out_dim(), out_dim=64, num_layers=8, layer_width=64, activation=nn.ReLU()
+        )
+        self.mlp_rgb = MLP(
+            in_dim=self.mlp_base.get_out_dim() + self.encoding_dir.get_out_dim(),
+            out_dim=64,
+            num_layers=2,
+            layer_width=64,
+            activation=nn.ReLU(),
+        )
+        self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_rgb.get_out_dim())
+        self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim())
+
+    def forward(self, ray_samples: RaySamples):
+        """Evaluates the field at points along the ray
+        Args:
+            xyz: ()
+        # TODO(ethan): change the input to be something more abstracted
+        e.g., a FieldInput structure
+        """
+        positions = ray_samples.positions
+        directions = ray_samples.directions
+        encoded_xyz = self.encoding_xyz(positions)
+        encoded_dir = self.encoding_dir(directions)
+        base_mlp_out = self.mlp_base(encoded_xyz)
+        rgb_mlp_out = self.mlp_rgb(torch.cat([encoded_dir, base_mlp_out], dim=-1))
+
+        field_rgb_output = self.field_output_rgb(rgb_mlp_out)
+        field_density_out = self.field_output_density(base_mlp_out)
+
+        field_outputs = field_rgb_output | field_density_out
+        return field_outputs
 
 
 class NeRFGraph(Graph):
-    """_summary_
-
-    Args:
-        Graph (_type_): _description_
-    """
+    """Vanilla NeRF graph"""
 
     def __init__(self, intrinsics=None, camera_to_world=None) -> None:
         super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world)