diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 4fc1822dd..d032c4d48 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -38,6 +38,7 @@ from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig +from nerfstudio.data.dataparsers.dycheck_dataparser import DycheckDataParserConfig from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig from nerfstudio.data.dataparsers.instant_ngp_dataparser import ( InstantNGPDataParserConfig, @@ -73,6 +74,7 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of "nuscenes-data": NuScenesDataParserConfig(), "dnerf-data": DNeRFDataParserConfig(), "phototourism-data": PhototourismDataParserConfig(), + "dycheck-data": DycheckDataParserConfig(), }, prefix_names=False, # Omit prefixes in subcommands themselves. ) diff --git a/nerfstudio/data/dataparsers/dycheck_dataparser.py b/nerfstudio/data/dataparsers/dycheck_dataparser.py new file mode 100644 index 000000000..92cca35ff --- /dev/null +++ b/nerfstudio/data/dataparsers/dycheck_dataparser.py @@ -0,0 +1,342 @@ +# 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 DyCheck (https://arxiv.org/abs/2210.13445) dataset of `iphone` subset""" +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, List, Tuple, Type + +import cv2 +import numpy as np +import torch +from rich.progress import Console + +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 + +CONSOLE = Console(width=120) + + +def downscale(img, scale: int) -> np.ndarray: + """Function from DyCheck's repo. Downscale an image. + + Args: + img: Input image + scale: Factor of the scale + + Returns: + New image + """ + if scale == 1: + return img + height, width = img.shape[:2] + if height % scale > 0 or width % scale > 0: + raise ValueError(f"Image shape ({height},{width}) must be divisible by the" f" scale ({scale}).") + out_height, out_width = height // scale, width // scale + resized = cv2.resize(img, (out_width, out_height), cv2.INTER_AREA) + return resized + + +def upscale(img, scale: int) -> np.ndarray: + """Function from DyCheck's repo. Upscale an image. + + Args: + img: Input image + scale: Factor of the scale + + Returns: + New image + """ + if scale == 1: + return img + height, width = img.shape[:2] + out_height, out_width = height * scale, width * scale + resized = cv2.resize(img, (out_width, out_height), cv2.INTER_AREA) + return resized + + +def rescale(img, scale_factor: float, interpolation=cv2.INTER_AREA) -> np.ndarray: + """Function from DyCheck's repo. Rescale an image. + + Args: + img: Input image + scale: Factor of the scale + interpolation: Interpolation method in opencv + + Returns: + New image + """ + scale_factor = float(scale_factor) + if scale_factor <= 0.0: + raise ValueError("scale_factor must be a non-negative number.") + if scale_factor == 1.0: + return img + + height, width = img.shape[:2] + if scale_factor.is_integer(): + return upscale(img, int(scale_factor)) + + inv_scale = 1.0 / scale_factor + if inv_scale.is_integer() and (scale_factor * height).is_integer() and (scale_factor * width).is_integer(): + return downscale(img, int(inv_scale)) + + print(f"Resizing image by non-integer factor {scale_factor}, this may lead to artifacts.") + height, width = img.shape[:2] + out_height = math.ceil(height * scale_factor) + out_height -= out_height % 2 + out_width = math.ceil(width * scale_factor) + out_width -= out_width % 2 + + return cv2.resize(img, (out_width, out_height), interpolation) + + +def _load_scene_info(data_dir: Path) -> Tuple[np.ndarray, float, float, float]: + """Function from DyCheck's repo. Load scene info from json. + + Args: + data_dir: data path + + Returns: + A tuple of scene info: center, scale, near, far + """ + scene_dict = load_from_json(data_dir / "scene.json") + center = np.array(scene_dict["center"], dtype=np.float32) + scale = scene_dict["scale"] + near = scene_dict["near"] + far = scene_dict["far"] + return center, scale, near, far + + +def _load_metadata_info(data_dir: Path) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Function from DyCheck's repo. Load scene metadata from json. + + Args: + data_dir: data path + + Returns: + A tuple of scene info: frame_names_map, time_ids, camera_ids + """ + dataset_dict = load_from_json(data_dir / "dataset.json") + _frame_names = np.array(dataset_dict["ids"]) + + metadata_dict = load_from_json(data_dir / "metadata.json") + time_ids = np.array([metadata_dict[k]["warp_id"] for k in _frame_names], dtype=np.uint32) + camera_ids = np.array([metadata_dict[k]["camera_id"] for k in _frame_names], dtype=np.uint32) + + frame_names_map = np.zeros((time_ids.max() + 1, camera_ids.max() + 1), _frame_names.dtype) + for i, (t, c) in enumerate(zip(time_ids, camera_ids)): + frame_names_map[t, c] = _frame_names[i] + + return frame_names_map, time_ids, camera_ids + + +def _rescale_depth(depth_raw: np.ndarray, cam: Dict) -> np.ndarray: + """Depth rescale function from DyCheck. + + Args: + depth: A numpy ndarray of the raw depth + cam: Dict of the camera + + Returns: + A numpy ndarray of the processed depth + """ + xx, yy = np.meshgrid(np.arange(cam["width"], dtype=np.float32), np.arange(cam["height"], dtype=np.float32)) + pixels = np.stack([xx, yy], axis=-1) + batch_shape = pixels.shape[:-1] + pixels = np.reshape(pixels, (-1, 2)) + y = (pixels[..., 1] - cam["cy"]) / cam["fy"] + x = (pixels[..., 0] - cam["cx"]) / cam["fx"] + # x = (pixels[..., 0] - self.principal_point_x - y * self.skew) / self.scale_factor_x + # assume skew = 0 + viewdirs = np.stack([x, y, np.ones_like(x)], axis=-1) + local_viewdirs = viewdirs / np.linalg.norm(viewdirs, axis=-1, keepdims=True) + viewdirs = (cam["camera_to_worlds"][:3, :3] @ local_viewdirs[..., None])[..., 0] + viewdirs /= np.linalg.norm(viewdirs, axis=-1, keepdims=True) + viewdirs = viewdirs.reshape((*batch_shape, 3)) + cosa = viewdirs @ (cam["camera_to_worlds"][:, 2]) + depth = depth_raw / cosa[..., None] + return depth + + +@dataclass +class DycheckDataParserConfig(DataParserConfig): + """Dycheck (https://arxiv.org/abs/2210.13445) dataset parser config""" + + _target: Type = field(default_factory=lambda: Dycheck) + """target class to instantiate""" + data: Path = Path("data/iphone/mochi-high-five") + """Directory specifying location of data.""" + scale_factor: float = 5.0 + """How much to scale the camera origins by.""" + alpha_color: str = "white" + """alpha color of background""" + downscale_factor: int = 1 + """How much to downscale images.""" + scene_box_bound: float = 1.5 + """Boundary of scene box.""" + + +@dataclass +class Dycheck(DataParser): + """Dycheck (https://arxiv.org/abs/2210.13445) Dataset `iphone` subset""" + + config: DycheckDataParserConfig + + def __init__(self, config: DycheckDataParserConfig): + super().__init__(config=config) + self.data: Path = config.data + self.scale_factor: float = config.scale_factor + self.alpha_color = config.alpha_color + # load extra info from "extra.json" + extra_path = self.data / "extra.json" + extra_dict = load_from_json(extra_path) + self._factor = extra_dict["factor"] + self._fps = extra_dict["fps"] + self._bbox = np.array(extra_dict["bbox"], dtype=np.float32) + self._lookat = np.array(extra_dict["lookat"], dtype=np.float32) + self._up = np.array(extra_dict["up"], dtype=np.float32) + self._center, self._scale, self._near, self._far = _load_scene_info(self.data) + self._frame_names_map, self._time_ids, self._camera_ids = _load_metadata_info(self.data) + + 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 + splits_dir = self.data / "splits" + + # scale the scene to fill the aabb bbox + sf = self.config.scene_box_bound / 4 / (self._scale * self._far) + # CONSOLE.print(f"scale factor changed from {self.config.scale_factor} to {sf}") + self.config.scale_factor = sf + + if not (splits_dir / f"{split}.json").exists(): + CONSOLE.print(f"split {split} not found, using split train") + split = "train" + split_dict = load_from_json(splits_dir / f"{split}.json") + frame_names = np.array(split_dict["frame_names"]) + time_ids = np.array(split_dict["time_ids"]) + if split != "train": + CONSOLE.print(f"split {split} is empty, using the 1st training image") + split_dict = load_from_json(splits_dir / "train.json") + frame_names = np.array(split_dict["frame_names"])[[0]] + time_ids = np.array(split_dict["time_ids"])[[0]] + + image_filenames, depth_filenames, cams = self.process_frames(frame_names, time_ids) + + scene_box = SceneBox( + aabb=torch.tensor( + [[-self.config.scene_box_bound] * 3, [self.config.scene_box_bound] * 3], dtype=torch.float32 + ) + ) + cam_dict = {} + for k in cams[0].keys(): + cam_dict[k] = torch.stack([torch.as_tensor(c[k]) for c in cams], dim=0) + cameras = Cameras(camera_type=CameraType.PERSPECTIVE, **cam_dict) + + scale = self._scale * self.config.scale_factor + dataparser_outputs = DataparserOutputs( + image_filenames=image_filenames, + cameras=cameras, + alpha_color=alpha_color_tensor, + scene_box=scene_box, + metadata={ + "depth_filenames": depth_filenames, + "depth_unit_scale_factor": scale, + "scale": scale, + "near": self._near * scale, + "far": self._far * scale, + }, + ) + + return dataparser_outputs + + def process_frames(self, frame_names: List[str], time_ids: np.ndarray) -> Tuple[List, List, List]: + """Read cameras and filenames from the name list. + + Args: + frame_names: list of file names. + time_ids: time id of each frame. + + Returns: + A list of camera, each entry is a dict of the camera. + """ + image_filenames, depth_filenames = [], [] + cams = [] + for idx, frame in enumerate(frame_names): + image_filenames.append(self.data / f"rgb/{self.config.downscale_factor}x/{frame}.png") + depth_filenames.append(self.data / f"processed_depth/{self.config.downscale_factor}x/{frame}.npy") + cam_json = load_from_json(self.data / f"camera/{frame}.json") + c2w = torch.as_tensor(cam_json["orientation"]).T + position = torch.as_tensor(cam_json["position"]) + position -= self._center # some scenes look weird (wheel) + position *= self._scale * self.config.scale_factor + pose = torch.zeros([3, 4]) + pose[:3, :3] = c2w + pose[:3, 3] = position + # from opencv coord to opengl coord (used by nerfstudio) + pose[0:3, 1:3] *= -1 # switch cam coord x,y + pose = pose[[1, 0, 2], :] # switch world x,y + pose[2, :] *= -1 # invert world z + # for aabb bbox usage + pose = pose[[1, 2, 0], :] # switch world xyz to zxy + cams.append( + { + "camera_to_worlds": pose, + "fx": cam_json["focal_length"] / self.config.downscale_factor, + "fy": cam_json["focal_length"] * cam_json["pixel_aspect_ratio"] / self.config.downscale_factor, + "cx": cam_json["principal_point"][0] / self.config.downscale_factor, + "cy": cam_json["principal_point"][1] / self.config.downscale_factor, + "height": cam_json["image_size"][1] // self.config.downscale_factor, + "width": cam_json["image_size"][0] // self.config.downscale_factor, + "times": torch.as_tensor(time_ids[idx] / self._time_ids.max()).float(), + } + ) + + d = self.config.downscale_factor + if not image_filenames[0].exists(): + CONSOLE.print(f"downscale factor {d}x not exist, converting") + ori_h, ori_w = cv2.imread(str(self.data / f"rgb/1x/{frame_names[0]}.png")).shape[:2] + (self.data / f"rgb/{d}x").mkdir(exist_ok=True) + h, w = ori_h // d, ori_w // d + for frame in frame_names: + cv2.imwrite( + str(self.data / f"rgb/{d}x/{frame}.png"), + cv2.resize(cv2.imread(str(self.data / f"rgb/1x/{frame}.png")), [h, w]), + ) + CONSOLE.print("finished") + + if not depth_filenames[0].exists(): + CONSOLE.print(f"processed depth downscale factor {d}x not exist, converting") + (self.data / f"processed_depth/{d}x").mkdir(exist_ok=True, parents=True) + for idx, frame in enumerate(frame_names): + depth = np.load(self.data / f"depth/1x/{frame}.npy") + mask = rescale((depth != 0).astype(np.uint8) * 255, 1 / d, cv2.INTER_AREA) + depth = rescale(depth, 1 / d, cv2.INTER_AREA) + depth[mask != 255] = 0 + depth = _rescale_depth(depth, cams[idx]) + np.save(str(self.data / f"processed_depth/{d}x/{frame}.npy"), depth) + CONSOLE.print("finished") + + return image_filenames, depth_filenames, cams diff --git a/nerfstudio/data/utils/data_utils.py b/nerfstudio/data/utils/data_utils.py index c4b21febd..f04fb368c 100644 --- a/nerfstudio/data/utils/data_utils.py +++ b/nerfstudio/data/utils/data_utils.py @@ -64,8 +64,8 @@ def get_depth_image_from_path( interpolation: int = cv2.INTER_NEAREST, ) -> torch.Tensor: """Loads, rescales and resizes depth images. + Filepath points to a 16-bit or 32-bit depth image, or a numpy array `*.npy`. - Assumes filepath points to a 16-bit or 32-bit depth image. Args: filepath: Path to depth image. height: Target depth image height. @@ -76,7 +76,11 @@ def get_depth_image_from_path( Returns: Depth image torch tensor with shape [width, height, 1]. """ - image = cv2.imread(str(filepath.absolute()), cv2.IMREAD_ANYDEPTH) - image = image.astype(np.float64) * scale_factor - image = cv2.resize(image, (width, height), interpolation=interpolation) + if filepath.suffix == ".npy": + image = np.load(filepath) * scale_factor + image = cv2.resize(image, (width, height), interpolation=interpolation) + else: + image = cv2.imread(str(filepath.absolute()), cv2.IMREAD_ANYDEPTH) + image = image.astype(np.float64) * scale_factor + image = cv2.resize(image, (width, height), interpolation=interpolation) return torch.from_numpy(image[:, :, np.newaxis])