diff --git a/Dockerfile b/Dockerfile index c2e3ab82b..94587f550 100644 --- a/Dockerfile +++ b/Dockerfile @@ -7,7 +7,8 @@ ENV DEBIAN_FRONTEND=noninteractive ## Set timezone as it is required by some packages. ENV TZ=Europe/Berlin ## CUDA architectures, required by tiny-cuda-nn. -ENV TCNN_CUDA_ARCHITECTURES=86 +## NOTE: All commonly used GPU architectures are included and supported here. To speedup the image build process remove all architectures but the one of your explicit GPU. Find details here: https://developer.nvidia.com/cuda-gpus (8.6 translates to 86 in the line below) or in the docs. +ENV TCNN_CUDA_ARCHITECTURES=90;89;86;80;75;70;61;52;37 ## CUDA Home, required to find CUDA in some packages. ENV CUDA_HOME="/usr/local/cuda" diff --git a/README.md b/README.md index d4c29ce9c..ca2bd48b4 100644 --- a/README.md +++ b/README.md @@ -81,7 +81,7 @@ For more complex changes (e.g., running with your own data/setting up a new NeRF ### Prerequisites -CUDA must be installed on the system. This library has been tested with version 11.3. You can find more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html) +You must have an NVIDIA video card with CUDA installed on the system. This library has been tested with version 11.3 of CUDA. You can find more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html) ### Create environment diff --git a/docs/developer_guides/viewer/viewer_overview.md b/docs/developer_guides/viewer/viewer_overview.md index b80463e75..4e10319b6 100644 --- a/docs/developer_guides/viewer/viewer_overview.md +++ b/docs/developer_guides/viewer/viewer_overview.md @@ -50,6 +50,10 @@ yarn install yarn start ``` +The local webserver runs on port 4000 by default, +so when `ns-train` is running, you can connect to the viewer locally at +[http://localhost:4000/?websocket_url=ws://localhost:7007](http://localhost:4000/?websocket_url=ws://localhost:7007) + #### Running the Bridge Server Manually The viewer bridge server runs on the same machine that you use for training. The training code will connect to this server with a lightweight TCP connection using ZMQ. The training job will launch the viewer bridge server is you specify `--viewer.launch-bridge-server` in the terminal. Otherwise, you can launch the bridge server manually with the following script. diff --git a/docs/quickstart/installation.md b/docs/quickstart/installation.md index ae6a17cc9..97c63aac8 100644 --- a/docs/quickstart/installation.md +++ b/docs/quickstart/installation.md @@ -95,13 +95,28 @@ docker run --gpus all -v /folder/of/your/data:/workspace/ -v /home//. ``` ### Note - The container works on Linux and Windows, depending on your OS some additional setup steps might be required to provide access to your GPU inside containers. -- Paths on Windows use backslash '\\\\' while unix based systems use a frontslash '/' for paths, where backslashes might require an escape character depending on where they are used (e.g. C:\\\\folder1\\\\folder2...). Ensure to use the correct paths when mounting folders or providing paths as parameters. +- Paths on Windows use backslash '\\' while unix based systems use a frontslash '/' for paths, where backslashes might require an escape character depending on where they are used (e.g. C:\\\\folder1\\\\folder2...). Alternatively, mounts can be quoted (e.g. ```-v 'C:\local_folder:/docker_folder'```). Ensure to use the correct paths when mounting folders or providing paths as parameters. +- Always use full paths, relative paths are known to create issues when being used in mounts into docker. - Everything inside the container, what is not in a mounted folder (workspace in the above example), will be permanently removed after destroying the container. Always do all your tasks and output folder in workdir! - The user inside the container is called 'user' and is mapped to the local user with ID 1000 (usually the first non-root user on Linux systems). - The container currently is based on nvidia/cuda:11.7.1-devel-ubuntu22.04, consequently it comes with CUDA 11.7 which must be supported by the nvidia driver. No local CUDA installation is required or will be affected by using the docker image. - The docker image (respectively Ubuntu 22.04) comes with Python3.10, no older version of Python is installed. - If you call the container with commands directly, you still might want to add the interactive terminal ('-it') flag to get live log outputs of the nerfstudio scripts. In case the container is used in an automated environment the flag should be discarded. +- The current version of docker is built for multi-architecture (CUDA architectures) use. The target architecture must be defined at build time for tinyCUDNN to be able to compile properly. If your GPU architecture is not covered by the following table you need to replace the number in the line ```ENV TCNN_CUDA_ARCHITECTURES=90;89;86;80;75;70;61;52;37``` to you specific architecture. It also is a good idea to remove all architectures but yours (e.g. ```ENV TCNN_CUDA_ARCHITECTURES=86```) to speedup the docker build a lot. +**Currently supported CUDA architectures in the docker image** + +GPU | CUDA arch +-- | -- +H100 | 90 +40X0 | 89 +30X0 | 86 +A100 | 80 +20X0 | 75 +TITAN V / V100 | 70 +10X0 / TITAN Xp | 61 +9X0 | 52 +K80 | 37 ## Installation FAQ diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index fee7495a8..d7d228bae 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -16,7 +16,6 @@ Camera Models """ import base64 -import importlib import math import os from dataclasses import dataclass @@ -147,6 +146,8 @@ class Cameras(TensorDataclass): self.__post_init__() # This will do the dataclass post_init and broadcast all the tensors + self._use_nerfacc = strtobool(os.environ.get("INTERSECT_WITH_NERFACC", "TRUE")) + def _init_get_fc_xy(self, fc_xy, name): """ Parses the input focal length / principle point x or y and returns a tensor of the correct shape @@ -470,11 +471,7 @@ class Cameras(TensorDataclass): rays_o = rays_o.reshape((-1, 3)) rays_d = rays_d.reshape((-1, 3)) - if strtobool(os.environ.get("INTERSECT_WITH_NERFACC", "TRUE")): - nerfacc = importlib.import_module("nerfacc") - t_min, t_max = nerfacc.ray_aabb_intersect(rays_o, rays_d, tensor_aabb) - else: - t_min, t_max = nerfstudio.utils.math.intersect_aabb(rays_o, rays_d, tensor_aabb) + t_min, t_max = nerfstudio.utils.math.intersect_aabb(rays_o, rays_d, tensor_aabb) t_min = t_min.reshape([shape[0], shape[1], 1]) t_max = t_max.reshape([shape[0], shape[1], 1]) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 3e6f3f80d..fee855fac 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -36,6 +36,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, @@ -54,6 +55,8 @@ from nerfstudio.models.dreamfusion import DreamFusionModelConfig 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 @@ -73,7 +76,9 @@ descriptions = { "tensorf": "tensorf", "dnerf": "Dynamic-NeRF model. (slow)", "phototourism": "Uses the Phototourism data.", - "dreamfusion": "NeRF Model used in implementation of stable dreamfusion.", + "dreamfusion": "Generative Text to NeRF model", + "nerfplayer-nerfacto": "NeRFPlayer with nerfacto backbone.", + "nerfplayer-ngp": "NeRFPlayer with InstantNGP backbone.", } method_configs["nerfacto"] = TrainerConfig( @@ -356,6 +361,58 @@ method_configs["dreamfusion"] = DreamfusionTrainerConfig( vis="viewer", ) +method_configs["nerfplayer-nerfacto"] = TrainerConfig( + method_name="nerfplayer-nerfacto", + 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[ 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/generative/__init__.py b/nerfstudio/generative/__init__.py new file mode 100644 index 000000000..d926cf95c --- /dev/null +++ b/nerfstudio/generative/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index d64bfdfeb..af3aa2196 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -12,7 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -# From https://github.com/ashawkey/stable-dreamfusion/blob/main/nerf/sd.py +"""Stable Diffusion helpers""" + +# Modified from https://github.com/ashawkey/stable-dreamfusion/blob/main/nerf/sd.py import sys from pathlib import Path @@ -77,15 +79,8 @@ class StableDiffusion(nn.Module): assert pipe is not None pipe = pipe.to(self.device) - # MEMORY IMPROVEMENTS pipe.enable_attention_slicing() - # Needs xformers package (installation is from source https://github.com/facebookresearch/xformers) - # pipe.enable_xformers_memory_efficient_attention() - - # More memory savings for 1/3rd performance - # pipe.enable_sequential_cpu_offload() - self.tokenizer = pipe.tokenizer self.unet = pipe.unet self.text_encoder = pipe.text_encoder @@ -138,7 +133,7 @@ class StableDiffusion(nn.Module): text_embeddings: TensorType["N", "max_length", "embed_dim"], image: TensorType["BS", 3, "H", "W"], guidance_scale: float = 100.0, - ) -> Tuple[torch.Tensor, TensorType["BS", "2N", "H", "W"], TensorType["BS", "2N", "H", "W"]]: + ) -> Tuple[torch.Tensor, TensorType["BS", 4, "H", "W"], TensorType["BS", 4, "H", "W"]]: """Score Distilation Sampling loss proposed in DreamFusion paper (https://dreamfusion3d.github.io/) Args: @@ -172,9 +167,9 @@ class StableDiffusion(nn.Module): grad = w * (noise_pred - noise) grad = torch.nan_to_num(grad) - noise_loss = torch.mean(torch.nan_to_num(torch.square(noise_pred - noise))) + sds_mse = torch.mean(torch.nan_to_num(torch.square(noise_pred - noise))) - return noise_loss, latents, grad + return sds_mse, latents, grad def produce_latents( self, @@ -183,8 +178,8 @@ class StableDiffusion(nn.Module): width: int = IMG_DIM, num_inference_steps: int = 50, guidance_scale: float = 7.5, - latents: Optional[TensorType["BS", "2N", "H", "W"]] = None, - ) -> TensorType["BS", "2N", "H", "W"]: + latents: Optional[TensorType["BS", 4, "H", "W"]] = None, + ) -> TensorType["BS", 4, "H", "W"]: """Produce latents for a given text embedding Args: @@ -224,7 +219,7 @@ class StableDiffusion(nn.Module): latents = self.scheduler.step(noise_pred, t, latents)["prev_sample"] # type: ignore return latents - def latents_to_img(self, latents: TensorType["BS", "2N", "H", "W"]) -> TensorType["BS", 3, "H", "W"]: + def latents_to_img(self, latents: TensorType["BS", 4, "H", "W"]) -> TensorType["BS", 3, "H", "W"]: """Convert latents to images Args: @@ -243,12 +238,12 @@ class StableDiffusion(nn.Module): return imgs - def imgs_to_latent(self, imgs: TensorType["BS", 3, "H", "W"]) -> TensorType["BS", "2N", "H", "W"]: + def imgs_to_latent(self, imgs: TensorType["BS", 3, "H", "W"]) -> TensorType["BS", 4, "H", "W"]: """Convert images to latents - + Args: imgs: Images to convert - + Returns: Latents """ 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/utils/math.py b/nerfstudio/utils/math.py index 8e1ebe24b..937866dcc 100644 --- a/nerfstudio/utils/math.py +++ b/nerfstudio/utils/math.py @@ -15,10 +15,13 @@ """ Math Helper Functions """ from dataclasses import dataclass +from typing import Tuple import torch from torchtyping import TensorType +_USE_NERFACC = True + def components_from_spherical_harmonics(levels: int, directions: TensorType[..., 3]) -> TensorType[..., "components"]: """ @@ -193,53 +196,92 @@ def expected_sin(x_means: torch.Tensor, x_vars: torch.Tensor) -> torch.Tensor: return torch.exp(-0.5 * x_vars) * torch.sin(x_means) -# From https://github.com/ashawkey/stable-dreamfusion -def safe_normalize(x, eps=1e-20): - """Safe normalization w/o divide by zero issues""" - return x / torch.sqrt(torch.clamp(torch.sum(x * x, -1, keepdim=True), min=eps)) - - -def intersect_aabb(origins, directions, aabb): +@torch.jit.script +def _intersect_aabb( + origins: torch.Tensor, + directions: torch.Tensor, + aabb: torch.Tensor, + max_bound: float = 1e10, + invalid_value: float = 1e10, +) -> Tuple[torch.Tensor, torch.Tensor]: """ - pytorch implementation of ray intersection with AABB box - :param origin: N,3 tensor of 3d positions - :param direction: N,3 tensor of normalized directions - :param aabb: 6,1 array of aabb box in the form of [x_min, y_min, z_min, x_max, y_max, z_max] - :return: t_min, t_max - two tensors of shapes N,1 representing the intersection signed distance from - the origin. t_min is clipped to 0 in case it is negative, in case of no intersection t_min is 1e10. + Implementation of ray intersection with AABB box + + Args: + origins: [N,3] tensor of 3d positions + directions: [N,3] tensor of normalized directions + aabb: [6] array of aabb box in the form of [x_min, y_min, z_min, x_max, y_max, z_max] + max_bound: Maximum value of t_max + invalid_value: Value to return in case of no intersection + + Returns: + t_min, t_max - two tensors of shapes N representing distance of intersection from the origin. """ - max_bound = 1e10 + tx_min = (aabb[:3] - origins) / directions + tx_max = (aabb[3:] - origins) / directions - tx_min = (aabb[0] - origins[:, 0]) / directions[:, 0] - tx_max = (aabb[3] - origins[:, 0]) / directions[:, 0] t_min = torch.min(tx_min, tx_max) t_max = torch.max(tx_min, tx_max) - ty_min_temp = (aabb[1] - origins[:, 1]) / directions[:, 1] - ty_max_temp = (aabb[4] - origins[:, 1]) / directions[:, 1] - ty_min = torch.min(ty_min_temp, ty_max_temp) - ty_max = torch.max(ty_min_temp, ty_max_temp) + t_min = torch.max(t_min, dim=-1).values + t_max = torch.min(t_max, dim=-1).values - cond = torch.logical_or((t_min > ty_max), (ty_min > t_max)) - t_min = torch.where(cond, max_bound, t_min) - t_max = torch.where(cond, max_bound, t_max) + t_min = torch.clamp(t_min, min=0, max=max_bound) + t_max = torch.clamp(t_max, min=0, max=max_bound) - t_min = torch.max(t_min, ty_min) - t_max = torch.min(t_max, ty_max) - - tz_min_temp = (aabb[2] - origins[:, 2]) / directions[:, 2] - tz_max_temp = (aabb[5] - origins[:, 2]) / directions[:, 2] - tz_min = torch.min(tz_min_temp, tz_max_temp) - tz_max = torch.max(tz_min_temp, tz_max_temp) - - cond = torch.logical_or((t_min > tz_max), tz_min > t_max) - t_min = torch.where(cond, max_bound, t_min) - t_max = torch.where(cond, max_bound, t_max) - - t_min = torch.max(t_min, tz_min) - t_max = torch.min(t_max, tz_max) - t_min = torch.where(t_min > 0, t_min, 0) - t_max = torch.where(t_min == max_bound, max_bound, t_max) + cond = t_max <= t_min + t_min = torch.where(cond, invalid_value, t_min) + t_max = torch.where(cond, invalid_value, t_max) return t_min, t_max + + +def intersect_aabb( + origins: TensorType["N", 3], + directions: TensorType["N", 3], + aabb: TensorType[6], +) -> Tuple[TensorType["N"], TensorType["N"]]: + """ + Implementation of ray intersection with AABB box + + Args: + origins: 3d positions + directions: Normalized directions + aabb: array of aabb box in the form of [x_min, y_min, z_min, x_max, y_max, z_max] + max_bound: Maximum value of t_max + invalid_value: Value to return in case of no intersection + + Returns: + t_min, t_max - two tensors of shapes N representing distance of intersection from the origin. + """ + + global _USE_NERFACC # pylint: disable=global-statement + if _USE_NERFACC: + try: + import nerfacc # pylint: disable=import-outside-toplevel + + t_min, t_max = nerfacc.ray_aabb_intersect(origins, directions, aabb) + except: # pylint: disable=bare-except + t_min, t_max = _intersect_aabb(origins, directions, aabb, max_bound=1e10, invalid_value=1e10) + _USE_NERFACC = False + else: + t_min, t_max = _intersect_aabb(origins, directions, aabb, max_bound=1e10, invalid_value=1e10) + + return t_min, t_max + + +def safe_normalize( + vectors: TensorType["batch_dim":..., "N"], + eps: float = 1e-10, +) -> TensorType["batch_dim":..., "N"]: + """Normalizes vectors. + + Args: + vectors: Vectors to normalize. + eps: Epsilon value to avoid division by zero. + + Returns: + Normalized vectors. + """ + return vectors / (torch.norm(vectors, dim=-1, keepdim=True) + eps) diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index f62cf9d8f..44056537f 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 @@ -395,13 +401,14 @@ class ViewerState: trainer.save_checkpoint(step) # get all camera paths camera_path_dir = os.path.join(self.datapath, "camera_paths") - camera_path_files = os.listdir(camera_path_dir) - all_path_dict = {} - for i in camera_path_files: - if i[-4:] == "json": - all_path_dict[i[:-5]] = load_from_json(Path(os.path.join(camera_path_dir, i))) - self.vis["renderingState/all_camera_paths"].write(all_path_dict) - self.vis["populate_paths_payload"].delete() + if os.path.exists(camera_path_dir): + camera_path_files = os.listdir(camera_path_dir) + all_path_dict = {} + for i in camera_path_files: + if i[-4:] == "json": + all_path_dict[i[:-5]] = load_from_json(Path(os.path.join(camera_path_dir, i))) + self.vis["renderingState/all_camera_paths"].write(all_path_dict) + self.vis["populate_paths_payload"].delete() def _check_webrtc_offer(self): """Check if there is a webrtc offer to respond to.""" @@ -553,12 +560,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 +919,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,7 +936,7 @@ 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) diff --git a/pyproject.toml b/pyproject.toml index 10b94e39b..21bc69456 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,12 +31,12 @@ dependencies = [ "mediapy>=1.1.0", "msgpack>=1.0.4", "msgpack_numpy>=0.4.8", - "nerfacc==0.3.3", + "nerfacc==0.3.4", "open3d>=0.16.0", "opencv-python==4.6.0.66", "Pillow>=9.3.0", "plotly>=5.7.0", - "protobuf<=3.20.3", + "protobuf<=3.20.3,!=3.20.0", "pymeshlab>=2022.2.post2", "pyngrok>=5.1.0", "python-socketio>=5.7.1", diff --git a/tests/test_train.py b/tests/test_train.py index ac0aa7bc3..f3f4e2c2c 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -15,7 +15,18 @@ from nerfstudio.data.dataparsers.minimal_dataparser import MinimalDataParserConf from nerfstudio.engine.trainer import TrainerConfig from scripts.train import train_loop -BLACKLIST = ["base", "semantic-nerfw", "instant-ngp", "instant-ngp-bounded", "nerfacto", "phototourism", "dreamfusion"] +BLACKLIST = [ + "base", + "semantic-nerfw", + "instant-ngp", + "instant-ngp-bounded", + "nerfacto", + "phototourism", + "depth-nerfacto", + "nerfplayer-ngp", + "nerfplayer-nerfacto", + "dreamfusion", +] def set_reduced_config(config: TrainerConfig):