merge

nerfstudio/data/datamanagers/base_datamanager.py

@@ -45,6 +45,7 @@ from nerfstudio.data.dataparsers.instant_ngp_dataparser import (
     InstantNGPDataParserConfig,
 )
 from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
+from nerfstudio.data.dataparsers.mipnerf360_dataparser import Mipnerf360DataParserConfig
 from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserConfig
 from nerfstudio.data.dataparsers.phototourism_dataparser import (
     PhototourismDataParserConfig,
@@ -63,6 +64,7 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[  # Omit prefixes of
     tyro.extras.subcommand_type_from_defaults(
         {
             "nerfstudio-data": NerfstudioDataParserConfig(),
+            "mipnerf360-data": Mipnerf360DataParserConfig(),
             "blender-data": BlenderDataParserConfig(),
             "friends-data": FriendsDataParserConfig(),
             "instant-ngp-data": InstantNGPDataParserConfig(),

nerfstudio/data/dataparsers/mipnerf360_dataparser.py

@@ -0,0 +1,313 @@
+# 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 mipnerf360 datasets. """
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from pathlib import Path, PurePath
+from typing import Optional, Type
+
+import numpy as np
+import torch
+from PIL import Image
+from rich.console import Console
+from typing_extensions import Literal
+
+from nerfstudio.cameras import camera_utils
+from nerfstudio.cameras.cameras import CAMERA_MODEL_TO_TYPE, Cameras, CameraType
+from nerfstudio.data.dataparsers.base_dataparser import (
+    DataParser,
+    DataParserConfig,
+    DataparserOutputs,
+)
+from nerfstudio.data.scene_box import SceneBox
+from nerfstudio.utils.io import load_from_json
+
+CONSOLE = Console(width=120)
+MAX_AUTO_RESOLUTION = 1600
+
+
+@dataclass
+class Mipnerf360DataParserConfig(DataParserConfig):
+    """Mipnerf360 dataset config"""
+
+    _target: Type = field(default_factory=lambda: Mipnerf360)
+    """target class to instantiate"""
+    data: Path = Path("data/nerfstudio-data-mipnerf360/garden")
+    """Directory or explicit json file path specifying location of data."""
+    scale_factor: float = 1.0
+    """How much to scale the camera origins by."""
+    downscale_factor: Optional[int] = None
+    """How much to downscale images. If not set, images are chosen such that the max dimension is <1600px."""
+    scene_scale: float = 1.0
+    """How much to scale the region of interest by."""
+    orientation_method: Literal["pca", "up", "none"] = "up"
+    """The method to use for orientation."""
+    center_poses: bool = True
+    """Whether to center the poses."""
+    auto_scale_poses: bool = True
+    """Whether to automatically scale the poses to fit in +/- 1 bounding box."""
+    train_split_percentage: float = 0.9
+    """The percent of images to use for training. The remaining images are for eval."""
+
+
+@dataclass
+class Mipnerf360(DataParser):
+    """Mipnerf360 DatasetParser"""
+
+    config: NerfstudioDataParserConfig
+    downscale_factor: Optional[int] = None
+
+    def _generate_dataparser_outputs(self, split="train"):
+        # pylint: disable=too-many-statements
+
+        if self.config.data.suffix == ".json":
+            meta = load_from_json(self.config.data)
+            data_dir = self.config.data.parent
+        else:
+            meta = load_from_json(self.config.data / "transforms.json")
+            data_dir = self.config.data
+
+        image_filenames = []
+        mask_filenames = []
+        poses = []
+        num_skipped_image_filenames = 0
+
+        fx_fixed = "fl_x" in meta
+        fy_fixed = "fl_y" in meta
+        cx_fixed = "cx" in meta
+        cy_fixed = "cy" in meta
+        height_fixed = "h" in meta
+        width_fixed = "w" in meta
+        distort_fixed = False
+        for distort_key in ["k1", "k2", "k3", "p1", "p2"]:
+            if distort_key in meta:
+                distort_fixed = True
+                break
+        fx = []
+        fy = []
+        cx = []
+        cy = []
+        height = []
+        width = []
+        distort = []
+
+        # sort the frames by fname
+        fnames = []
+        for frame in meta["frames"]:
+            filepath = PurePath(frame["file_path"])
+            fname = self._get_fname(filepath, data_dir)
+            fnames.append(fname)
+        inds = np.argsort(fnames)
+        frames = [meta["frames"][ind] for ind in inds]
+
+        for frame in frames:
+            filepath = PurePath(frame["file_path"])
+            fname = self._get_fname(filepath, data_dir)
+            if not fname.exists():
+                num_skipped_image_filenames += 1
+                continue
+
+            if not fx_fixed:
+                assert "fl_x" in frame, "fx not specified in frame"
+                fx.append(float(frame["fl_x"]))
+            if not fy_fixed:
+                assert "fl_y" in frame, "fy not specified in frame"
+                fy.append(float(frame["fl_y"]))
+            if not cx_fixed:
+                assert "cx" in frame, "cx not specified in frame"
+                cx.append(float(frame["cx"]))
+            if not cy_fixed:
+                assert "cy" in frame, "cy not specified in frame"
+                cy.append(float(frame["cy"]))
+            if not height_fixed:
+                assert "h" in frame, "height not specified in frame"
+                height.append(int(frame["h"]))
+            if not width_fixed:
+                assert "w" in frame, "width not specified in frame"
+                width.append(int(frame["w"]))
+            if not distort_fixed:
+                distort.append(
+                    camera_utils.get_distortion_params(
+                        k1=float(meta["k1"]) if "k1" in meta else 0.0,
+                        k2=float(meta["k2"]) if "k2" in meta else 0.0,
+                        k3=float(meta["k3"]) if "k3" in meta else 0.0,
+                        k4=float(meta["k4"]) if "k4" in meta else 0.0,
+                        p1=float(meta["p1"]) if "p1" in meta else 0.0,
+                        p2=float(meta["p2"]) if "p2" in meta else 0.0,
+                    )
+                )
+
+            image_filenames.append(fname)
+            poses.append(np.array(frame["transform_matrix"]))
+            if "mask_path" in frame:
+                mask_filepath = PurePath(frame["mask_path"])
+                mask_fname = self._get_fname(
+                    mask_filepath,
+                    data_dir,
+                    downsample_folder_prefix="masks_",
+                )
+                mask_filenames.append(mask_fname)
+        if num_skipped_image_filenames >= 0:
+            CONSOLE.log(f"Skipping {num_skipped_image_filenames} files in dataset split {split}.")
+        assert (
+            len(image_filenames) != 0
+        ), """
+        No image files found. 
+        You should check the file_paths in the transforms.json file to make sure they are correct.
+        """
+        assert len(mask_filenames) == 0 or (
+            len(mask_filenames) == len(image_filenames)
+        ), """
+        Different number of image and mask filenames.
+        You should check that mask_path is specified for every frame (or zero frames) in transforms.json.
+        """
+
+        # sort the filenames according to image_filenames
+
+
+        # filter image_filenames and poses based on train/eval split percentage
+        num_images = len(image_filenames)
+        
+        # mipnerf360 code...
+        # TODO: add parameters for this...
+        all_indices = np.arange(num_images)
+        train_indices = all_indices[all_indices % 8 != 0]
+        eval_indices = all_indices[all_indices % 8 == 0]
+        i_train = train_indices
+        i_eval = eval_indices
+        
+        if split == "train":
+            indices = i_train
+        elif split in ["val", "test"]:
+            indices = i_eval
+        else:
+            raise ValueError(f"Unknown dataparser split {split}")
+
+        if "orientation_override" in meta:
+            orientation_method = meta["orientation_override"]
+            CONSOLE.log(f"[yellow] Dataset is overriding orientation method to {orientation_method}")
+        else:
+            orientation_method = self.config.orientation_method
+
+        poses = torch.from_numpy(np.array(poses).astype(np.float32))
+        poses, transform_matrix = camera_utils.auto_orient_and_center_poses(
+            poses,
+            method=orientation_method,
+            center_poses=self.config.center_poses,
+        )
+
+        # Scale poses
+        scale_factor = 1.0
+        if self.config.auto_scale_poses:
+            scale_factor /= float(torch.max(torch.abs(poses[:, :3, 3])))
+        scale_factor *= self.config.scale_factor
+
+        poses[:, :3, 3] *= scale_factor
+
+        # Choose image_filenames and poses based on split, but after auto orient and scaling the poses.
+        image_filenames = [image_filenames[i] for i in indices]
+        mask_filenames = [mask_filenames[i] for i in indices] if len(mask_filenames) > 0 else []
+        poses = poses[indices]
+
+        # 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
+            )
+        )
+
+        if "camera_model" in meta:
+            camera_type = CAMERA_MODEL_TO_TYPE[meta["camera_model"]]
+        else:
+            camera_type = CameraType.PERSPECTIVE
+
+        idx_tensor = torch.tensor(indices, dtype=torch.long)
+        fx = float(meta["fl_x"]) if fx_fixed else torch.tensor(fx, dtype=torch.float32)[idx_tensor]
+        fy = float(meta["fl_y"]) if fy_fixed else torch.tensor(fy, dtype=torch.float32)[idx_tensor]
+        cx = float(meta["cx"]) if cx_fixed else torch.tensor(cx, dtype=torch.float32)[idx_tensor]
+        cy = float(meta["cy"]) if cy_fixed else torch.tensor(cy, dtype=torch.float32)[idx_tensor]
+        height = int(meta["h"]) if height_fixed else torch.tensor(height, dtype=torch.int32)[idx_tensor]
+        width = int(meta["w"]) if width_fixed else torch.tensor(width, dtype=torch.int32)[idx_tensor]
+        if distort_fixed:
+            distortion_params = camera_utils.get_distortion_params(
+                k1=float(meta["k1"]) if "k1" in meta else 0.0,
+                k2=float(meta["k2"]) if "k2" in meta else 0.0,
+                k3=float(meta["k3"]) if "k3" in meta else 0.0,
+                k4=float(meta["k4"]) if "k4" in meta else 0.0,
+                p1=float(meta["p1"]) if "p1" in meta else 0.0,
+                p2=float(meta["p2"]) if "p2" in meta else 0.0,
+            )
+        else:
+            distortion_params = torch.stack(distort, dim=0)[idx_tensor]
+
+        cameras = Cameras(
+            fx=fx,
+            fy=fy,
+            cx=cx,
+            cy=cy,
+            distortion_params=distortion_params,
+            height=height,
+            width=width,
+            camera_to_worlds=poses[:, :3, :4],
+            camera_type=camera_type,
+        )
+
+        assert self.downscale_factor is not None
+        cameras.rescale_output_resolution(scaling_factor=1.0 / self.downscale_factor)
+
+        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=scale_factor,
+            dataparser_transform=transform_matrix,
+        )
+        return dataparser_outputs
+
+    def _get_fname(self, filepath: PurePath, data_dir: PurePath, downsample_folder_prefix="images_") -> Path:
+        """Get the filename of the image file.
+        downsample_folder_prefix can be used to point to auxillary image data, e.g. masks
+
+        filepath: the base file name of the transformations.
+        data_dir: the directory of the data that contains the transform file
+        downsample_folder_prefix: prefix of the newly generated downsampled images
+        """
+
+        if self.downscale_factor is None:
+            if self.config.downscale_factor is None:
+                test_img = Image.open(data_dir / filepath)
+                h, w = test_img.size
+                max_res = max(h, w)
+                df = 0
+                while True:
+                    if (max_res / 2 ** (df)) < MAX_AUTO_RESOLUTION:
+                        break
+                    if not (data_dir / f"{downsample_folder_prefix}{2**(df+1)}" / filepath.name).exists():
+                        break
+                    df += 1
+
+                self.downscale_factor = 2**df
+                CONSOLE.log(f"Auto image downscale factor of {self.downscale_factor}")
+            else:
+                self.downscale_factor = self.config.downscale_factor
+
+        if self.downscale_factor > 1:
+            return data_dir / f"{downsample_folder_prefix}{self.downscale_factor}" / filepath.name
+        return data_dir / filepath

nerfstudio/data/dataparsers/nerfstudio_dataparser.py

@@ -105,16 +105,7 @@ class Nerfstudio(DataParser):
         width = []
         distort = []
 
-        # sort the frames by fname
-        fnames = []
         for frame in meta["frames"]:
-            filepath = PurePath(frame["file_path"])
-            fname = self._get_fname(filepath, data_dir)
-            fnames.append(fname)
-        inds = np.argsort(fnames)
-        frames = [meta["frames"][ind] for ind in inds]
-
-        for frame in frames:
             filepath = PurePath(frame["file_path"])
             fname = self._get_fname(filepath, data_dir)
             if not fname.exists():
@@ -176,39 +167,16 @@ class Nerfstudio(DataParser):
         You should check that mask_path is specified for every frame (or zero frames) in transforms.json.
         """
 
-        # sort the filenames according to image_filenames
-
-
         # filter image_filenames and poses based on train/eval split percentage
         num_images = len(image_filenames)
         num_train_images = math.ceil(num_images * self.config.train_split_percentage)
         num_eval_images = num_images - num_train_images
         i_all = np.arange(num_images)
-        i_eval = np.linspace(
-            0, num_images - 1, num_eval_images, dtype=int
+        i_train = np.linspace(
+            0, num_images - 1, num_train_images, dtype=int
         )  # equally spaced training images starting and ending at 0 and num_images-1
-        i_train = np.setdiff1d(i_all, i_eval)  # eval images are the remaining images
-        assert len(i_train) == num_train_images
+        i_eval = np.setdiff1d(i_all, i_train)  # eval images are the remaining images
         assert len(i_eval) == num_eval_images
-
-        # # print(image_filenames)
-        # inds = np.argsort(image_filenames)
-        # image_filenames = [image_filenames[i] for i in inds]
-        # poses = [poses[i] for i in inds]
-        # if len(mask_filenames) != 0:
-        #     mask_filenames = [mask_filenames[i] for i in inds]
-        
-        # mipnerf360 code...
-        all_indices = np.arange(num_images)
-        train_indices = all_indices[all_indices % 8 != 0]
-        eval_indices = all_indices[all_indices % 8 == 0]
-        i_train = train_indices
-        i_eval = eval_indices
-
-        print(image_filenames)
-        print(i_train)
-        print(i_eval)
-        
         if split == "train":
             indices = i_train
         elif split in ["val", "test"]:
@@ -216,24 +184,6 @@ class Nerfstudio(DataParser):
         else:
             raise ValueError(f"Unknown dataparser split {split}")
 
-        # # print("\nhere\n")
-        # # print(i_train)
-        # # print(i_eval)
-
-        # # mipnerf360 code...
-        # eval_interval = 1.0 - self.config.train_split_percentage
-        # all_indices = np.arange(num_images)
-        # train_indices = all_indices[all_indices % 8 != 0]
-        # eval_indices = all_indices[all_indices % 8 == 0]
-        # # print("\nhere\n")
-        # # print(train_indices)
-        # # print(eval_indices)
-
-        # # import sys
-        # # sys.exit()
-
-        # indices
-
         if "orientation_override" in meta:
             orientation_method = meta["orientation_override"]
             CONSOLE.log(f"[yellow] Dataset is overriding orientation method to {orientation_method}")

scripts/benchmarking/benchmark_nerfstudio_paper.py

@@ -19,7 +19,7 @@ mipnerf360_table_rows = [
     (
         "nerfacto-w/o-pose-app",
         "nerfacto",
-        "--pipeline.eval_optimize_cameras False --pipeline.eval_optimize_appearance False --pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-appearance-embedding False nerfstudio-data --downscale-factor 4 --train-split-percentage 0.875",
+        "--pipeline.eval_optimize_cameras False --pipeline.eval_optimize_appearance False --pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-appearance-embedding False nerfstudio-data --downscale-factor 4 --train-split-percentage 0.875 mipnerf360-data",
     ),
 ]
 
@@ -94,12 +94,6 @@ def main(capture_names, table_rows, data_path: Path = Path("data/nerfstudio")):
     # make a list of all the jobs that need to be fun
     jobs = []
     for capture_name in capture_names:
-
-        # extra_string = ""
-        # if str(data_path) == "data/nerfstudio-data-mipnerf360" and capture_name == "stump":
-        #     extra_string = "--pipeline.model.near_plane 0.15"
-
-        # extra_string = "--pipeline.model.near_plane 0.15"
         
         for table_row_name, method, table_row_command in table_rows:
             command = " ".join(
@@ -154,5 +148,5 @@ def main(capture_names, table_rows, data_path: Path = Path("data/nerfstudio")):
 
 if __name__ == "__main__":
     # pass
-    main(mipnerf360_capture_names, mipnerf360_table_rows, data_path=Path("data/nerfstudio-data-mipnerf360"))
-    # main(ablations_capture_names, ablations_table_rows)
+    # main(mipnerf360_capture_names, mipnerf360_table_rows, data_path=Path("data/nerfstudio-data-mipnerf360"))
+    main(ablations_capture_names, ablations_table_rows)