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 <tancik@berkeley.edu>

.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/
+camera_paths/

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)

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

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)

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
 

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))

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"""

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)

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()