diff --git a/mattport/nerf/field_modules/encoding.py b/mattport/nerf/field_modules/encoding.py index 2eb0b7375..748664792 100644 --- a/mattport/nerf/field_modules/encoding.py +++ b/mattport/nerf/field_modules/encoding.py @@ -2,6 +2,7 @@ Encoding functions """ +from abc import abstractmethod import numpy as np import torch import torch.nn.functional as F @@ -25,6 +26,7 @@ class Encoding(FieldModule): raise ValueError("Input dimension should be greater than zero") self.in_dim = in_dim + @abstractmethod def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: """Encodes an input tensor. @@ -41,6 +43,16 @@ class Encoding(FieldModule): return self.encode(in_tensor) +class Identity(Encoding): + """Identity encoding (Does not modify input)""" + + def get_out_dim(self) -> int: + return self.in_dim + + def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + return in_tensor + + class ScalingAndOffset(Encoding): """Simple scaling and offet to input""" diff --git a/mattport/nerf/field_modules/mlp.py b/mattport/nerf/field_modules/mlp.py index 483da5228..532184db3 100644 --- a/mattport/nerf/field_modules/mlp.py +++ b/mattport/nerf/field_modules/mlp.py @@ -16,9 +16,9 @@ class MLP(FieldModule): def __init__( self, in_dim: int, - out_dim: int, num_layers: int, - layer_width: int = None, + layer_width: int, + out_dim: Optional[int] = None, skip_connections: Tuple[int] = (), activation: Optional[nn.Module] = nn.ReLU(), out_activation: Optional[nn.Module] = None, @@ -27,16 +27,16 @@ class MLP(FieldModule): Args: in_dim (int): Input layer dimension - out_dim (int): Ouput layer dimension num_layers (int): Number of network layers layer_width (int): Width of each MLP layer + out_dim (int, optional): Ouput layer dimension. Defaults to layer_width activation (Optional[nn.Module], optional): intermediate layer activation function. Defaults to nn.ReLU. out_activation (Optional[nn.Module], optional): output activation function. Defaults to None. """ super().__init__() self.in_dim = in_dim assert self.in_dim > 0 - self.out_dim = out_dim + self.out_dim = out_dim if out_dim is not None else layer_width self.num_layers = num_layers self.layer_width = layer_width self.skip_connections = skip_connections diff --git a/mattport/nerf/fields/nerf_field.py b/mattport/nerf/fields/nerf_field.py new file mode 100644 index 000000000..6aac19caf --- /dev/null +++ b/mattport/nerf/fields/nerf_field.py @@ -0,0 +1,89 @@ +""" Classic NeRF field""" + + +from typing import Tuple + +import torch +from mattport.nerf.field_modules.encoding import Encoding, Identity +from mattport.nerf.field_modules.field_heads import DensityFieldHead, RGBFieldHead +from mattport.nerf.field_modules.mlp import MLP +from mattport.nerf.fields.base import Field +from mattport.structures.rays import PointSamples +from mattport.utils.misc import is_not_none + + +class NeRFField(Field): + """NeRF Field""" + + def __init__( + self, + position_encoding: Encoding = Identity(in_dim=3), + direction_encoding: Encoding = Identity(in_dim=3), + base_mlp_num_layers: int = 8, + base_mlp_layer_width: int = 256, + head_mlp_num_layers: int = 2, + head_mlp_layer_width: int = 128, + skip_connections: Tuple = (4,), + ) -> None: + """Create boilerplate NeRF field. + + Args: + position_encoding (Encoding, optional): Position encoder. Defaults to Identity(in_dim=3). + direction_encoding (Encoding, optional): Direction encoder. Defaults to Identity(in_dim=3). + base_mlp_num_layers (int, optional): Number of layers for base MLP. Defaults to 8. + base_mlp_layer_width (int, optional): Width of base MLP layers. Defaults to 256. + head_mlp_num_layers (int, optional): Number of layer for ourput head MLP. Defaults to 2. + head_mlp_layer_width (int, optional): Width of output head MLP layers. Defaults to 128. + skip_connections (Tuple, optional): Where to add skip connection in base MLP. Defaults to (4,). + """ + super().__init__() + self.position_encoding = position_encoding + self.direction_encoding = direction_encoding + + self.mlp_base = MLP( + in_dim=self.position_encoding.get_out_dim(), + num_layers=base_mlp_num_layers, + layer_width=base_mlp_layer_width, + skip_connections=skip_connections, + ) + + self.mlp_head = MLP( + in_dim=self.mlp_base.get_out_dim() + self.direction_encoding.get_out_dim(), + num_layers=head_mlp_num_layers, + layer_width=head_mlp_layer_width, + ) + + self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim()) + self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_head.get_out_dim()) + + def get_density(self, point_samples: PointSamples): + """Computes and returns the densities.""" + positions = point_samples.positions + valid_mask = point_samples.valid_mask + if not is_not_none(valid_mask): + valid_mask = torch.ones_like(positions[..., 0]).bool() + # placeholders for values to return + density = torch.zeros(*valid_mask.shape, 1, dtype=torch.float32, device=positions.device) + base_mlp_out = torch.zeros( + *valid_mask.shape, self.mlp_base.out_dim, dtype=torch.float32, device=positions.device + ) + if not valid_mask.any(): # empty mask + return density, base_mlp_out + + encoded_xyz = self.position_encoding(positions[valid_mask]) + base_mlp_out[valid_mask] = self.mlp_base(encoded_xyz) + density[valid_mask] = self.field_output_density(base_mlp_out[valid_mask]) + return density, base_mlp_out + + def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None): + directions = point_samples.directions + if not is_not_none(valid_mask): + valid_mask = torch.ones_like(directions[..., 0]).bool() + # placeholders for values to return + rgb = torch.zeros(*valid_mask.shape, 3, dtype=torch.float32, device=directions.device) + if not valid_mask.any(): # empty mask + return {"rgb": rgb} + encoded_dir = self.direction_encoding(directions[valid_mask]) + rgb_mlp_out = self.mlp_head(torch.cat([encoded_dir, density_embedding[valid_mask]], dim=-1)) + rgb[valid_mask] = self.field_output_rgb(rgb_mlp_out) + return {"rgb": rgb} diff --git a/mattport/nerf/graph/instant_ngp.py b/mattport/nerf/graph/instant_ngp.py index 7a9bcd340..908d8df4b 100644 --- a/mattport/nerf/graph/instant_ngp.py +++ b/mattport/nerf/graph/instant_ngp.py @@ -2,35 +2,10 @@ Implementation of Instant NGP. """ -from torch import nn - -from mattport.nerf.field_modules.encoding import HashEncoding, NeRFEncoding -from mattport.nerf.field_modules.mlp import MLP +from mattport.nerf.field_modules.encoding import HashEncoding, SHEncoding from mattport.nerf.graph.vanilla_nerf import NeRFField, NeRFGraph -class NGPField(NeRFField): - """NeRF module""" - - def __init__(self, num_layers=3, layer_width=64) -> None: - super().__init__(num_layers=num_layers, layer_width=layer_width) - - def build_encodings(self): - self.encoding_xyz = HashEncoding() - self.encoding_dir = NeRFEncoding( - in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True - ) - - def build_mlp_base(self): - self.mlp_base = MLP( - in_dim=self.encoding_xyz.get_out_dim(), - out_dim=self.layer_width, - num_layers=self.num_layers, - layer_width=self.layer_width, - activation=nn.ReLU(), - ) - - class NGPGraph(NeRFGraph): """NeRF-W graph""" @@ -39,5 +14,23 @@ class NGPGraph(NeRFGraph): def populate_fields(self): """Set the fields.""" - self.field_coarse = NGPField() - self.field_fine = NGPField() + + position_encoding = HashEncoding() + direction_encoding = SHEncoding() + + self.field_coarse = NeRFField( + position_encoding=position_encoding, + direction_encoding=direction_encoding, + base_mlp_num_layers=3, + base_mlp_layer_width=64, + head_mlp_num_layers=2, + head_mlp_layer_width=32, + ) + self.field_fine = NeRFField( + position_encoding=position_encoding, + direction_encoding=direction_encoding, + base_mlp_num_layers=3, + base_mlp_layer_width=64, + head_mlp_num_layers=2, + head_mlp_layer_width=32, + ) diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 2bb77c109..4f2275957 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -3,111 +3,23 @@ Implementation of vanilla nerf. """ -from typing import Dict, List, Tuple +from typing import Dict, List import torch -from torch import nn from torch.nn import Parameter from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity - 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.fields.base import Field + +from mattport.nerf.fields.nerf_field import NeRFField from mattport.nerf.graph.base import Graph from mattport.nerf.loss import MSELoss from mattport.nerf.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer from mattport.nerf.sampler import PDFSampler, UniformSampler from mattport.structures import colors -from mattport.structures.rays import PointSamples, RayBundle +from mattport.structures.rays import RayBundle from mattport.nerf.colliders import NearFarCollider from mattport.utils import visualization, writer -from mattport.utils.misc import is_not_none - - -class NeRFField(Field): - """NeRF module""" - - OUTPUT_RGB = "rgb" - - def __init__(self, num_layers=8, layer_width=256, skip_connections: Tuple = (4,)) -> None: - super().__init__() - self.num_layers = num_layers - self.layer_width = layer_width - self.skip_connections = skip_connections - - self.build_encodings() - self.build_mlp_base() - self.build_mlp_rgb() - self.build_heads() - - def build_encodings(self): - """Build the encodings.""" - self.encoding_xyz = NeRFEncoding( - in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True - ) - self.encoding_dir = NeRFEncoding( - in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True - ) - - def build_mlp_base(self): - """Build the MLP base.""" - self.mlp_base = MLP( - in_dim=self.encoding_xyz.get_out_dim(), - out_dim=self.layer_width, - num_layers=self.num_layers, - layer_width=self.layer_width, - skip_connections=self.skip_connections, - activation=nn.ReLU(), - ) - - def build_mlp_rgb(self): - """Build the MLP for RGB.""" - self.mlp_rgb = MLP( - in_dim=self.mlp_base.get_out_dim() + self.encoding_dir.get_out_dim(), - out_dim=self.layer_width // 2, - num_layers=2, - layer_width=self.layer_width // 2, - activation=nn.ReLU(), - ) - - def build_heads(self): - """Build the heads.""" - 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 get_density(self, point_samples: PointSamples): - """Computes and returns the densities.""" - positions = point_samples.positions - valid_mask = point_samples.valid_mask - if not is_not_none(valid_mask): - valid_mask = torch.ones_like(positions[..., 0]).bool() - # placeholders for values to return - density = torch.zeros(*valid_mask.shape, 1, dtype=torch.float32, device=positions.device) - base_mlp_out = torch.zeros( - *valid_mask.shape, self.mlp_base.out_dim, dtype=torch.float32, device=positions.device - ) - if not valid_mask.any(): # empty mask - return density, base_mlp_out - - encoded_xyz = self.encoding_xyz(positions[valid_mask]) - base_mlp_out[valid_mask] = self.mlp_base(encoded_xyz) - density[valid_mask] = self.field_output_density(base_mlp_out[valid_mask]) - return density, base_mlp_out - - def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None): - directions = point_samples.directions - if not is_not_none(valid_mask): - valid_mask = torch.ones_like(directions[..., 0]).bool() - # placeholders for values to return - rgb = torch.zeros(*valid_mask.shape, 3, dtype=torch.float32, device=directions.device) - if not valid_mask.any(): # empty mask - return {"rgb": rgb} - encoded_dir = self.encoding_dir(directions[valid_mask]) - rgb_mlp_out = self.mlp_rgb(torch.cat([encoded_dir, density_embedding[valid_mask]], dim=-1)) - rgb[valid_mask] = self.field_output_rgb(rgb_mlp_out) - return {"rgb": rgb} class NeRFGraph(Graph): @@ -136,8 +48,16 @@ class NeRFGraph(Graph): def populate_fields(self): """Set the fields.""" - self.field_coarse = NeRFField() - self.field_fine = NeRFField() + + position_encoding = NeRFEncoding( + in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True + ) + direction_encoding = NeRFEncoding( + in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True + ) + + self.field_coarse = NeRFField(position_encoding=position_encoding, direction_encoding=direction_encoding) + self.field_fine = NeRFField(position_encoding=position_encoding, direction_encoding=direction_encoding) def populate_misc_modules(self): # samplers