DNeRF Data Parser (#817)

* Add dnerf dataparser

* Switch time to be stored per camera

I realized it would be better downstream to store times on the ray bundle, and in order to store
them on the ray bundle it's much easier to store them on the camera. Rather than having to
revert this change, I just made it here.

nerfstudio/cameras/cameras.py

@@ -81,6 +81,7 @@ class Cameras:
         camera_type: Optional[
             Union[TensorType["num_cameras"], int, List[CameraType], CameraType]
         ] = CameraType.PERSPECTIVE,
+        times: Optional[TensorType["num_cameras"]] = None,
     ):
         self._num_cameras = camera_to_worlds.shape[0]
         self.camera_to_worlds = camera_to_worlds  # This comes first since it determines @property self.device
@@ -121,6 +122,8 @@ class Cameras:
 
         self.camera_type = self._init_get_camera_type(camera_type)
 
+        self.times = times
+
     def _init_get_camera_type(
         self, camera_type: Union[TensorType["num_cameras"], int, List[CameraType], CameraType]
     ) -> TensorType["num_cameras"]:
@@ -223,6 +226,7 @@ class Cameras:
             height=self.image_height.to(device),
             distortion_params=distortion_params.to(device) if distortion_params is not None else None,
             camera_type=self.camera_type.to(device),
+            times=self.times.to(device) if self.times is not None else None,
         )
 
     def get_image_coords(self, pixel_offset: float = 0.5) -> TensorType["height", "width", 2]:
@@ -305,11 +309,7 @@ class Cameras:
 
             sin_theta = torch.sin(theta)
             directions_stack = torch.stack(
-                [
-                    coord_stack[..., 0] * sin_theta / theta,
-                    coord_stack[..., 1] * sin_theta / theta,
-                    -torch.cos(theta),
-                ],
+                [coord_stack[..., 0] * sin_theta / theta, coord_stack[..., 1] * sin_theta / theta, -torch.cos(theta)],
                 dim=-1,
             )
         else:
@@ -338,17 +338,11 @@ class Cameras:
             ray_bundle_camera_indices = camera_indices.view(pixel_area.shape)
 
         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
         )
 
     def to_json(
-        self,
-        camera_idx: int,
-        image: Optional[TensorType["height", "width", 2]] = None,
-        max_size: Optional[int] = None,
+        self, camera_idx: int, image: Optional[TensorType["height", "width", 2]] = None, max_size: Optional[int] = None
     ) -> Dict:
         """Convert a camera to a json dictionary.
 
@@ -368,6 +362,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,
         }
         if image is not None:
             image_uint8 = (image * 255).detach().type(torch.uint8)
@@ -422,6 +417,7 @@ class Cameras:
                 width=self._image_widths[indices],
                 distortion_params=self.distortion_params[indices] if self.distortion_params is not None else None,
                 camera_type=self.camera_type[indices],
+                times=self.times[indices] if self.times is not None else None,
             )
         if isinstance(indices, (int, slice)):
             indices = (indices,)
@@ -437,4 +433,5 @@ class Cameras:
             if self.distortion_params is not None
             else None,
             camera_type=self.camera_type[indices + (slice(None),)],
+            times=self.times[indices + (slice(None),)] if self.times is not None else None,
         )

nerfstudio/data/datamanagers.py

@@ -34,6 +34,7 @@ from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.configs.base_config import InstantiateConfig
 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.instant_ngp_dataparser import (
     InstantNGPDataParserConfig,
@@ -61,6 +62,7 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[  # Omit prefixes of
             "friends-data": FriendsDataParserConfig(),
             "instant-ngp-data": InstantNGPDataParserConfig(),
             "record3d-data": Record3DDataParserConfig(),
+            "dnerf-data": DNeRFDataParserConfig(),
         },
         prefix_names=False,  # Omit prefixes in subcommands themselves.
     )

nerfstudio/data/dataparsers/base_dataparser.py

@@ -63,6 +63,8 @@ class DataparserOutputs:
     """Scene box of dataset. Used to bound the scene or provide the scene scale depending on model."""
     semantics: Optional[Semantics] = None
     """Semantics information."""
+    times: Optional[TensorType[1]] = None
+    """Time in range [0,1] for when each image was taken."""
     additional_inputs: Dict[str, Any] = to_immutable_dict({})
     """Dictionary of additional dataset information (e.g. semantics/point clouds/masks).
     {input_name:

nerfstudio/data/dataparsers/dnerf_dataparser.py

@@ -0,0 +1,108 @@
+# 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.
+
+"""Data parser for blender dataset"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Type
+
+import imageio
+import numpy as np
+import torch
+
+from nerfstudio.cameras.cameras import Cameras, CameraType
+from nerfstudio.data.dataparsers.base_dataparser import (
+    DataParser,
+    DataParserConfig,
+    DataparserOutputs,
+)
+from nerfstudio.data.scene_box import SceneBox
+from nerfstudio.utils.colors import get_color
+from nerfstudio.utils.io import load_from_json
+
+
+@dataclass
+class DNeRFDataParserConfig(DataParserConfig):
+    """D-NeRF dataset parser config"""
+
+    _target: Type = field(default_factory=lambda: DNeRF)
+    """target class to instantiate"""
+    data: Path = Path("data/dnerf/lego")
+    """Directory specifying location of data."""
+    scale_factor: float = 1.0
+    """How much to scale the camera origins by."""
+    alpha_color: str = "white"
+    """alpha color of background"""
+
+
+@dataclass
+class DNeRF(DataParser):
+    """DNeRF Dataset"""
+
+    config: DNeRFDataParserConfig
+
+    def __init__(self, config: DNeRFDataParserConfig):
+        super().__init__(config=config)
+        self.data: Path = config.data
+        self.scale_factor: float = config.scale_factor
+        self.alpha_color = config.alpha_color
+
+    def _generate_dataparser_outputs(self, split="train"):
+        if self.alpha_color is not None:
+            alpha_color_tensor = get_color(self.alpha_color)
+        else:
+            alpha_color_tensor = None
+
+        meta = load_from_json(self.data / f"transforms_{split}.json")
+        image_filenames = []
+        poses = []
+        times = []
+        for frame in meta["frames"]:
+            fname = self.data / Path(frame["file_path"].replace("./", "") + ".png")
+            image_filenames.append(fname)
+            poses.append(np.array(frame["transform_matrix"]))
+            times.append(frame["time"])
+        poses = np.array(poses).astype(np.float32)
+        times = torch.tensor(times, dtype=torch.float32)
+
+        img_0 = imageio.imread(image_filenames[0])
+        image_height, image_width = img_0.shape[:2]
+        camera_angle_x = float(meta["camera_angle_x"])
+        focal_length = 0.5 * image_width / np.tan(0.5 * camera_angle_x)
+
+        cx = image_width / 2.0
+        cy = image_height / 2.0
+        camera_to_world = torch.from_numpy(poses[:, :3])  # camera to world transform
+
+        # in x,y,z order
+        camera_to_world[..., 3] *= self.scale_factor
+        scene_box = SceneBox(aabb=torch.tensor([[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]], dtype=torch.float32))
+
+        cameras = Cameras(
+            camera_to_worlds=camera_to_world,
+            fx=focal_length,
+            fy=focal_length,
+            cx=cx,
+            cy=cy,
+            camera_type=CameraType.PERSPECTIVE,
+            times=times,
+        )
+
+        dataparser_outputs = DataparserOutputs(
+            image_filenames=image_filenames, cameras=cameras, alpha_color=alpha_color_tensor, scene_box=scene_box
+        )
+
+        return dataparser_outputs