mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 10:34:29 +08:00
Merge pull request #92 from ethanweber/tancik/fields_refactor
Refactor nerf fields to separate file.
This commit is contained in:
Matthew Tancik
GitHub
commit
08c181cd5e
@ -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"""
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
89
mattport/nerf/fields/nerf_field.py
Normal file
89
mattport/nerf/fields/nerf_field.py
Normal file
@ -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}
|
||||
@ -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,
|
||||
)
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user