tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 18:44:21 +08:00
Code Activity

Fix use_appearance_embedding for instant-ngp (#1999)

Browse Source 
* Attempt fixes for instant-ngp

* Run `ns-dev-test`

* Remove NGP field implementation

* Test nerfacto field instead of ngp field

* Black
This commit is contained in:
Brent Yi authored and GitHub committed 2023-07-01 18:54:27 +00:00
1 parent 5c4b0a1dd6
commit d6bb8bd8e5
5 files changed
+58 -211
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/reference/api/fields.rst
+37 -2
View File
@@ -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:
nerfstudio/fields/instant_ngp_field.py
-198
View File
@@ -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
nerfstudio/fields/nerfw_field.py
+3 -3
View File
@@ -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)
nerfstudio/models/instant_ngp.py
+1
View File
@@ -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,
tests/field_components/test_fields.py
+17 -8
View File
@@ -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()