mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
282 lines
10 KiB
Python
282 lines
10 KiB
Python
# Copyright 2022 The Plenoptix Team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""
|
|
Field for compound nerf model, adds scene contraction and image embeddings to instant ngp
|
|
"""
|
|
|
|
|
|
from typing import Dict, Optional, Tuple
|
|
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn.parameter import Parameter
|
|
from torchtyping import TensorType
|
|
|
|
from nerfactory.cameras.rays import RaySamples
|
|
from nerfactory.datamanagers.structs import SceneBounds
|
|
from nerfactory.fields.base import Field
|
|
from nerfactory.fields.modules.embedding import Embedding
|
|
from nerfactory.fields.modules.encoding import Encoding, HashEncoding, SHEncoding
|
|
from nerfactory.fields.modules.field_heads import (
|
|
DensityFieldHead,
|
|
FieldHead,
|
|
FieldHeadNames,
|
|
RGBFieldHead,
|
|
)
|
|
from nerfactory.fields.modules.mlp import MLP
|
|
from nerfactory.fields.modules.spatial_distortions import (
|
|
SceneContraction,
|
|
SpatialDistortion,
|
|
)
|
|
from nerfactory.utils.activations import trunc_exp
|
|
|
|
try:
|
|
import tinycudann as tcnn
|
|
except ImportError:
|
|
# tinycudann module doesn't exist
|
|
pass
|
|
|
|
|
|
def get_normalized_directions(directions: TensorType["bs":..., 3]):
|
|
"""SH encoding must be in the range [0, 1]
|
|
|
|
Args:
|
|
directions: batch of directions
|
|
"""
|
|
return (directions + 1.0) / 2.0
|
|
|
|
|
|
class TCNNCompoundField(Field):
|
|
"""Compound Field that uses TCNN
|
|
|
|
Args:
|
|
aabb: parameters of scene aabb bounds
|
|
num_layers: number of hidden layers
|
|
hidden_dim: dimension of hidden layers
|
|
geo_feat_dim: output geo feat dimensions
|
|
num_layers_color: number of hidden layers for color network
|
|
hidden_dim_color: dimension of hidden layers for color network
|
|
appearance_embedding_dim: dimension of appearance embedding
|
|
spatial_distortion: spatial distortion to apply to the scene
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
aabb,
|
|
num_images: int,
|
|
num_layers: int = 2,
|
|
hidden_dim: int = 64,
|
|
geo_feat_dim: int = 15,
|
|
num_layers_color: int = 3,
|
|
hidden_dim_color: int = 64,
|
|
appearance_embedding_dim: int = 32,
|
|
spatial_distortion: SpatialDistortion = SceneContraction(),
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self.aabb = Parameter(aabb, requires_grad=False)
|
|
self.geo_feat_dim = geo_feat_dim
|
|
|
|
self.spatial_distortion = spatial_distortion
|
|
self.num_images = num_images
|
|
self.appearance_embedding_dim = appearance_embedding_dim
|
|
self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim)
|
|
|
|
# TODO: set this properly based on the aabb
|
|
per_level_scale = 1.4472692012786865
|
|
|
|
self.direction_encoding = tcnn.Encoding(
|
|
n_input_dims=3,
|
|
encoding_config={
|
|
"otype": "SphericalHarmonics",
|
|
"degree": 4,
|
|
},
|
|
)
|
|
|
|
self.mlp_base = tcnn.NetworkWithInputEncoding(
|
|
n_input_dims=3,
|
|
n_output_dims=1 + self.geo_feat_dim,
|
|
encoding_config={
|
|
"otype": "HashGrid",
|
|
"n_levels": 16,
|
|
"n_features_per_level": 2,
|
|
"log2_hashmap_size": 19,
|
|
"base_resolution": 16,
|
|
"per_level_scale": per_level_scale,
|
|
},
|
|
network_config={
|
|
"otype": "FullyFusedMLP",
|
|
"activation": "ReLU",
|
|
"output_activation": "None",
|
|
"n_neurons": hidden_dim,
|
|
"n_hidden_layers": num_layers - 1,
|
|
},
|
|
)
|
|
|
|
self.mlp_head = tcnn.Network(
|
|
n_input_dims=self.direction_encoding.n_output_dims + self.geo_feat_dim + self.appearance_embedding_dim,
|
|
n_output_dims=3,
|
|
network_config={
|
|
"otype": "FullyFusedMLP",
|
|
"activation": "ReLU",
|
|
"output_activation": "Sigmoid",
|
|
"n_neurons": hidden_dim_color,
|
|
"n_hidden_layers": num_layers_color - 1,
|
|
},
|
|
)
|
|
|
|
def get_density(self, ray_samples: RaySamples):
|
|
"""Computes and returns the densities."""
|
|
if self.spatial_distortion is not None:
|
|
positions = ray_samples.frustums.get_positions()
|
|
positions = self.spatial_distortion(positions)
|
|
positions = (positions + 2.0) / 4.0
|
|
else:
|
|
positions = SceneBounds.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
|
|
positions_flat = positions.view(-1, 3)
|
|
h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1)
|
|
density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1)
|
|
|
|
# Rectifying the density with an exponential is much more stable than a ReLU or
|
|
# softplus, because it enables high post-activation (float32) density outputs
|
|
# from smaller internal (float16) parameters.
|
|
density = trunc_exp(density_before_activation.to(positions))
|
|
return density, base_mlp_out
|
|
|
|
def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
|
|
if ray_samples.camera_indices is None:
|
|
raise AttributeError("Camera indices are not provided.")
|
|
camera_indices = ray_samples.camera_indices.squeeze()
|
|
directions = get_normalized_directions(ray_samples.frustums.directions)
|
|
directions_flat = directions.view(-1, 3)
|
|
d = self.direction_encoding(directions_flat)
|
|
|
|
if self.training:
|
|
embedded_appearance = self.embedding_appearance(camera_indices)
|
|
else:
|
|
embedded_appearance = torch.zeros(
|
|
(*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
|
|
)
|
|
|
|
if density_embedding is None:
|
|
positions = SceneBounds.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
|
|
h = torch.cat([d, positions.view(-1, 3), embedded_appearance], dim=-1)
|
|
else:
|
|
h = torch.cat(
|
|
[
|
|
d,
|
|
density_embedding.view(-1, self.geo_feat_dim),
|
|
embedded_appearance.view(-1, self.appearance_embedding_dim),
|
|
],
|
|
dim=-1,
|
|
)
|
|
rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions)
|
|
|
|
return {FieldHeadNames.RGB: rgb}
|
|
|
|
|
|
class TorchCompoundField(Field):
|
|
"""
|
|
PyTorch implementation of the compound field.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
aabb,
|
|
num_images: int,
|
|
position_encoding: Encoding = HashEncoding(),
|
|
direction_encoding: Encoding = SHEncoding(),
|
|
base_mlp_num_layers: int = 3,
|
|
base_mlp_layer_width: int = 64,
|
|
head_mlp_num_layers: int = 2,
|
|
head_mlp_layer_width: int = 32,
|
|
appearance_embedding_dim: int = 40,
|
|
skip_connections: Tuple = (4,),
|
|
field_heads: Tuple[FieldHead] = (RGBFieldHead(),),
|
|
spatial_distortion: SpatialDistortion = SceneContraction(),
|
|
) -> None:
|
|
super().__init__()
|
|
self.aabb = Parameter(aabb, requires_grad=False)
|
|
self.spatial_distortion = spatial_distortion
|
|
self.num_images = num_images
|
|
self.appearance_embedding_dim = appearance_embedding_dim
|
|
self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim)
|
|
|
|
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() + self.appearance_embedding_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_heads = nn.ModuleList(field_heads)
|
|
for field_head in self.field_heads:
|
|
field_head.set_in_dim(self.mlp_head.get_out_dim()) # type: ignore
|
|
|
|
def get_density(self, ray_samples: RaySamples):
|
|
if self.spatial_distortion is not None:
|
|
positions = ray_samples.frustums.get_positions()
|
|
positions = self.spatial_distortion(positions)
|
|
else:
|
|
positions = ray_samples.frustums.get_positions()
|
|
encoded_xyz = self.position_encoding(positions)
|
|
base_mlp_out = self.mlp_base(encoded_xyz)
|
|
density = self.field_output_density(base_mlp_out)
|
|
return density, base_mlp_out
|
|
|
|
def get_outputs(
|
|
self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None
|
|
) -> Dict[FieldHeadNames, TensorType]:
|
|
|
|
if ray_samples.camera_indices is None:
|
|
raise AttributeError("Camera indices are not provided.")
|
|
camera_indices = ray_samples.camera_indices.squeeze()
|
|
if self.training:
|
|
embedded_appearance = self.embedding_appearance(camera_indices)
|
|
else:
|
|
embedded_appearance = torch.zeros(
|
|
(*ray_samples.frustums.directions.shape[:-1], self.appearance_embedding_dim),
|
|
device=ray_samples.frustums.directions.device,
|
|
)
|
|
|
|
outputs = {}
|
|
for field_head in self.field_heads:
|
|
encoded_dir = self.direction_encoding(ray_samples.frustums.directions)
|
|
mlp_out = self.mlp_head(
|
|
torch.cat(
|
|
[
|
|
encoded_dir,
|
|
density_embedding, # type:ignore
|
|
embedded_appearance.view(-1, self.appearance_embedding_dim),
|
|
],
|
|
dim=-1, # type:ignore
|
|
)
|
|
)
|
|
outputs[field_head.field_head_name] = field_head(mlp_out)
|
|
return outputs
|
|
|
|
|
|
field_implementation_to_class = {"tcnn": TCNNCompoundField, "torch": TorchCompoundField}
|