diff --git a/mattport/nerf/field/__init__.py b/mattport/nerf/field/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/mattport/nerf/field/nerf.py b/mattport/nerf/field/nerf.py new file mode 100644 index 000000000..d250634d8 --- /dev/null +++ b/mattport/nerf/field/nerf.py @@ -0,0 +1,52 @@ +""" +Multi Layer Perceptron +""" +from typing import Optional + +from torch import nn +from torchtyping import TensorType +from mattport.nerf.field_modules.encoding import NeRFEncoding +from mattport.nerf.field_modules.field_outputs import DensityFieldOutput, RGBFieldOutput +from mattport.nerf.field_modules.mlp import MLP + + +class NeRFField(): + """Multilayer perceptron""" + + def __init__(self) -> None: + self.build_nn_modules() + + def build_nn_modules(self) -> None: + """Initialize mulilayer perceptron.""" + 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.base_mlp.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 = RGBFieldOutput(in_dim=self.mlp_rgb.get_out_dim()) + self.field_output_density = DensityFieldOutput(in_dim=self.mlp_base.get_out_dim()) + + def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType[..., "out_dim"]: + """Process input with a multilayer perceptron. + + Args: + in_tensor (TensorType[..., "in_dim]): Network input + + Returns: + TensorType[..., "out_dim"]: Network output + """ + + encoded_xyz = self.encoding_xyz(xyz) + encoded_dir = self.encoding_dir(rays.direction) + 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) + return rgb_coarse diff --git a/mattport/nerf/field_modules/base.py b/mattport/nerf/field_modules/base.py index b3c0a8e52..521db07a7 100644 --- a/mattport/nerf/field_modules/base.py +++ b/mattport/nerf/field_modules/base.py @@ -19,6 +19,7 @@ class FieldModule(nn.Module): def build_nn_modules(self) -> None: """Function instantiates any torch.nn members within the module. If none exist, do nothing.""" + raise ValueError("This function shouldn't be called.") def set_in_dim(self, in_dim: int) -> None: """Sets input dimension of encoding diff --git a/mattport/nerf/field_modules/field_outputs.py b/mattport/nerf/field_modules/field_outputs.py index 262b53a00..ad138773e 100644 --- a/mattport/nerf/field_modules/field_outputs.py +++ b/mattport/nerf/field_modules/field_outputs.py @@ -23,8 +23,7 @@ class FieldOutput(FieldModule): self.in_dim = in_dim self.out_dim = out_dim self.activation = activation - - self.net = None + self.net = self.build_nn_modules() def build_nn_modules(self) -> None: layers = [nn.Linear(self.in_dim, self.out_dim)] diff --git a/mattport/nerf/field_modules/mlp.py b/mattport/nerf/field_modules/mlp.py index 57c99a900..283027e46 100644 --- a/mattport/nerf/field_modules/mlp.py +++ b/mattport/nerf/field_modules/mlp.py @@ -37,7 +37,7 @@ class MLP(FieldModule): self.layer_width = layer_width self.activation = activation self.out_activation = out_activation - self.net = None + self.net = self.build_nn_modules() def build_nn_modules(self) -> None: """Initialize mulilayer perceptron.""" diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph/base.py similarity index 93% rename from mattport/nerf/graph.py rename to mattport/nerf/graph/base.py index 1f951e8bd..332534e7c 100644 --- a/mattport/nerf/graph.py +++ b/mattport/nerf/graph/base.py @@ -4,6 +4,7 @@ The Graph module contains all trainable parameters. import importlib from dataclasses import dataclass from typing import Dict, List, Optional, Set +from black import Priority from omegaconf import DictConfig from torch import nn @@ -146,4 +147,10 @@ class Graph(nn.ModuleDict): # x = self.modules[name](x) ## can index into previous results using dictionary - raise NotImplementedError + # raise NotImplementedError + for inputs, node_name, out_name in priority_queue: + self.outputs[out_name] = self.modules[node_name](inputs) + + # possibly delete things from output if no longer needed + for old_name in : + del self.outputs[old_name] diff --git a/mattport/nerf/graph/nerf.py b/mattport/nerf/graph/nerf.py new file mode 100644 index 000000000..c8eac50e6 --- /dev/null +++ b/mattport/nerf/graph/nerf.py @@ -0,0 +1,79 @@ +""" +Implementation of vanilla nerf. +""" + + +import torch +from omegaconf import DictConfig +from torch import TensorType, nn + +from mattport.nerf.field_modules.encoding import NeRFEncoding +from mattport.nerf.field_modules.field_outputs import DensityFieldOutput, RGBFieldOutput +from mattport.nerf.field_modules.mlp import MLP +from mattport.nerf.field.nerf import NeRFField +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 UniformSampler +from mattport.structures.cameras import Rays +from mattport.nerf.dataset.utils import DatasetInputs + + +class NeRFGraph(Graph): + def __init__(self, modules_config: DictConfig, dataset_inputs: DatasetInputs) -> None: + super().__init__(modules_config) + + self.populate_modules(dataset_inputs=dataset_inputs) + # TODO(ethan): build all modules here instead of in the code below + + def populate_modules(self, dataset_inputs: DatasetInputs): + + # ray generator + self.ray_generator = RayGenerator(dataset_inputs.intrinsics, dataset_inputs.camera_to_world) + + # samplers + self.sampler_uniform = UniformSampler(near_plane=0.1, far_plane=4.0, num_samples=64) + self.sampler_pdf = UniformSampler(near_plane=0.1, far_plane=4.0, num_samples=64) + + # field + self.coarse = NeRFField() + self.fine = NeRFField() + + # renderers + self.renderer_rgb = RGBRenderer() + + # losses + self.rgb_loss = MSELoss() + + def forward(self, batch_indices: TensorType["num_rays", 3]): + + ray_bundle = self.ray_generator.forward(batch_indices) + rays = Rays(origin=ray_bundle.origins, direction=ray_bundle.directions) # TODO(ethan): consolidate this + ray_samples = self.sampler_uniform(rays) + + xyz = ray_samples.locations + deltas = ray_samples.deltas + + encoded_xyz = self.encoding_xyz(xyz) + encoded_dir = self.encoding_dir(rays.direction) + 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) + rgb_coarse = self.renderer_rgb(rgb=field_rgb_output, density=field_density_out, deltas=deltas) + + xyz_pdf = self.sampler_pdf(rays) # TODO(ethan): implement pdf + encoded_xyz_pdf = self.encoding_xyz(xyz_pdf) + base_mlp_out = self.mlp_base(encoded_xyz_pdf) + 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) + rgb_fine = self.renderer_rgb(rgb=field_rgb_output, density=field_density_out, deltas=deltas) + + return {"rgb_coarse": rgb_coarse, "rgb_fine": rgb_fine} + + def get_loss(self, batch, render_outputs): + rgb_loss_coarse = self.rgb_loss(batch.rgb, render_outputs["rgb_coarse"]) + rgb_loss_fine = self.rgb_loss(batch.rgb, render_outputs["rgb_fine"]) + return {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine} diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py index 1e78e04aa..6a54d0e3f 100644 --- a/mattport/nerf/renderers.py +++ b/mattport/nerf/renderers.py @@ -37,20 +37,19 @@ class Renderer(nn.Module): raise NotImplementedError -class RGB(Renderer): +class RGBRenderer(Renderer): """Standard volumetic rendering.""" - def __init__(self, rgb_name, density_name) -> None: + def __init__(self) -> None: super().__init__() - self.rgb_name = rgb_name - self.density_name = density_name def required_field_outputs(self) -> Set[field_outputs.FieldOutput]: return set(field_outputs.RGBFieldOutput, field_outputs.DensityFieldOutput) def forward( self, - field_output_dict: Dict[str, TensorType[..., "num_samples", -1]], + rgb: TensorType[..., "num_samples", 3], + density: TensorType[..., "num_samples", 1], deltas: TensorType[..., "num_samples"], ) -> TensorType[..., "out_dim"]: """Composite samples along ray and render color image @@ -62,7 +61,7 @@ class RGB(Renderer): Returns: TensorType[..., "out_dim"]: Composited RGB ray """ - delta_density = deltas * field_output_dict[self.density_name][..., 0] + delta_density = deltas * density[..., 0] alphas = 1 - torch.exp(-delta_density) transmittance = torch.cumsum(delta_density[..., :-1], dim=-1) @@ -71,12 +70,12 @@ class RGB(Renderer): weights = alphas * transmittance # [..., "num_samples"] - rgb = torch.sum(weights[..., None] * field_output_dict[self.rgb_name], dim=-2) + rgb = torch.sum(weights[..., None] * rgb, dim=-2) return rgb -class Depth(Renderer): +class DepthRenderer(Renderer): """_summary_ Args: diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 6d339404b..9d85f06db 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -1,11 +1,15 @@ """ Code to train model. """ -from omegaconf import DictConfig import torch.distributed as dist +from omegaconf import DictConfig from torch.nn.parallel import DistributedDataParallel as DDP +from torch.utils.data import DataLoader +from tqdm import tqdm -from mattport.nerf.graph import Graph +from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch +from mattport.nerf.dataset.utils import get_dataset_inputs +from mattport.nerf.graph.base import Graph from mattport.nerf.optimizer import Optimizer @@ -23,8 +27,20 @@ class Trainer: def setup_dataset(self): """_summary_""" - self.train_dataset = None - self.test_dataset = None + dataset_inputs = get_dataset_inputs(**self.config.dataset) + self.train_dataset = ImageDataset( + image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor + ) + self.train_dataloader = DataLoader( + self.train_dataset, + batch_size=self.config.dataloader.num_images_to_sample_from, + collate_fn=lambda batch: collate_batch( + batch, self.config.dataloader.num_rays_per_batch, keep_full_image=True + ), + num_workers=self.config.dataloader.num_workers, + shuffle=True, + ) + # TODO(ethan): implement test data def setup_graph(self): """_summary_""" @@ -60,10 +76,17 @@ class Trainer: def train_epoch(self): """_summary_""" - raise NotImplementedError + for _ in range(len(self.train_dataset)): + batch = next(iter(self.train_dataloader)) + self.train_iteration(batch) - def train_iteration(self): + def train_iteration(self, batch): """_summary_""" # TODO(): save checkpoints, and do logging only on rank 0 device only # if gpu == 0: - raise NotImplementedError + output_dict = self.graph(batch.indices) + + loss = self.graph.get_loss(batch, ) + + loss.backward() + # update model