From 35610cb9e2cfe1e8872b7fc40328ef71c16bed68 Mon Sep 17 00:00:00 2001 From: Julian Knodt Date: Wed, 2 Nov 2022 12:23:13 -0700 Subject: [PATCH] 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 | 23 ++-- nerfstudio/data/datamanagers.py | 2 + .../data/dataparsers/base_dataparser.py | 2 + .../data/dataparsers/dnerf_dataparser.py | 108 ++++++++++++++++++ 4 files changed, 122 insertions(+), 13 deletions(-) create mode 100644 nerfstudio/data/dataparsers/dnerf_dataparser.py 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