diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 555d618e4..d4742efd4 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/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(), diff --git a/nerfstudio/data/dataparsers/mipnerf360_dataparser.py b/nerfstudio/data/dataparsers/mipnerf360_dataparser.py new file mode 100644 index 000000000..d3884108f --- /dev/null +++ b/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 diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index b13362b9d..e311a08cc 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/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}") diff --git a/scripts/benchmarking/benchmark_nerfstudio_paper.py b/scripts/benchmarking/benchmark_nerfstudio_paper.py index 017d1301d..af780df2a 100644 --- a/scripts/benchmarking/benchmark_nerfstudio_paper.py +++ b/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)