Adding NeRF-OSR Dataset (#1626)

* Added NeRFOSR Dataparser

* updating nerfosr dataparser

* fix import order in datamanager

* updating nerfosr dataparser

* update nerfosr dataparser

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.

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

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")],
 ]