From a32c681a7a6b295925b9ec9d61006c9193cb57d8 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Thu, 21 Apr 2022 14:51:51 -0700 Subject: [PATCH] Rearrange nerf files --- .github/workflows/code_checks.yml | 3 +- .vscode/settings.json | 3 +- environment/requirements.txt | 1 + mattport/nerf/field/__init__.py | 0 mattport/nerf/field/nerf.py | 56 ------------------- .../nerf/graph/{nerf.py => vanilla_nerf.py} | 55 ++++++++++++++++-- 6 files changed, 54 insertions(+), 64 deletions(-) delete mode 100644 mattport/nerf/field/__init__.py delete mode 100644 mattport/nerf/field/nerf.py rename mattport/nerf/graph/{nerf.py => vanilla_nerf.py} (60%) diff --git a/.github/workflows/code_checks.yml b/.github/workflows/code_checks.yml index 270a67066..b83fa2ecf 100644 --- a/.github/workflows/code_checks.yml +++ b/.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 diff --git a/.vscode/settings.json b/.vscode/settings.json index 80164259e..ed3e7382f 100644 --- a/.vscode/settings.json +++ b/.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" ], } \ No newline at end of file diff --git a/environment/requirements.txt b/environment/requirements.txt index 9110d4604..9d0aa8b6e 100644 --- a/environment/requirements.txt +++ b/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 diff --git a/mattport/nerf/field/__init__.py b/mattport/nerf/field/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/mattport/nerf/field/nerf.py b/mattport/nerf/field/nerf.py deleted file mode 100644 index 52e9c7880..000000000 --- a/mattport/nerf/field/nerf.py +++ /dev/null @@ -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 diff --git a/mattport/nerf/graph/nerf.py b/mattport/nerf/graph/vanilla_nerf.py similarity index 60% rename from mattport/nerf/graph/nerf.py rename to mattport/nerf/graph/vanilla_nerf.py index 20af5aa7b..ddac3f76c 100644 --- a/mattport/nerf/graph/nerf.py +++ b/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)