From d6bb8bd8e5d278634b5fa73c18ae8bcc0daad986 Mon Sep 17 00:00:00 2001 From: Brent Yi Date: Sat, 1 Jul 2023 11:54:27 -0700 Subject: [PATCH] Fix `use_appearance_embedding` for instant-ngp (#1999) * Attempt fixes for instant-ngp * Run `ns-dev-test` * Remove NGP field implementation * Test nerfacto field instead of ngp field * Black --- docs/reference/api/fields.rst | 39 ++++- nerfstudio/fields/instant_ngp_field.py | 198 ------------------------- nerfstudio/fields/nerfw_field.py | 6 +- nerfstudio/models/instant_ngp.py | 1 + tests/field_components/test_fields.py | 25 +++- 5 files changed, 58 insertions(+), 211 deletions(-) delete mode 100644 nerfstudio/fields/instant_ngp_field.py diff --git a/docs/reference/api/fields.rst b/docs/reference/api/fields.rst index 4dea6f8e9..33bde75c1 100644 --- a/docs/reference/api/fields.rst +++ b/docs/reference/api/fields.rst @@ -10,10 +10,45 @@ Base :members: :show-inheritance: -Instant NGP +Density ---------------- -.. automodule:: nerfstudio.fields.instant_ngp_field +.. automodule:: nerfstudio.fields.density_fields + :members: + :show-inheritance: + +Nerfacto +---------------- + +.. automodule:: nerfstudio.fields.nerfacto_field + :members: + :show-inheritance: + +Nerf-W +---------------- + +.. automodule:: nerfstudio.fields.nerfw_field + :members: + :show-inheritance: + +SDF +---------------- + +.. automodule:: nerfstudio.fields.sdf_field + :members: + :show-inheritance: + +Semantic NeRF +---------------- + +.. automodule:: nerfstudio.fields.semantic_nerf_field + :members: + :show-inheritance: + +TensoRF +---------------- + +.. automodule:: nerfstudio.fields.tensorf_field :members: :show-inheritance: diff --git a/nerfstudio/fields/instant_ngp_field.py b/nerfstudio/fields/instant_ngp_field.py deleted file mode 100644 index 63b01eff6..000000000 --- a/nerfstudio/fields/instant_ngp_field.py +++ /dev/null @@ -1,198 +0,0 @@ -# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. 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. - -""" -Instant-NGP field implementations using tiny-cuda-nn, torch, .... -""" - - -from typing import Dict, Optional, Tuple - -import numpy as np -import torch -from jaxtyping import Float -from torch import Tensor -from torch.nn.parameter import Parameter - -from nerfstudio.cameras.rays import RaySamples -from nerfstudio.data.scene_box import SceneBox -from nerfstudio.field_components.activations import trunc_exp -from nerfstudio.field_components.embedding import Embedding -from nerfstudio.field_components.field_heads import FieldHeadNames -from nerfstudio.field_components.spatial_distortions import ( - SceneContraction, - SpatialDistortion, -) -from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn - -try: - import tinycudann as tcnn -except ModuleNotFoundError: - # tinycudann module doesn't exist - pass - - -class TCNNInstantNGPField(Field): - """TCNN implementation of the Instant-NGP field. - - 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 - use_appearance_embedding: whether to use appearance embedding - num_images: number of images, required if use_appearance_embedding is True - appearance_embedding_dim: dimension of appearance embedding - contraction_type: type of contraction - num_levels: number of levels of the hashmap for the base mlp - log2_hashmap_size: size of the hashmap for the base mlp - max_res: maximum resolution of the hashmap for the base mlp - mode: which implementation to use for the field - """ - - def __init__( - self, - aabb: Tensor, - num_layers: int = 2, - hidden_dim: int = 64, - geo_feat_dim: int = 15, - num_layers_color: int = 3, - hidden_dim_color: int = 64, - use_appearance_embedding: Optional[bool] = False, - num_images: Optional[int] = None, - appearance_embedding_dim: int = 32, - num_levels: int = 16, - log2_hashmap_size: int = 19, - max_res: int = 2048, - spatial_distortion: Optional[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.use_appearance_embedding = use_appearance_embedding - if use_appearance_embedding: - assert num_images is not None - self.appearance_embedding_dim = appearance_embedding_dim - self.appearance_embedding = Embedding(num_images, appearance_embedding_dim) - - # TODO: set this properly based on the aabb - base_res: int = 16 - features_per_level: int = 2 - growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1)) - - 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": num_levels, - "n_features_per_level": features_per_level, - "log2_hashmap_size": log2_hashmap_size, - "base_resolution": base_res, - "per_level_scale": growth_factor, - }, - network_config={ - "otype": "FullyFusedMLP", - "activation": "ReLU", - "output_activation": "None", - "n_neurons": hidden_dim, - "n_hidden_layers": num_layers - 1, - }, - ) - - in_dim = self.direction_encoding.n_output_dims + self.geo_feat_dim - if self.use_appearance_embedding: - in_dim += self.appearance_embedding_dim - self.mlp_head = tcnn.Network( - n_input_dims=in_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) -> Tuple[Tensor, Tensor]: - 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 = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) - # Make sure the tcnn gets inputs between 0 and 1. - selector = ((positions > 0.0) & (positions < 1.0)).all(dim=-1) - positions = positions * selector[..., None] - 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)) - density = density * selector[..., None] - return density, base_mlp_out - - def get_outputs( - self, ray_samples: RaySamples, density_embedding: Optional[Tensor] = None - ) -> Dict[FieldHeadNames, Tensor]: - assert density_embedding is not None - directions = shift_directions_for_tcnn(ray_samples.frustums.directions) - directions_flat = directions.view(-1, 3) - d = self.direction_encoding(directions_flat) - - h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) - - if self.use_appearance_embedding: - 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.appearance_embedding(camera_indices) - else: - embedded_appearance = torch.zeros( - (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device - ) - h = torch.cat([h, 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} - - def get_opacity(self, positions: Float[Tensor, "*bs 3"], step_size) -> Float[Tensor, "*bs 1"]: - """Returns the opacity for a position. Used primarily by the occupancy grid. - - Args: - positions: the positions to evaluate the opacity at. - step_size: the step size to use for the opacity evaluation. - """ - density = self.density_fn(positions) - opacity = density * step_size - return opacity diff --git a/nerfstudio/fields/nerfw_field.py b/nerfstudio/fields/nerfw_field.py index be71f15bb..4f385a371 100644 --- a/nerfstudio/fields/nerfw_field.py +++ b/nerfstudio/fields/nerfw_field.py @@ -123,11 +123,11 @@ class VanillaNerfWField(Field): """Returns the outputs of the NeRF-W field. Args: - ray_samples (RaySamples): Ray samples. - density_embedding (Float[Tensor, "*batch embedding_size"], optional): Density embedding. + ray_samples: Ray samples. + density_embedding: Density embedding. Returns: - Dict[FieldHeadNames, Tensor]: Outputs of the NeRF-W field. + Outputs of the NeRF-W field. """ outputs = {} encoded_dir = self.direction_encoding(ray_samples.frustums.directions) diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index e8ac52b96..95f059b60 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/nerfstudio/models/instant_ngp.py @@ -110,6 +110,7 @@ class NGPModel(Model): self.field = NerfactoField( aabb=self.scene_box.aabb, + appearance_embedding_dim=0 if self.config.use_appearance_embedding else 32, num_images=self.num_train_data, log2_hashmap_size=self.config.log2_hashmap_size, max_res=self.config.max_res, diff --git a/tests/field_components/test_fields.py b/tests/field_components/test_fields.py index 9e5611fcc..74d7d6c3a 100644 --- a/tests/field_components/test_fields.py +++ b/tests/field_components/test_fields.py @@ -4,11 +4,11 @@ Test the fields import torch from nerfstudio.cameras.rays import Frustums, RaySamples -from nerfstudio.fields.instant_ngp_field import TCNNInstantNGPField +from nerfstudio.fields.nerfacto_field import NerfactoField -def test_tcnn_instant_ngp_field(): - """Test the tiny-cuda-nn field""" +def test_nerfacto_field(): + """Test the Nerfacto field""" try: import tinycudann as tcnn # noqa: F401 except ModuleNotFoundError as e: @@ -18,9 +18,11 @@ def test_tcnn_instant_ngp_field(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") aabb_scale = 1.0 aabb = torch.tensor( - [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32 - ).to(device) - field = TCNNInstantNGPField(aabb) + [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], + dtype=torch.float32, + device=device, + ) + field = NerfactoField(aabb, num_images=1).to(device) num_rays = 1024 num_samples = 256 positions = torch.rand((num_rays, num_samples, 3), dtype=torch.float32, device=device) @@ -32,9 +34,16 @@ def test_tcnn_instant_ngp_field(): ends=torch.zeros((*directions.shape[:-1], 1), device=device), pixel_area=torch.ones((*directions.shape[:-1], 1), device=device), ) - ray_samples = RaySamples(frustums=frustums) + ray_samples = RaySamples( + frustums=frustums, + camera_indices=torch.zeros( + (num_rays, 1, 1), + device=device, + dtype=torch.int32, + ), + ) field.forward(ray_samples) if __name__ == "__main__": - test_tcnn_instant_ngp_field() + test_nerfacto_field()