From e73692e599ca7047703874ec2d5bf66fe1f6d703 Mon Sep 17 00:00:00 2001 From: LSong <12553253+lsongx@users.noreply.github.com> Date: Tue, 31 Jan 2023 11:56:28 -0800 Subject: [PATCH] nerfplayer with ngp and nerfacto (#1277) * nerfplayer with ngp and nerfacto * update comments * update cpp extension build * avoid Windows repeated CUDA compile --------- Co-authored-by: machenmusik --- nerfstudio/configs/method_configs.py | 62 +++ nerfstudio/field_components/cuda/_backend.py | 3 +- nerfstudio/field_components/temporal_grid.py | 2 +- .../fields/nerfplayer_nerfacto_field.py | 415 ++++++++++++++++++ nerfstudio/fields/nerfplayer_ngp_field.py | 236 ++++++++++ nerfstudio/model_components/ray_samplers.py | 2 + nerfstudio/models/nerfplayer_nerfacto.py | 251 +++++++++++ nerfstudio/models/nerfplayer_ngp.py | 241 ++++++++++ nerfstudio/viewer/server/viewer_utils.py | 39 +- tests/test_train.py | 2 + 10 files changed, 1238 insertions(+), 15 deletions(-) create mode 100644 nerfstudio/fields/nerfplayer_nerfacto_field.py create mode 100644 nerfstudio/fields/nerfplayer_ngp_field.py create mode 100644 nerfstudio/models/nerfplayer_nerfacto.py create mode 100644 nerfstudio/models/nerfplayer_ngp.py diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 30682a242..947dad880 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -33,6 +33,7 @@ from nerfstudio.data.datamanagers.variable_res_datamanager import ( ) from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig +from nerfstudio.data.dataparsers.dycheck_dataparser import DycheckDataParserConfig from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig from nerfstudio.data.dataparsers.instant_ngp_dataparser import ( InstantNGPDataParserConfig, @@ -49,6 +50,8 @@ from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig from nerfstudio.models.instant_ngp import InstantNGPModelConfig from nerfstudio.models.mipnerf import MipNerfModel from nerfstudio.models.nerfacto import NerfactoModelConfig +from nerfstudio.models.nerfplayer_nerfacto import NerfplayerNerfactoModelConfig +from nerfstudio.models.nerfplayer_ngp import NerfplayerNGPModelConfig from nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig from nerfstudio.models.tensorf import TensoRFModelConfig from nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig @@ -67,6 +70,8 @@ descriptions = { "tensorf": "tensorf", "dnerf": "Dynamic-NeRF model. (slow)", "phototourism": "Uses the Phototourism data.", + "nerfplayer-nerfacto": "NeRFPlayer with nerfacto backbone.", + "nerfplayer-ngp": "NeRFPlayer with InstantNGP backbone.", } method_configs["nerfacto"] = TrainerConfig( @@ -325,6 +330,63 @@ method_configs["phototourism"] = TrainerConfig( vis="viewer", ) +method_configs["nerfplayer-nerfacto"] = TrainerConfig( + method_name="nerfplayer-nerfacto", + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=30000, + mixed_precision=True, + pipeline=VanillaPipelineConfig( + datamanager=DepthDataManagerConfig( + dataparser=DycheckDataParserConfig(), + train_num_rays_per_batch=4096, + eval_num_rays_per_batch=4096, + camera_optimizer=CameraOptimizerConfig( + mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) + ), + ), + model=NerfplayerNerfactoModelConfig(eval_num_rays_per_chunk=1 << 15), + ), + optimizers={ + "proposal_networks": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + }, + "fields": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + }, + }, + viewer=ViewerConfig(num_rays_per_chunk=1 << 15), + vis="viewer", +) + +method_configs["nerfplayer-ngp"] = TrainerConfig( + method_name="nerfplayer-ngp", + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=30000, + mixed_precision=True, + pipeline=DynamicBatchPipelineConfig( + datamanager=DepthDataManagerConfig(dataparser=DycheckDataParserConfig(), train_num_rays_per_batch=8192), + model=NerfplayerNGPModelConfig( + eval_num_rays_per_chunk=8192, + contraction_type=ContractionType.AABB, + render_step_size=0.001, + max_num_samples_per_ray=48, + near_plane=0.01, + ), + ), + optimizers={ + "fields": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + } + }, + viewer=ViewerConfig(num_rays_per_chunk=64000), + vis="viewer", +) + AnnotatedBaseConfigUnion = tyro.conf.SuppressFixed[ # Don't show unparseable (fixed) arguments in helptext. tyro.conf.FlagConversionOff[ tyro.extras.subcommand_type_from_defaults(defaults=method_configs, descriptions=descriptions) diff --git a/nerfstudio/field_components/cuda/_backend.py b/nerfstudio/field_components/cuda/_backend.py index d4269b9ee..17732a196 100644 --- a/nerfstudio/field_components/cuda/_backend.py +++ b/nerfstudio/field_components/cuda/_backend.py @@ -27,9 +27,10 @@ BUILD_DIR = _get_build_directory(NAME, verbose=False) _C = None if cuda_toolkit_available(): - if os.path.exists(os.path.join(BUILD_DIR, f"{NAME}.so")): + if os.listdir(BUILD_DIR) != []: # If the build exists, we assume the extension has been built # and we can load it. + print("nerfstudio field components: CUDA set up, loading (should be quick)") _C = load( name=NAME, sources=glob.glob(os.path.join(PATH, "csrc/*.cu")), diff --git a/nerfstudio/field_components/temporal_grid.py b/nerfstudio/field_components/temporal_grid.py index c4f5e20c5..6bccc110a 100644 --- a/nerfstudio/field_components/temporal_grid.py +++ b/nerfstudio/field_components/temporal_grid.py @@ -182,7 +182,7 @@ class TemporalGridEncoder(nn.Module): input_dim: int = 3, num_levels: int = 16, level_dim: int = 2, - per_level_scale: int = 2, + per_level_scale: float = 2.0, base_resolution: int = 16, log2_hashmap_size: int = 19, desired_resolution: Optional[int] = None, diff --git a/nerfstudio/fields/nerfplayer_nerfacto_field.py b/nerfstudio/fields/nerfplayer_nerfacto_field.py new file mode 100644 index 000000000..e87df3b03 --- /dev/null +++ b/nerfstudio/fields/nerfplayer_nerfacto_field.py @@ -0,0 +1,415 @@ +# Copyright 2022 The Nerfstudio 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 implementations for NeRFPlayer (https://arxiv.org/abs/2210.15947) implementation with nerfacto backbone +""" + + +from typing import Optional + +import numpy as np +import torch +from torch.nn.parameter import Parameter +from torchtyping import TensorType + +from nerfstudio.cameras.rays import Frustums, 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, + PredNormalsFieldHead, + SemanticFieldHead, + TransientDensityFieldHead, + TransientRGBFieldHead, + UncertaintyFieldHead, +) +from nerfstudio.field_components.spatial_distortions import SpatialDistortion +from nerfstudio.field_components.temporal_grid import TemporalGridEncoder +from nerfstudio.fields.base_field import Field + +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 TemporalHashMLPDensityField(Field): + """A lightweight temporal density field module. + + Args: + aabb: Parameters of scene aabb bounds + temporal_dim: Hashing grid parameter. A higher temporal dim means a higher temporal frequency. + num_layers: Number of hidden layers + hidden_dim: Dimension of hidden layers + spatial_distortion: Spatial distortion module + num_levels: Hashing grid parameter. Used for initialize TemporalGridEncoder class. + max_res: Hashing grid parameter. Used for initialize TemporalGridEncoder class. + base_res: Hashing grid parameter. Used for initialize TemporalGridEncoder class. + log2_hashmap_size: Hashing grid parameter. Used for initialize TemporalGridEncoder class. + features_per_level: Hashing grid parameter. Used for initialize TemporalGridEncoder class. + """ + + def __init__( + self, + aabb, + temporal_dim: int = 64, + num_layers: int = 2, + hidden_dim: int = 64, + spatial_distortion: Optional[SpatialDistortion] = None, + num_levels=8, + max_res=1024, + base_res=16, + log2_hashmap_size=18, + features_per_level=2, + ) -> None: + super().__init__() + # from .temporal_grid import test; test() # DEBUG + self.aabb = Parameter(aabb, requires_grad=False) + self.spatial_distortion = spatial_distortion + growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1)) + + self.encoding = TemporalGridEncoder( + input_dim=3, + temporal_dim=temporal_dim, + num_levels=num_levels, + level_dim=features_per_level, + per_level_scale=growth_factor, + base_resolution=base_res, + log2_hashmap_size=log2_hashmap_size, + ) + self.linear = tcnn.Network( + n_input_dims=num_levels * features_per_level, + n_output_dims=1, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": hidden_dim, + "n_hidden_layers": num_layers - 1, + }, + ) + + # pylint: disable=arguments-differ + def density_fn(self, positions: TensorType["bs":..., 3], times: TensorType["bs", 1]) -> TensorType["bs":..., 1]: + """Returns only the density. Used primarily with the density grid. + + Args: + positions: the origin of the samples/frustums + times: the time of rays + """ + if len(positions.shape) == 3 and len(times.shape) == 2: + # position is [ray, sample, 3]; times is [ray, 1] + times = times[:, None] # RaySamples can handle the shape + # Need to figure out a better way to descibe positions with a ray. + ray_samples = RaySamples( + frustums=Frustums( + origins=positions, + directions=torch.ones_like(positions), + starts=torch.zeros_like(positions[..., :1]), + ends=torch.zeros_like(positions[..., :1]), + pixel_area=torch.ones_like(positions[..., :1]), + ), + times=times, + ) + density, _ = self.get_density(ray_samples) + return density + + def get_density(self, ray_samples: RaySamples): + if self.spatial_distortion is not None: + positions = self.spatial_distortion(ray_samples.frustums.get_positions()) + positions = (positions + 2.0) / 4.0 + else: + positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + positions_flat = positions.view(-1, 3) + time_flat = ray_samples.times.reshape(-1, 1) + x = self.encoding(positions_flat, time_flat).to(positions) + density_before_activation = self.linear(x).view(*ray_samples.frustums.shape, -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) + return density, None + + def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None): + return {} + + +class NerfplayerNerfactoField(Field): + """NeRFPlayer (https://arxiv.org/abs/2210.15947) field with nerfacto backbone. + + Args: + aabb: parameters of scene aabb bounds + num_images: number of images in the dataset + num_layers: number of hidden layers + hidden_dim: dimension of hidden layers + geo_feat_dim: output geo feat dimensions + num_levels: number of levels of the hashmap for the base mlp + max_res: maximum resolution of the hashmap for the base mlp + log2_hashmap_size: size of the hashmap for the base mlp + num_layers_color: number of hidden layers for color network + num_layers_transient: number of hidden layers for transient network + hidden_dim_color: dimension of hidden layers for color network + hidden_dim_transient: dimension of hidden layers for transient network + appearance_embedding_dim: dimension of appearance embedding + transient_embedding_dim: dimension of transient embedding + use_transient_embedding: whether to use transient embedding + use_semantics: whether to use semantic segmentation + num_semantic_classes: number of semantic classes + use_pred_normals: whether to use predicted normals + use_average_appearance_embedding: whether to use average appearance embedding or zeros for inference + 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, + temporal_dim: int = 64, + num_levels: int = 16, + features_per_level: int = 2, + log2_hashmap_size: int = 19, + num_layers_color: int = 3, + num_layers_transient: int = 2, + hidden_dim_color: int = 64, + hidden_dim_transient: int = 64, + appearance_embedding_dim: int = 32, + transient_embedding_dim: int = 16, + use_transient_embedding: bool = False, + use_semantics: bool = False, + num_semantic_classes: int = 100, + use_pred_normals: bool = False, + use_average_appearance_embedding: bool = False, + spatial_distortion: Optional[SpatialDistortion] = None, + ) -> 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) + self.use_average_appearance_embedding = use_average_appearance_embedding + self.use_transient_embedding = use_transient_embedding + self.use_semantics = use_semantics + self.use_pred_normals = use_pred_normals + + self.direction_encoding = tcnn.Encoding( + n_input_dims=3, + encoding_config={ + "otype": "SphericalHarmonics", + "degree": 4, + }, + ) + + self.position_encoding = tcnn.Encoding( + n_input_dims=3, + encoding_config={"otype": "Frequency", "n_frequencies": 2}, + ) + + self.mlp_base = TemporalGridEncoder( + input_dim=3, + temporal_dim=temporal_dim, + num_levels=num_levels, + level_dim=features_per_level, + log2_hashmap_size=log2_hashmap_size, + desired_resolution=1024 * (self.aabb.max() - self.aabb.min()), + ) + self.mlp_base_decode = tcnn.Network( + n_input_dims=num_levels * features_per_level, + n_output_dims=1 + self.geo_feat_dim, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": hidden_dim, + "n_hidden_layers": num_layers - 1, + }, + ) + + # transients + if self.use_transient_embedding: + self.transient_embedding_dim = transient_embedding_dim + self.embedding_transient = Embedding(self.num_images, self.transient_embedding_dim) + self.mlp_transient = tcnn.Network( + n_input_dims=self.geo_feat_dim + self.transient_embedding_dim, + n_output_dims=hidden_dim_transient, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": hidden_dim_transient, + "n_hidden_layers": num_layers_transient - 1, + }, + ) + self.field_head_transient_uncertainty = UncertaintyFieldHead(in_dim=self.mlp_transient.n_output_dims) + self.field_head_transient_rgb = TransientRGBFieldHead(in_dim=self.mlp_transient.n_output_dims) + self.field_head_transient_density = TransientDensityFieldHead(in_dim=self.mlp_transient.n_output_dims) + + # semantics + if self.use_semantics: + self.mlp_semantics = tcnn.Network( + n_input_dims=self.geo_feat_dim, + n_output_dims=hidden_dim_transient, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": 64, + "n_hidden_layers": 1, + }, + ) + self.field_head_semantics = SemanticFieldHead( + in_dim=self.mlp_semantics.n_output_dims, num_classes=num_semantic_classes + ) + + # predicted normals + if self.use_pred_normals: + self.mlp_pred_normals = tcnn.Network( + n_input_dims=self.geo_feat_dim + self.position_encoding.n_output_dims, + n_output_dims=hidden_dim_transient, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": 64, + "n_hidden_layers": 2, + }, + ) + self.field_head_pred_normals = PredNormalsFieldHead(in_dim=self.mlp_pred_normals.n_output_dims) + + 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 = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + positions_flat = positions.view(-1, 3) + assert ray_samples.times is not None, "Time should be included in the input for NeRFPlayer" + time_flat = ray_samples.times.reshape(-1, 1) + h = self.mlp_base(positions_flat, time_flat) + h = self.mlp_base_decode(h).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): + assert density_embedding is not None + outputs = {} + 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) + + outputs_shape = ray_samples.frustums.directions.shape[:-1] + + # appearance + if self.training: + embedded_appearance = self.embedding_appearance(camera_indices) + else: + if self.use_average_appearance_embedding: + embedded_appearance = torch.ones( + (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device + ) * self.embedding_appearance.mean(dim=0) + else: + embedded_appearance = torch.zeros( + (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device + ) + + # transients + if self.use_transient_embedding and self.training: + embedded_transient = self.embedding_transient(camera_indices) + transient_input = torch.cat( + [ + density_embedding.view(-1, self.geo_feat_dim), + embedded_transient.view(-1, self.transient_embedding_dim), + ], + dim=-1, + ) + x = self.mlp_transient(transient_input).view(*outputs_shape, -1).to(directions) + outputs[FieldHeadNames.UNCERTAINTY] = self.field_head_transient_uncertainty(x) + outputs[FieldHeadNames.TRANSIENT_RGB] = self.field_head_transient_rgb(x) + outputs[FieldHeadNames.TRANSIENT_DENSITY] = self.field_head_transient_density(x) + + # semantics + if self.use_semantics: + density_embedding_copy = density_embedding.clone().detach() + semantics_input = torch.cat( + [ + density_embedding_copy.view(-1, self.geo_feat_dim), + ], + dim=-1, + ) + x = self.mlp_semantics(semantics_input).view(*outputs_shape, -1).to(directions) + outputs[FieldHeadNames.SEMANTICS] = self.field_head_semantics(x) + + # predicted normals + if self.use_pred_normals: + positions = ray_samples.frustums.get_positions() + + positions_flat = self.position_encoding(positions.view(-1, 3)) + pred_normals_inp = torch.cat([positions_flat, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) + + x = self.mlp_pred_normals(pred_normals_inp).view(*outputs_shape, -1).to(directions) + outputs[FieldHeadNames.PRED_NORMALS] = self.field_head_pred_normals(x) + + 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(*outputs_shape, -1).to(directions) + outputs.update({FieldHeadNames.RGB: rgb}) + + return outputs diff --git a/nerfstudio/fields/nerfplayer_ngp_field.py b/nerfstudio/fields/nerfplayer_ngp_field.py new file mode 100644 index 000000000..ed1cf250e --- /dev/null +++ b/nerfstudio/fields/nerfplayer_ngp_field.py @@ -0,0 +1,236 @@ +# Copyright 2022 The Nerfstudio 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. + +""" +NeRFPlayer (https://arxiv.org/abs/2210.15947) field implementations with InstantNGP backbone +""" + + +from typing import Optional + +import torch +from nerfacc import ContractionType, contract +from torch.nn.parameter import Parameter +from torchtyping import TensorType + +from nerfstudio.cameras.rays import Frustums, 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.temporal_grid import TemporalGridEncoder +from nerfstudio.fields.base_field import Field +from nerfstudio.fields.instant_ngp_field import get_normalized_directions + +try: + import tinycudann as tcnn +except ImportError: + # tinycudann module doesn't exist + pass + + +class NerfplayerNGPField(Field): + """NeRFPlayer (https://arxiv.org/abs/2210.15947) field with InstantNGP backbone. + + Args: + aabb: parameters of scene aabb bounds + temporal_dim: the dimension of temporal axis, a higher dimension indicates a higher temporal frequency + please refer to the implementation of TemporalGridEncoder for more details + num_levels: the number of multi-resolution levels; same as InstantNGP + features_per_level: the dim of output feature vector for each level; same as InstantNGP + log2_hashmap_size: the size of the table; same as InstantNGP + base_resolution: base resolution for the table; same as InstantNGP + num_layers: number of hidden layers (occupancy decoder network after sampling) + hidden_dim: dimension of hidden layers (occupancy decoder network after sampling) + 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 + disable_viewing_dependent: if true, disable the viewing dependent effect (no viewing direction as inputs) + Sometimes we need to disable viewing dependent effects in a dynamic scene, because there is + ambiguity between being dynamic and viewing dependent effects. For example, the shadow of the camera + should be a dynamic effect, but may be reconstructed as viewing dependent effects. + num_images: number of images, requried if use_appearance_embedding is True + appearance_embedding_dim: dimension of appearance embedding + contraction_type: type of contraction + """ + + def __init__( + self, + aabb, + temporal_dim: int = 16, + num_levels: int = 16, + features_per_level: int = 2, + log2_hashmap_size: int = 19, + base_resolution: int = 16, + 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: bool = False, + disable_viewing_dependent: bool = False, + num_images: Optional[int] = None, + appearance_embedding_dim: int = 32, + contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE, + ) -> None: + super().__init__() + + self.aabb = Parameter(aabb, requires_grad=False) + self.geo_feat_dim = geo_feat_dim + self.contraction_type = contraction_type + + 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) + + self.direction_encoding = tcnn.Encoding( + n_input_dims=3, + encoding_config={ + "otype": "SphericalHarmonics", + "degree": 4, + }, + ) + + self.mlp_base = TemporalGridEncoder( + input_dim=3, + temporal_dim=temporal_dim, + num_levels=num_levels, + level_dim=features_per_level, + base_resolution=base_resolution, + log2_hashmap_size=log2_hashmap_size, + desired_resolution=1024 * (self.aabb.max() - self.aabb.min()), + ) + self.mlp_base_decode = tcnn.Network( + n_input_dims=num_levels * features_per_level, + n_output_dims=1 + self.geo_feat_dim, + 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 disable_viewing_dependent: + in_dim = self.geo_feat_dim + self.direction_encoding = None + 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): + positions = ray_samples.frustums.get_positions() + positions_flat = positions.view(-1, 3) + positions_flat = contract(x=positions_flat, roi=self.aabb, type=self.contraction_type) + assert ray_samples.times is not None, "Time should be included in the input for NeRFPlayer" + times_flat = ray_samples.times.view(-1, 1) + + h = self.mlp_base(positions_flat, times_flat) + h = self.mlp_base_decode(h).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): + directions = get_normalized_directions(ray_samples.frustums.directions) + directions_flat = directions.view(-1, 3) + + if self.direction_encoding is not None: + d = self.direction_encoding(directions_flat) + if density_embedding is None: + positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + h = torch.cat([d, positions.view(-1, 3)], dim=-1) + else: + h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) + else: + # viewing direction is disabled + if density_embedding is None: + positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + h = positions.view(-1, 3) + else: + h = density_embedding.view(-1, self.geo_feat_dim) + + 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} + + # pylint: disable=arguments-differ + def density_fn(self, positions: TensorType["bs":..., 3], times: TensorType["bs":..., 1]) -> TensorType["bs":..., 1]: + """Returns only the density. Used primarily with the density grid. + Overwrite this function since density is time dependent now. + + Args: + positions: the origin of the samples/frustums + times: the time of each position + """ + ray_samples = RaySamples( + frustums=Frustums( + origins=positions, + directions=torch.ones_like(positions), + starts=torch.zeros_like(positions[..., :1]), + ends=torch.zeros_like(positions[..., :1]), + pixel_area=torch.ones_like(positions[..., :1]), + ), + times=times, + ) + density, _ = self.get_density(ray_samples) + return density + + def get_opacity(self, positions: TensorType["bs":..., 3], step_size, time_intervals=10) -> TensorType["bs":..., 1]: + """Returns the opacity for a position and time. Used primarily by the occupancy grid. + This will return the maximum opacity of the points in the space in a dynamic sequence. + + Args: + positions: the positions to evaluate the opacity at. + step_size: the step size to use for the opacity evaluation. + time_intervals: sample density on N time stamps + """ + # TODO: Converting opacity by time intervals is slow, and may lead to temporal artifacts. + # (Maybe random sample time and EMA?) + opacity = [] + for t in range(0, time_intervals): + density = self.density_fn(positions, t / (time_intervals - 1) * torch.ones_like(positions)[..., [0]]) + opacity.append(density * step_size) + opacity = torch.stack(opacity, dim=0).max(dim=0).values + return opacity diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index e41416053..5df52d2fa 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -503,6 +503,8 @@ class VolumetricSampler(Sampler): ), camera_indices=camera_indices, ) + if ray_bundle.times is not None: + ray_samples.times = ray_bundle.times[ray_indices] return ray_samples, ray_indices diff --git a/nerfstudio/models/nerfplayer_nerfacto.py b/nerfstudio/models/nerfplayer_nerfacto.py new file mode 100644 index 000000000..f35a748b1 --- /dev/null +++ b/nerfstudio/models/nerfplayer_nerfacto.py @@ -0,0 +1,251 @@ +# Copyright 2022 The Nerfstudio 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. + +""" +NeRFPlayer (https://arxiv.org/abs/2210.15947) implementation with nerfacto backbone. +""" + +from __future__ import annotations + +import functools +from dataclasses import dataclass, field +from typing import Dict, List, Type + +import numpy as np +import torch +from torchmetrics import PeakSignalNoiseRatio +from torchmetrics.functional import structural_similarity_index_measure +from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity +from typing_extensions import Literal + +from nerfstudio.cameras.rays import RayBundle +from nerfstudio.field_components.field_heads import FieldHeadNames +from nerfstudio.field_components.spatial_distortions import SceneContraction +from nerfstudio.fields.nerfplayer_nerfacto_field import ( + NerfplayerNerfactoField, + TemporalHashMLPDensityField, +) +from nerfstudio.model_components.losses import ( + MSELoss, + interlevel_loss, + orientation_loss, + pred_normal_loss, +) +from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler +from nerfstudio.model_components.renderers import ( + AccumulationRenderer, + DepthRenderer, + NormalsRenderer, + RGBRenderer, +) +from nerfstudio.model_components.scene_colliders import NearFarCollider +from nerfstudio.models.base_model import Model +from nerfstudio.models.nerfacto import NerfactoModel, NerfactoModelConfig + + +@dataclass +class NerfplayerNerfactoModelConfig(NerfactoModelConfig): + """Nerfplayer Model Config with Nerfacto backbone""" + + _target: Type = field(default_factory=lambda: NerfplayerNerfactoModel) + near_plane: float = 0.05 + """How far along the ray to start sampling.""" + far_plane: float = 1000.0 + """How far along the ray to stop sampling.""" + background_color: Literal["random", "last_sample"] = "random" + """Whether to randomize the background color. (Random is reported to be better on DyCheck.)""" + num_levels: int = 16 + """Hashing grid parameter.""" + features_per_level: int = 2 + """Hashing grid parameter.""" + log2_hashmap_size: int = 18 + """Hashing grid parameter.""" + temporal_dim: int = 32 + """Hashing grid parameter. A higher temporal dim means a higher temporal frequency.""" + proposal_net_args_list: List[Dict] = field( + default_factory=lambda: [ + {"hidden_dim": 16, "temporal_dim": 32, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 64}, + {"hidden_dim": 16, "temporal_dim": 32, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256}, + ] + ) + """Arguments for the proposal density fields.""" + distortion_loss_mult: float = 1e-2 + """Distortion loss multiplier.""" + temporal_tv_weight: float = 1 + """Temporal TV balancing weight for feature channels.""" + depth_weight: float = 1e-1 + """depth loss balancing weight for feature channels.""" + + +class NerfplayerNerfactoModel(NerfactoModel): + """Nerfplayer model with Nerfacto backbone. + + Args: + config: Nerfplayer configuration to instantiate model + """ + + config: NerfplayerNerfactoModelConfig + + def populate_modules(self): + """Set the fields and modules.""" + Model.populate_modules(self) + + scene_contraction = SceneContraction(order=float("inf")) + + # Fields + self.field = NerfplayerNerfactoField( + self.scene_box.aabb, + temporal_dim=self.config.temporal_dim, + num_levels=self.config.num_levels, + features_per_level=self.config.features_per_level, + log2_hashmap_size=self.config.log2_hashmap_size, + spatial_distortion=scene_contraction, + num_images=self.num_train_data, + use_pred_normals=self.config.predict_normals, + use_average_appearance_embedding=self.config.use_average_appearance_embedding, + ) + + self.density_fns = [] + num_prop_nets = self.config.num_proposal_iterations + # Build the proposal network(s) + self.proposal_networks = torch.nn.ModuleList() + if self.config.use_same_proposal_network: + assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed." + prop_net_args = self.config.proposal_net_args_list[0] + network = TemporalHashMLPDensityField( + self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args + ) + self.proposal_networks.append(network) + self.density_fns.extend([network.density_fn for _ in range(num_prop_nets)]) + else: + for i in range(num_prop_nets): + prop_net_args = self.config.proposal_net_args_list[min(i, len(self.config.proposal_net_args_list) - 1)] + network = TemporalHashMLPDensityField( + self.scene_box.aabb, + spatial_distortion=scene_contraction, + **prop_net_args, + ) + self.proposal_networks.append(network) + self.density_fns.extend([network.density_fn for network in self.proposal_networks]) + + # Samplers + update_schedule = lambda step: np.clip( + np.interp(step, [0, self.config.proposal_warmup], [0, self.config.proposal_update_every]), + 1, + self.config.proposal_update_every, + ) + self.proposal_sampler = ProposalNetworkSampler( + num_nerf_samples_per_ray=self.config.num_nerf_samples_per_ray, + num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray, + num_proposal_network_iterations=self.config.num_proposal_iterations, + single_jitter=self.config.use_single_jitter, + update_sched=update_schedule, + ) + + # Collider + self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane) + + # renderers + self.renderer_rgb = RGBRenderer(background_color=self.config.background_color) + self.renderer_accumulation = AccumulationRenderer() + self.renderer_depth = DepthRenderer(method="expected") # for depth loss + self.renderer_normals = NormalsRenderer() + + # losses + self.rgb_loss = MSELoss() + + # metrics + self.psnr = PeakSignalNoiseRatio(data_range=1.0) + self.ssim = structural_similarity_index_measure + self.lpips = LearnedPerceptualImagePatchSimilarity() + self.temporal_distortion = True # for viewer + + def get_outputs(self, ray_bundle: RayBundle): + assert ray_bundle.times is not None, "Time not provided." + ray_samples, weights_list, ray_samples_list = self.proposal_sampler( + ray_bundle, density_fns=[functools.partial(f, times=ray_bundle.times) for f in self.density_fns] + ) + field_outputs = self.field(ray_samples, compute_normals=self.config.predict_normals) + weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY]) + weights_list.append(weights) + ray_samples_list.append(ray_samples) + + rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights) + depth = self.renderer_depth(weights=weights, ray_samples=ray_samples) + accumulation = self.renderer_accumulation(weights=weights) + + outputs = { + "rgb": rgb, + "accumulation": accumulation, + "depth": depth, + } + + if self.config.predict_normals: + outputs["normals"] = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) + outputs["pred_normals"] = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights) + + # These use a lot of GPU memory, so we avoid storing them for eval. + if self.training: + outputs["weights_list"] = weights_list + outputs["ray_samples_list"] = ray_samples_list + + if self.training and self.config.predict_normals: + outputs["rendered_orientation_loss"] = orientation_loss( + weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions + ) + + outputs["rendered_pred_normal_loss"] = pred_normal_loss( + weights.detach(), + field_outputs[FieldHeadNames.NORMALS].detach(), + field_outputs[FieldHeadNames.PRED_NORMALS], + ) + + for i in range(self.config.num_proposal_iterations): + outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i]) + + return outputs + + def get_loss_dict(self, outputs, batch, metrics_dict=None): + loss_dict = {} + image = batch["image"].to(self.device) + loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"]) + if self.training: + loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss( + outputs["weights_list"], outputs["ray_samples_list"] + ) + assert metrics_dict is not None and "distortion" in metrics_dict + loss_dict["distortion_loss"] = self.config.distortion_loss_mult * metrics_dict["distortion"] + if self.config.predict_normals: + # orientation loss for computed normals + loss_dict["orientation_loss"] = self.config.orientation_loss_mult * torch.mean( + outputs["rendered_orientation_loss"] + ) + + # ground truth supervision for normals + loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean( + outputs["rendered_pred_normal_loss"] + ) + if "depth_image" in batch.keys() and self.config.depth_weight > 0: + mask = (batch["depth_image"] != 0).view([-1]) + loss_dict["depth_loss"] = 0 + l = lambda x: self.config.depth_weight * (x - batch["depth_image"][mask]).pow(2).mean() + loss_dict["depth_loss"] = l(outputs["depth"][mask]) + for i in range(self.config.num_proposal_iterations): + loss_dict["depth_loss"] += l(outputs[f"prop_depth_{i}"][mask]) + if self.config.temporal_tv_weight > 0: + loss_dict["temporal_tv_loss"] = self.field.mlp_base.get_temporal_tv_loss() + for net in self.proposal_networks: + loss_dict["temporal_tv_loss"] += net.encoding.get_temporal_tv_loss() + loss_dict["temporal_tv_loss"] *= self.config.temporal_tv_weight + return loss_dict diff --git a/nerfstudio/models/nerfplayer_ngp.py b/nerfstudio/models/nerfplayer_ngp.py new file mode 100644 index 000000000..39507a538 --- /dev/null +++ b/nerfstudio/models/nerfplayer_ngp.py @@ -0,0 +1,241 @@ +# Copyright 2022 The Nerfstudio 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. + +""" +Implementation of NeRFPlayer (https://arxiv.org/abs/2210.15947) with InstantNGP backbone. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Type + +import nerfacc +import torch +from nerfacc import ContractionType +from torch.nn import Parameter +from torchmetrics import PeakSignalNoiseRatio +from torchmetrics.functional import structural_similarity_index_measure +from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity +from typing_extensions import Literal + +from nerfstudio.cameras.rays import RayBundle +from nerfstudio.field_components.field_heads import FieldHeadNames +from nerfstudio.fields.nerfplayer_ngp_field import NerfplayerNGPField +from nerfstudio.model_components.losses import MSELoss +from nerfstudio.model_components.ray_samplers import VolumetricSampler +from nerfstudio.model_components.renderers import ( + AccumulationRenderer, + DepthRenderer, + RGBRenderer, +) +from nerfstudio.models.base_model import Model +from nerfstudio.models.instant_ngp import InstantNGPModelConfig, NGPModel +from nerfstudio.utils import colors + + +@dataclass +class NerfplayerNGPModelConfig(InstantNGPModelConfig): + """NeRFPlayer Model Config with InstantNGP backbone. + Tips for tuning the performance: + 1. If the scene is flickering, this is caused by unwanted high-freq on the temporal dimension. + Try reducing `temporal_dim` first, but don't be too small, otherwise the dynamic object is blurred. + Then try increasing the `temporal_tv_weight`. This is the loss for promoting smoothness among the + temporal channels. + 2. If a faster rendering is preferred, then try reducing `log2_hashmap_size`. If more details are + wanted, try increasing `log2_hashmap_size`. + 3. If the input cameras are of limited numbers, try reducing `num_levels`. `num_levels` is for + multi-resolution volume sampling, and has a similar behavior to the freq in NeRF. With a small + `num_levels`, a blurred rendering will be generated, but it is unlikely to overfit the training views. + """ + + _target: Type = field(default_factory=lambda: NerfplayerNGPModel) + temporal_dim: int = 64 + """Hashing grid parameter. A higher temporal dim means a higher temporal frequency.""" + num_levels: int = 16 + """Hashing grid parameter.""" + features_per_level: int = 2 + """Hashing grid parameter.""" + log2_hashmap_size: int = 17 + """Hashing grid parameter.""" + base_resolution: int = 16 + """Hashing grid parameter.""" + temporal_tv_weight: float = 1 + """Temporal TV loss balancing weight for feature channels.""" + depth_weight: float = 1e-1 + """depth loss balancing weight for feature channels.""" + train_background_color: Literal["random", "black", "white"] = "random" + """The training background color that is given to untrained areas.""" + eval_background_color: Literal["random", "black", "white"] = "white" + """The training background color that is given to untrained areas.""" + disable_viewing_dependent: bool = True + """Disable viewing dependent effects.""" + + +class NerfplayerNGPModel(NGPModel): + """NeRFPlayer Model with InstantNGP backbone. + + Args: + config: NeRFPlayer NGP configuration to instantiate model + """ + + config: NerfplayerNGPModelConfig + field: NerfplayerNGPField + + def populate_modules(self): + """Set the fields and modules.""" + Model.populate_modules(self) + + self.field = NerfplayerNGPField( + aabb=self.scene_box.aabb, + contraction_type=self.config.contraction_type, + use_appearance_embedding=self.config.use_appearance_embedding, + num_images=self.num_train_data, + temporal_dim=self.config.temporal_dim, + num_levels=self.config.num_levels, + features_per_level=self.config.features_per_level, + log2_hashmap_size=self.config.log2_hashmap_size, + base_resolution=self.config.base_resolution, + disable_viewing_dependent=self.config.disable_viewing_dependent, + ) + + self.scene_aabb = Parameter(self.scene_box.aabb.flatten(), requires_grad=False) + + # Occupancy Grid + self.occupancy_grid = nerfacc.OccupancyGrid( + roi_aabb=self.scene_aabb, + resolution=self.config.grid_resolution, + contraction_type=self.config.contraction_type, + ) + + # Sampler + vol_sampler_aabb = self.scene_box.aabb if self.config.contraction_type == ContractionType.AABB else None + self.sampler = VolumetricSampler( + scene_aabb=vol_sampler_aabb, + occupancy_grid=self.occupancy_grid, + density_fn=self.field.density_fn, + ) # need to update the density_fn later during forward (for input time) + + # renderers + self.train_background_color = self.config.train_background_color + self.eval_background_color = self.config.eval_background_color + if self.config.train_background_color in ["white", "black"]: + self.train_background_color = colors.COLORS_DICT[self.config.train_background_color] + if self.config.eval_background_color in ["white", "black"]: + self.eval_background_color = colors.COLORS_DICT[self.config.eval_background_color] + + self.renderer_rgb = RGBRenderer() # will update bgcolor later during forward + self.renderer_accumulation = AccumulationRenderer() + self.renderer_depth = DepthRenderer(method="expected") + + # losses + self.rgb_loss = MSELoss() + + # metrics + self.psnr = PeakSignalNoiseRatio(data_range=1.0) + self.ssim = structural_similarity_index_measure + self.lpips = LearnedPerceptualImagePatchSimilarity() + self.temporal_distortion = True # for viewer + + def get_outputs(self, ray_bundle: RayBundle): + num_rays = len(ray_bundle) + + # update the density_fn of the sampler so that the density is time aware + self.sampler.density_fn = lambda x: self.field.density_fn(x, ray_bundle.times) + with torch.no_grad(): + ray_samples, ray_indices = self.sampler( + ray_bundle=ray_bundle, + near_plane=self.config.near_plane, + far_plane=self.config.far_plane, + render_step_size=self.config.render_step_size, + cone_angle=self.config.cone_angle, + ) + + field_outputs = self.field(ray_samples) + + # accumulation + packed_info = nerfacc.pack_info(ray_indices, num_rays) + weights = nerfacc.render_weight_from_density( + packed_info=packed_info, + sigmas=field_outputs[FieldHeadNames.DENSITY], + t_starts=ray_samples.frustums.starts, + t_ends=ray_samples.frustums.ends, + ) + + # update bgcolor in the renderer; usually random color for training and fixed color for inference + if self.training: + self.renderer_rgb.background_color = self.train_background_color + else: + self.renderer_rgb.background_color = self.eval_background_color + rgb = self.renderer_rgb( + rgb=field_outputs[FieldHeadNames.RGB], + weights=weights, + ray_indices=ray_indices, + num_rays=num_rays, + ) + depth = self.renderer_depth( + weights=weights, ray_samples=ray_samples, ray_indices=ray_indices, num_rays=num_rays + ) + accumulation = self.renderer_accumulation(weights=weights, ray_indices=ray_indices, num_rays=num_rays) + alive_ray_mask = accumulation.squeeze(-1) > 0 + + outputs = { + "rgb": rgb, + "accumulation": accumulation, + "depth": depth, + "alive_ray_mask": alive_ray_mask, # the rays we kept from sampler + "num_samples_per_ray": packed_info[:, 1], + } + # adding them to outputs for calculating losses + if self.training and self.config.depth_weight > 0: + outputs["ray_indices"] = ray_indices + outputs["ray_samples"] = ray_samples + outputs["weights"] = weights + outputs["sigmas"] = field_outputs[FieldHeadNames.DENSITY] + return outputs + + def get_loss_dict(self, outputs, batch, metrics_dict=None): + image = batch["image"].to(self.device) + mask = outputs["alive_ray_mask"] + rgb_loss = self.rgb_loss(image[mask], outputs["rgb"][mask]) + loss_dict = {"rgb_loss": rgb_loss} + if "depth_image" in batch.keys() and self.config.depth_weight > 0: + mask = batch["depth_image"] != 0 + # First we calculate the depth value, just like most of the papers. + loss_dict["depth_loss"] = (outputs["depth"][mask] - batch["depth_image"][mask]).abs().mean() + # But this is not enough -- it will lead to fog like reconstructions, even with depth supervision. + # Because this loss only cares about the mean value. + # (Feel free to try it and see the fog-like effects on DyCheck by commenting out the following loss.) + # > The fog can be effectively penalized by multiview inputs (i.e., from another viewing point). + # > However, it is hard to get multiple views under the setting of dynamic scenes. + # > A new view always indicates another camera, which is expensive and brings sync problems. + # In DyCheck (https://arxiv.org/abs/2210.13445), surface sparsity regularizer + # (i.e., `distortion_loss` in nerfstudio) is used to make surface tight. + # The `distortion_loss` can be used here as well, but personally find it hard to implement... + # (Due to volume sampling, seems that cuda kernels are needed for efficiently computing the loss.) + # (Try nerfplayer with nerfacto backbone for distortion loss.) + # Instead, directly penalizing the empty space is found effective here. Perhaps due to a more + # "clear" loss (as it is directly applied to the network outputs, rather than post-processed values). + # But such a loss also has drawbacks: it tends to overfit wrong (or noise) presented in the depth map. + # Some structures are floating in the air with this loss... + gt_depth_packed = batch["depth_image"][outputs["ray_indices"]] + steps = (outputs["ray_samples"].frustums.starts + outputs["ray_samples"].frustums.ends) / 2 + # empty area should not be too close to the given depth, so lets add a margin to the gt depth + margin = (self.scene_box.aabb.max() - self.scene_box.aabb.min()) / 128 + density_min_mask = (gt_depth_packed - steps > margin) & (gt_depth_packed != 0) + loss_dict["depth_loss"] += (outputs["sigmas"][density_min_mask[..., 0]].pow(2)).mean() * 1e-2 + loss_dict["depth_loss"] *= self.config.depth_weight + if self.config.temporal_tv_weight > 0: + loss_dict["temporal_tv_loss"] = self.config.temporal_tv_weight * self.field.mlp_base.get_temporal_tv_loss() + return loss_dict diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index f62cf9d8f..d27548d20 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -196,11 +196,16 @@ class CheckThread(threading.Thread): while not self.state.check_done_render: # check camera data = self.state.vis["renderingState/camera"].read() + render_time = self.state.vis["renderingState/render_time"].read() if data is not None: camera_object = data["object"] - if self.state.prev_camera_matrix is None or ( - not np.allclose(camera_object["matrix"], self.state.prev_camera_matrix) - and not self.state.prev_moving + if ( + self.state.prev_camera_matrix is None + or ( + not np.allclose(camera_object["matrix"], self.state.prev_camera_matrix) + and not self.state.prev_moving + ) + or (render_time is not None and render_time != self.state.prev_render_time) ): self.state.check_interrupt_vis = True self.state.prev_moving = True @@ -280,6 +285,7 @@ class ViewerState: # viewer specific variables self.prev_camera_matrix = None + self.prev_render_time = 0 self.prev_output_type = OutputTypes.INIT self.prev_colormap_type = ColormapTypes.INIT self.prev_moving = False @@ -553,12 +559,23 @@ class ViewerState: return None camera_object = data["object"] + render_time = self.vis["renderingState/render_time"].read() - if self.prev_camera_matrix is not None and np.allclose(camera_object["matrix"], self.prev_camera_matrix): - self.camera_moving = False + if render_time is not None: + if ( + self.prev_camera_matrix is not None and np.allclose(camera_object["matrix"], self.prev_camera_matrix) + ) and (self.prev_render_time == render_time): + self.camera_moving = False + else: + self.prev_camera_matrix = camera_object["matrix"] + self.prev_render_time = render_time + self.camera_moving = True else: - self.prev_camera_matrix = camera_object["matrix"] - self.camera_moving = True + if self.prev_camera_matrix is not None and np.allclose(camera_object["matrix"], self.prev_camera_matrix): + self.camera_moving = False + else: + self.prev_camera_matrix = camera_object["matrix"] + self.camera_moving = True output_type = self.vis["renderingState/output_choice"].read() if output_type is None: @@ -901,10 +918,6 @@ class ViewerState: dim=0, ) - times = self.vis["renderingState/render_time"].read() - if times is not None: - times = torch.tensor([float(times)]) - camera_type_msg = camera_object["camera_type"] if camera_type_msg == "perspective": camera_type = CameraType.PERSPECTIVE @@ -922,11 +935,11 @@ class ViewerState: cy=intrinsics_matrix[1, 2], camera_type=camera_type, camera_to_worlds=camera_to_world[None, ...], - times=times, + times=torch.tensor([float(self.prev_render_time)]), ) camera = camera.to(graph.device) - camera_ray_bundle = camera.generate_rays(camera_indices=0, aabb_box=graph.render_aabb) + camera_ray_bundle = camera.generate_rays(camera_indices=0) graph.eval() diff --git a/tests/test_train.py b/tests/test_train.py index a3f1b6c60..5c52e50e8 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -23,6 +23,8 @@ BLACKLIST = [ "nerfacto", "phototourism", "depth-nerfacto", + "nerfplayer-ngp", + "nerfplayer-nerfacto", ]