# 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.""" # TODO: add spatial distortion to ngp sampler # positions = self.spatial_distortion(ray_samples.frustums.get_positions()) 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.spatial_distortion = spatial_distortion 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): # TODO: add spatial distortion to ngp sampler # positions = self.spatial_distortion(ray_samples.frustums.get_positions()) 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}