From 9ff167d7050c158b7ccf7359d58c0e10b583c9c7 Mon Sep 17 00:00:00 2001 From: Julian Knodt Date: Tue, 15 Nov 2022 14:29:29 -0800 Subject: [PATCH] Add temporal distortion and D-NeRF model (#825) * Add D-NeRF temporal distortion and D-NeRF model This allows for temporal distortion to be enabled on a variety of models through config parameters, without an explicit model. Temporal distortion and MIP-NeRF are not currently jointly runnable but can be changed in the future. * Simplifying changes * Simplifying changes Co-authored-by: Matthew Tancik --- .gitignore | 3 +- nerfstudio/cameras/cameras.py | 10 +- nerfstudio/cameras/rays.py | 19 +++- nerfstudio/configs/method_configs.py | 31 ++++++- nerfstudio/data/utils/dataloaders.py | 1 + .../field_components/temporal_distortions.py | 92 +++++++++++++++++++ nerfstudio/models/base_model.py | 7 ++ nerfstudio/models/vanilla_nerf.py | 25 ++++- .../test_temporal_distortions.py | 25 +++++ 9 files changed, 206 insertions(+), 7 deletions(-) create mode 100644 nerfstudio/field_components/temporal_distortions.py create mode 100644 tests/field_components/test_temporal_distortions.py diff --git a/.gitignore b/.gitignore index 84ed2e9ce..72951da16 100644 --- a/.gitignore +++ b/.gitignore @@ -174,6 +174,7 @@ data # Misc old/ temp* +!temporal* .nfs* external/ __MACOSX/ @@ -183,4 +184,4 @@ bash/ cache/ renders/ package-lock.json -camera_paths/ \ No newline at end of file +camera_paths/ diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index 7cf107f2f..5057509c9 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -337,8 +337,14 @@ class Cameras: else: ray_bundle_camera_indices = camera_indices.view(pixel_area.shape) + times = self.times[camera_indices, None] if self.times is not None else None + return RayBundle( - origins=origins, directions=directions, pixel_area=pixel_area, camera_indices=ray_bundle_camera_indices + origins=origins, + directions=directions, + pixel_area=pixel_area, + camera_indices=ray_bundle_camera_indices, + times=times, ) def to_json( @@ -362,7 +368,7 @@ class Cameras: "fy": self.fy[camera_idx].tolist(), "camera_to_world": self.camera_to_worlds[camera_idx].tolist(), "camera_index": camera_idx, - "times": self.times[camera_idx] if self.times is not None else None, + "times": self.times[camera_idx].item() if self.times is not None else None, } if image is not None: image_uint8 = (image * 255).detach().type(torch.uint8) diff --git a/nerfstudio/cameras/rays.py b/nerfstudio/cameras/rays.py index 16b28eae8..d9dbaba0b 100644 --- a/nerfstudio/cameras/rays.py +++ b/nerfstudio/cameras/rays.py @@ -40,6 +40,8 @@ class Frustums(TensorDataclass): """Where the frustum ends along a ray.""" pixel_area: TensorType["bs":..., 1] """Projected area of pixel a distance 1 away from origin.""" + offsets: Optional[TensorType["bs":..., 3]] = None + """Offsets for each sample position""" def get_positions(self) -> TensorType[..., 3]: """Calulates "center" position of frustum. Not weighted by mass. @@ -47,7 +49,14 @@ class Frustums(TensorDataclass): Returns: xyz positions. """ - return self.origins + self.directions * (self.starts + self.ends) / 2 + pos = self.origins + self.directions * (self.starts + self.ends) / 2 + if self.offsets is not None: + pos = pos + self.offsets + return pos + + def set_offsets(self, offsets): + """Sets offsets for this frustum for computing positions""" + self.offsets = offsets def get_gaussian_blob(self) -> Gaussians: """Calculates guassian approximation of conical frustum. @@ -57,6 +66,8 @@ class Frustums(TensorDataclass): """ # Cone radius is set such that the square pixel_area matches the cone area. cone_radius = torch.sqrt(self.pixel_area) / 1.7724538509055159 # r = sqrt(pixel_area / pi) + if self.offsets is not None: + raise NotImplementedError() return conical_frustum_to_gaussian( origins=self.origins, directions=self.directions, @@ -100,6 +111,9 @@ class RaySamples(TensorDataclass): metadata: Optional[Dict[str, TensorType["bs":..., "latent_dims"]]] = None """addtional information relevant to generating ray samples""" + times: Optional[TensorType[..., 1]] = None + """Times at which rays are sampled""" + def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples", 1]: """Return weights based on predicted densities @@ -143,6 +157,8 @@ class RayBundle(TensorDataclass): """Rays Distance along ray to stop sampling""" metadata: Optional[Dict[str, TensorType["num_rays", "latent_dims"]]] = None """Additional metadata or data needed for interpolation, will mimic shape of rays""" + times: Optional[TensorType[..., 1]] = None + """Times at which rays are sampled""" def set_camera_indices(self, camera_index: int) -> None: """Sets all of the the camera indices to a specific camera index. @@ -223,6 +239,7 @@ class RayBundle(TensorDataclass): spacing_ends=spacing_ends, # [..., num_samples, 1] spacing_to_euclidean_fn=spacing_to_euclidean_fn, metadata=shaped_raybundle_fields.metadata, + times=None if self.times is None else self.times[..., None], # [..., 1, 1] ) return ray_samples diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 2423984c7..a3649458a 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -32,9 +32,11 @@ from nerfstudio.configs.base_config import ( from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManagerConfig from nerfstudio.data.datamanagers.semantic_datamanager import SemanticDataManagerConfig from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig +from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig +from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind from nerfstudio.models.base_model import VanillaModelConfig from nerfstudio.models.instant_ngp import InstantNGPModelConfig from nerfstudio.models.mipnerf import MipNerfModel @@ -53,6 +55,7 @@ descriptions = { "semantic-nerfw": "Predicts semantic segmentations and filters out transient objects.", "vanilla-nerf": "Original NeRF model. (slow)", "tensorf": "tensorf", + "dnerf": "Dynamic-NeRF model. (slow)", } method_configs["nerfacto"] = Config( @@ -159,7 +162,11 @@ method_configs["vanilla-nerf"] = Config( "fields": { "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), "scheduler": None, - } + }, + "temporal_distortion": { + "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), + "scheduler": None, + }, }, ) @@ -184,6 +191,28 @@ method_configs["tensorf"] = Config( }, ) +method_configs["dnerf"] = Config( + method_name="dnerf", + pipeline=VanillaPipelineConfig( + datamanager=VanillaDataManagerConfig(dataparser=DNeRFDataParserConfig()), + model=VanillaModelConfig( + _target=NeRFModel, + enable_temporal_distortion=True, + temporal_distortion_params={"kind": TemporalDistortionKind.DNERF}, + ), + ), + optimizers={ + "fields": { + "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), + "scheduler": None, + }, + "temporal_distortion": { + "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), + "scheduler": None, + }, + }, +) + 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/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py index 8f8873d65..5eec5e530 100644 --- a/nerfstudio/data/utils/dataloaders.py +++ b/nerfstudio/data/utils/dataloaders.py @@ -166,6 +166,7 @@ class EvalDataloader(DataLoader): camera_to_worlds=self.cameras.camera_to_worlds[image_idx : image_idx + 1], distortion_params=distortion_params, camera_type=self.cameras.camera_type[image_idx : image_idx + 1], + times=self.cameras.times[image_idx : image_idx + 1] if self.cameras.time else None, ) return camera diff --git a/nerfstudio/field_components/temporal_distortions.py b/nerfstudio/field_components/temporal_distortions.py new file mode 100644 index 000000000..cd50893e4 --- /dev/null +++ b/nerfstudio/field_components/temporal_distortions.py @@ -0,0 +1,92 @@ +# 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. + +"""Space distortions which occur as a function of time.""" + +from enum import Enum +from typing import Any, Dict, Optional, Tuple + +import torch +from torch import nn +from torchtyping import TensorType + +from nerfstudio.field_components.encodings import Encoding, NeRFEncoding +from nerfstudio.field_components.mlp import MLP + + +class TemporalDistortion(nn.Module): + """Apply spatial distortions as a function of time""" + + def forward(self, positions: TensorType["bs":..., 3], times: Optional[TensorType[1]]) -> TensorType["bs":..., 3]: + """ + Args: + positions: Samples to translate as a function of time + times: times for each sample + + Returns: + Translated positions. + """ + + +class TemporalDistortionKind(Enum): + """Possible temporal distortion names""" + + DNERF = "dnerf" + + def to_temporal_distortion(self, config: Dict[str, Any]) -> TemporalDistortion: + """Converts this kind to a temporal distortion""" + if self == TemporalDistortionKind.DNERF: + return DNeRFDistortion(**config) + raise NotImplementedError(f"Unknown temporal distortion kind {self}") + + +class DNeRFDistortion(TemporalDistortion): + """Optimizable temporal deformation using an MLP. + Args: + position_encoding: An encoding for the XYZ of distortion + temporal_encoding: An encoding for the time of distortion + mlp_num_layers: Number of layers in distortion MLP + mlp_layer_width: Size of hidden layer for the MLP + skip_connections: Number of layers for skip connections in the MLP + """ + + def __init__( + self, + position_encoding: Encoding = NeRFEncoding( + in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True + ), + temporal_encoding: Encoding = NeRFEncoding( + in_dim=1, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True + ), + mlp_num_layers: int = 4, + mlp_layer_width: int = 256, + skip_connections: Tuple[int] = (4,), + ) -> None: + super().__init__() + self.position_encoding = position_encoding + self.temporal_encoding = temporal_encoding + self.mlp_deform = MLP( + in_dim=self.position_encoding.get_out_dim() + self.temporal_encoding.get_out_dim(), + out_dim=3, + num_layers=mlp_num_layers, + layer_width=mlp_layer_width, + skip_connections=skip_connections, + ) + + def forward(self, positions, times=None): + if times is None: + return None + p = self.position_encoding(positions) + t = self.temporal_encoding(times) + return self.mlp_deform(torch.cat([p, t], dim=-1)) diff --git a/nerfstudio/models/base_model.py b/nerfstudio/models/base_model.py index ef36f53c8..57d0d3bea 100644 --- a/nerfstudio/models/base_model.py +++ b/nerfstudio/models/base_model.py @@ -32,6 +32,7 @@ from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.configs.config_utils import to_immutable_dict from nerfstudio.data.scene_box import SceneBox from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes +from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind from nerfstudio.model_components.scene_colliders import NearFarCollider @@ -77,6 +78,7 @@ class Model(nn.Module): self.num_train_data = num_train_data self.kwargs = kwargs self.collider = None + self.populate_modules() # populate the modules self.callbacks = None # to keep track of which device the nn.Module is on @@ -218,3 +220,8 @@ class VanillaModelConfig(ModelConfig): """Number of samples in coarse field evaluation""" num_importance_samples: int = 128 """Number of samples in fine field evaluation""" + + enable_temporal_distortion: bool = False + """Specifies whether or not to include ray warping based on time.""" + temporal_distortion_params: Dict[str, Any] = to_immutable_dict({"kind": TemporalDistortionKind.DNERF}) + """Parameters to instantiate temporal distortion with""" diff --git a/nerfstudio/models/vanilla_nerf.py b/nerfstudio/models/vanilla_nerf.py index 827bdfa26..2c5227ee4 100644 --- a/nerfstudio/models/vanilla_nerf.py +++ b/nerfstudio/models/vanilla_nerf.py @@ -55,6 +55,7 @@ class NeRFModel(Model): ) -> None: self.field_coarse = None self.field_fine = None + self.temporal_distortion = None super().__init__( config=config, @@ -73,8 +74,15 @@ class NeRFModel(Model): in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True ) - self.field_coarse = NeRFField(position_encoding=position_encoding, direction_encoding=direction_encoding) - self.field_fine = NeRFField(position_encoding=position_encoding, direction_encoding=direction_encoding) + self.field_coarse = NeRFField( + position_encoding=position_encoding, + direction_encoding=direction_encoding, + ) + + self.field_fine = NeRFField( + position_encoding=position_encoding, + direction_encoding=direction_encoding, + ) # samplers self.sampler_uniform = UniformSampler(num_samples=self.config.num_coarse_samples) @@ -93,11 +101,18 @@ class NeRFModel(Model): self.ssim = structural_similarity_index_measure self.lpips = LearnedPerceptualImagePatchSimilarity() + if getattr(self.config, "enable_temporal_distortion", False): + params = self.config.temporal_distortion_params + kind = params.pop("kind") + self.temporal_distortion = kind.to_temporal_distortion(params) + def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} if self.field_coarse is None or self.field_fine is None: raise ValueError("populate_fields() must be called before get_param_groups") param_groups["fields"] = list(self.field_coarse.parameters()) + list(self.field_fine.parameters()) + if self.temporal_distortion is not None: + param_groups["temporal_distortion"] = list(self.temporal_distortion.parameters()) return param_groups def get_outputs(self, ray_bundle: RayBundle): @@ -107,6 +122,9 @@ class NeRFModel(Model): # uniform sampling ray_samples_uniform = self.sampler_uniform(ray_bundle) + if self.temporal_distortion is not None: + offsets = self.temporal_distortion(ray_samples_uniform.frustums.get_positions(), ray_samples_uniform.times) + ray_samples_uniform.frustums.set_offsets(offsets) # coarse field: field_outputs_coarse = self.field_coarse.forward(ray_samples_uniform) @@ -120,6 +138,9 @@ class NeRFModel(Model): # pdf sampling ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) + if self.temporal_distortion is not None: + offsets = self.temporal_distortion(ray_samples_pdf.frustums.get_positions(), ray_samples_pdf.times) + ray_samples_pdf.frustums.set_offsets(offsets) # fine field: field_outputs_fine = self.field_fine.forward(ray_samples_pdf) diff --git a/tests/field_components/test_temporal_distortions.py b/tests/field_components/test_temporal_distortions.py new file mode 100644 index 000000000..61c240fd0 --- /dev/null +++ b/tests/field_components/test_temporal_distortions.py @@ -0,0 +1,25 @@ +""" +Test if temporal distortions run properly +""" +import torch + +from nerfstudio.field_components.temporal_distortions import DNeRFDistortion + + +def test_dnerf_distortion(): + """Test dnerf distortion""" + # pylint: disable=import-outside-toplevel + # pylint: disable=unused-import + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + distortion = DNeRFDistortion().to(device) + + num_rays = 1024 + num_samples = 256 + positions = torch.rand((num_rays, num_samples, 3), dtype=torch.float32, device=device) + assert distortion.forward(positions, None) is None + times = torch.rand_like(positions[..., :1]) + assert distortion.forward(positions, times).shape == positions.shape + + +if __name__ == "__main__": + test_dnerf_distortion()