diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index edc8da0a0..7cf107f2f 100644 --- a/nerfstudio/cameras/cameras.py +++ b/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, ) diff --git a/nerfstudio/data/datamanagers.py b/nerfstudio/data/datamanagers.py index 774808bde..2e908ad35 100644 --- a/nerfstudio/data/datamanagers.py +++ b/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. ) diff --git a/nerfstudio/data/dataparsers/base_dataparser.py b/nerfstudio/data/dataparsers/base_dataparser.py index f671f17ea..5f19f126c 100644 --- a/nerfstudio/data/dataparsers/base_dataparser.py +++ b/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: diff --git a/nerfstudio/data/dataparsers/dnerf_dataparser.py b/nerfstudio/data/dataparsers/dnerf_dataparser.py new file mode 100644 index 000000000..56a9f232b --- /dev/null +++ b/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