From 9a01738b0e472dc6b93578a765edbcbb371a03e5 Mon Sep 17 00:00:00 2001 From: James Gardner Date: Fri, 17 Mar 2023 22:53:19 +0000 Subject: [PATCH] Adding NeRF-OSR Dataset (#1626) * Added NeRFOSR Dataparser * updating nerfosr dataparser * fix import order in datamanager * updating nerfosr dataparser * update nerfosr dataparser --- .../data/datamanagers/base_datamanager.py | 2 + .../data/dataparsers/nerfosr_dataparser.py | 235 ++++++++++++++++++ scripts/downloads/download_data.py | 60 +++++ 3 files changed, 297 insertions(+) create mode 100644 nerfstudio/data/dataparsers/nerfosr_dataparser.py diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 7ff2f9aaf..dc96e22dc 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -47,6 +47,7 @@ from nerfstudio.data.dataparsers.instant_ngp_dataparser import ( InstantNGPDataParserConfig, ) from nerfstudio.data.dataparsers.minimal_dataparser import MinimalDataParserConfig +from nerfstudio.data.dataparsers.nerfosr_dataparser import NeRFOSRDataParserConfig from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserConfig from nerfstudio.data.dataparsers.phototourism_dataparser import ( @@ -87,6 +88,7 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of "dycheck-data": DycheckDataParserConfig(), "scannet-data": ScanNetDataParserConfig(), "sdfstudio-data": SDFStudioDataParserConfig(), + "nerfosr-data": NeRFOSRDataParserConfig(), "sitcoms3d-data": Sitcoms3DDataParserConfig(), }, prefix_names=False, # Omit prefixes in subcommands themselves. diff --git a/nerfstudio/data/dataparsers/nerfosr_dataparser.py b/nerfstudio/data/dataparsers/nerfosr_dataparser.py new file mode 100644 index 000000000..6b24ca957 --- /dev/null +++ b/nerfstudio/data/dataparsers/nerfosr_dataparser.py @@ -0,0 +1,235 @@ +# 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 NeRF-OSR datasets + + Presented in the paper: https://4dqv.mpi-inf.mpg.de/NeRF-OSR/ + +""" + +from __future__ import annotations + +import glob +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import List, Tuple, Type + +import numpy as np +import torch +from rich.console import Console +from typing_extensions import Literal + +from nerfstudio.cameras import camera_utils +from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.data.dataparsers.base_dataparser import ( + DataParser, + DataParserConfig, + DataparserOutputs, +) +from nerfstudio.data.scene_box import SceneBox + +CONSOLE = Console(width=120) + + +def find_files(directory: str, exts: List[str]): + """Find all files in a directory that have a certain file extension. + + Args: + directory : The directory to search for files. + exts : A list of file extensions to search for. Each file extension should be in the form '*.ext'. + + Returns: + A list of file paths for all the files that were found. The list is sorted alphabetically. + """ + if os.path.isdir(directory): + # types should be ['*.png', '*.jpg', '*.JPG', '*.PNG'] + files_grabbed = [] + for ext in exts: + files_grabbed.extend(glob.glob(os.path.join(directory, ext))) + if len(files_grabbed) > 0: + files_grabbed = sorted(files_grabbed) + return files_grabbed + return [] + + +def _parse_osm_txt(filename: str): + """Parse a text file containing numbers and return a 4x4 numpy array of float32 values. + + Args: + filename : a file containing numbers in a 4x4 matrix. + + Returns: + A numpy array of shape [4, 4] containing the numbers from the file. + """ + assert os.path.isfile(filename) + with open(filename, encoding="UTF-8") as f: + nums = f.read().split() + return np.array([float(x) for x in nums]).reshape([4, 4]).astype(np.float32) + + +def get_camera_params(scene_dir: str, split: str) -> Tuple[torch.Tensor, torch.Tensor, int]: + """Load camera intrinsic and extrinsic parameters for a given scene split. + + Args" + scene_dir : The directory containing the scene data. + split : The split for which to load the camera parameters. Either "train", "val", or "test". + + Returns + A tuple containing the intrinsic parameters (as a torch.Tensor of shape [N, 4, 4]), + the camera-to-world matrices (as a torch.Tensor of shape [N, 4, 4]), and the number of cameras (N). + """ + split = "validation" if split == "val" else split + split_dir = f"{scene_dir}/{split}" + + # camera parameters files + intrinsics_files = find_files(f"{split_dir}/intrinsics", exts=["*.txt"]) + pose_files = find_files(f"{split_dir}/pose", exts=["*.txt"]) + + num_cams = len(pose_files) + + intrinsics = [] + camera_to_worlds = [] + for i in range(num_cams): + intrinsics.append(_parse_osm_txt(intrinsics_files[i])) + + pose = _parse_osm_txt(pose_files[i]) + + # convert from COLMAP/OpenCV to nerfstudio camera (OpenGL/Blender) + pose[0:3, 1:3] *= -1 + + camera_to_worlds.append(pose) + + intrinsics = torch.from_numpy(np.stack(intrinsics).astype(np.float32)) # [N, 4, 4] + camera_to_worlds = torch.from_numpy(np.stack(camera_to_worlds).astype(np.float32)) # [N, 4, 4] + + return intrinsics, camera_to_worlds, num_cams + + +@dataclass +class NeRFOSRDataParserConfig(DataParserConfig): + """Nerfstudio dataset config""" + + _target: Type = field(default_factory=lambda: NeRFOSR) + """target class to instantiate""" + data: Path = Path("data/NeRF-OSR/Data/") + """Directory specifying location of data.""" + scene: str = "stjacob" + """Which scene to load""" + scene_scale: float = 1.0 + """How much to scale the region of interest by.""" + scale_factor: float = 1.0 + """How much to scale the camera origins by.""" + use_masks: bool = False + """Whether to use masks.""" + orientation_method: Literal["pca", "up", "vertical", "none"] = "vertical" + """The method to use for orientation.""" + center_method: Literal["poses", "focus", "none"] = "focus" + """The method to use for centering.""" + auto_scale_poses: bool = True + """Whether to automatically scale the poses to fit in +/- 1 bounding box.""" + + +@dataclass +class NeRFOSR(DataParser): + """NeRFOSR Dataparser + Presented in the paper: https://4dqv.mpi-inf.mpg.de/NeRF-OSR/ + + Some of this code comes from https://github.com/r00tman/NeRF-OSR/blob/main/data_loader_split.py + + Source data convention is: + camera coordinate system: x-->right, y-->down, z-->scene (opencv/colmap convention) + poses is camera-to-world + masks are 0 for dynamic content, 255 for static content + """ + + config: NeRFOSRDataParserConfig + + def _generate_dataparser_outputs(self, split="train"): + data = self.config.data + scene = self.config.scene + split = "validation" if split == "val" else split + + if scene == "trevi": + scene_dir = f"{data}/{scene}/final_clean" + split_dir = f"{data}/{scene}/final_clean/{split}" + else: + scene_dir = f"{data}/{scene}/final" + split_dir = f"{data}/{scene}/final/{split}" + + # get all split cam params + intrinsics_train, camera_to_worlds_train, n_train = get_camera_params(scene_dir, "train") + intrinsics_val, camera_to_worlds_val, n_val = get_camera_params(scene_dir, "val") + intrinsics_test, camera_to_worlds_test, _ = get_camera_params(scene_dir, "test") + + # combine all cam params + intrinsics = torch.cat([intrinsics_train, intrinsics_val, intrinsics_test], dim=0) + camera_to_worlds = torch.cat([camera_to_worlds_train, camera_to_worlds_val, camera_to_worlds_test], dim=0) + + camera_to_worlds, _ = camera_utils.auto_orient_and_center_poses( + camera_to_worlds, + method=self.config.orientation_method, + center_method=self.config.center_method, + ) + + # Scale poses + scale_factor = 1.0 + if self.config.auto_scale_poses: + scale_factor /= torch.max(torch.abs(camera_to_worlds[:, :3, 3])) + + camera_to_worlds[:, :3, 3] *= scale_factor * self.config.scale_factor + + if split == "train": + camera_to_worlds = camera_to_worlds[:n_train] + intrinsics = intrinsics[:n_train] + elif split == "val": + camera_to_worlds = camera_to_worlds[n_train : n_train + n_val] + intrinsics = intrinsics[n_train : n_train + n_val] + elif split == "test": + camera_to_worlds = camera_to_worlds[n_train + n_val :] + intrinsics = intrinsics[n_train + n_val :] + + cameras = Cameras( + camera_to_worlds=camera_to_worlds[:, :3, :4], + fx=intrinsics[:, 0, 0], + fy=intrinsics[:, 1, 1], + cx=intrinsics[:, 0, 2], + cy=intrinsics[:, 1, 2], + camera_type=CameraType.PERSPECTIVE, + ) + + # in x,y,z order + # assumes that the scene is centered at the origin + aabb_scale = self.config.scene_scale + scene_box = SceneBox( + aabb=torch.tensor( + [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32 + ) + ) + + # --- images --- + image_filenames = find_files(f"{split_dir}/rgb", exts=["*.png", "*.jpg", "*.JPG", "*.PNG"]) + + # --- masks --- + mask_filenames = [] + if self.config.use_masks: + mask_filenames = find_files(f"{split_dir}/mask", exts=["*.png", "*.jpg", "*.JPG", "*.PNG"]) + + dataparser_outputs = DataparserOutputs( + image_filenames=image_filenames, + cameras=cameras, + scene_box=scene_box, + mask_filenames=mask_filenames if len(mask_filenames) > 0 else None, + dataparser_scale=self.config.scale_factor, + ) + return dataparser_outputs diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py index 3339b50bc..a92798bbb 100644 --- a/scripts/downloads/download_data.py +++ b/scripts/downloads/download_data.py @@ -359,6 +359,65 @@ class SDFstudioDemoDownload(DatasetDownload): os.remove(download_path) +# pylint: disable=line-too-long +nerfosr_downloads = { + "europa": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=europa&downloadStartSecret=0k2r95c1fdej", + "lk2": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=lk2&downloadStartSecret=w8kuvjzmchc", + "lwp": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=lwp&downloadStartSecret=gtnc4vmkcjq", + "rathaus": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=rathaus&downloadStartSecret=7372aewy6rr", + "schloss": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=schloss&downloadStartSecret=y8t00nqx0h", + "st": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=st&downloadStartSecret=kl9ptuxe8v", + "stjacob": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=stjacob&downloadStartSecret=sntsim6ebvm", + "stjohann": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=stjohann&downloadStartSecret=g80ug1fsbmh", + "trevi": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=trevi&downloadStartSecret=ot1483bigjm", + "all": None, +} + + +NeRFOSRCaptureName = tyro.extras.literal_type_from_choices(nerfosr_downloads.keys()) + + +@dataclass +class NeRFOSRDownload(DatasetDownload): + """Download the NeRF-OSR dataset.""" + + capture_name: NeRFOSRCaptureName = "europa" + + def download(self, save_dir: Path): + """Download the NeRF-OSR dataset: https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk""" + + if self.capture_name == "all": + for capture_name in nerfosr_downloads: + if capture_name != "all": + NeRFOSRDownload(capture_name=capture_name).download(save_dir) + return + + assert ( + self.capture_name in nerfosr_downloads + ), f"Capture name {self.capture_name} not found in {nerfosr_downloads.keys()}" + url = nerfosr_downloads[self.capture_name] + target_path = str(save_dir / f"NeRF-OSR/Data/{self.capture_name}") + os.makedirs(target_path, exist_ok=True) + download_path = Path(f"{target_path}.zip") + tmp_path = str(save_dir / ".temp") + shutil.rmtree(tmp_path, ignore_errors=True) + os.makedirs(tmp_path, exist_ok=True) + + os.system(f"curl -L '{url}' > {download_path}") + + # Extract the zip file + with zipfile.ZipFile(download_path, "r") as zip_ref: + zip_ref.extractall(tmp_path) + + inner_folders = os.listdir(tmp_path) + assert len(inner_folders) == 1, "There is more than one folder inside this zip file." + folder = os.path.join(tmp_path, inner_folders[0]) + shutil.rmtree(target_path) + shutil.move(folder, target_path) + shutil.rmtree(tmp_path) + os.remove(download_path) + + Commands = Union[ Annotated[BlenderDownload, tyro.conf.subcommand(name="blender")], Annotated[Sitcoms3DDownload, tyro.conf.subcommand(name="sitcoms3d")], @@ -367,6 +426,7 @@ Commands = Union[ Annotated[DNerfDownload, tyro.conf.subcommand(name="dnerf")], Annotated[PhototourismDownload, tyro.conf.subcommand(name="phototourism")], Annotated[SDFstudioDemoDownload, tyro.conf.subcommand(name="sdfstudio")], + Annotated[NeRFOSRDownload, tyro.conf.subcommand(name="nerfosr")], ]