From d9a391f79a7301a091148003b57d1ab86b5a808a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=C3=A1=C5=A1=20Kulh=C3=A1nek?= Date: Sat, 11 Mar 2023 22:14:45 +0100 Subject: [PATCH] Fixed: MinimalDataParser add options, allow Nerfstudio data format to specify test indices (#1586) * MinimalDataParser add options, allow Nerfstudio data format to specify test indices * Switch from indices to filenames and add tests * Fix #1581 * Fix changed interface --- nerfstudio/cameras/camera_utils.py | 2 +- .../data/dataparsers/minimal_dataparser.py | 18 +++- .../data/dataparsers/nerfstudio_dataparser.py | 53 +++++++--- nerfstudio/process_data/colmap_utils.py | 5 + .../dataparsers/test_nerfstudio_dataparser.py | 99 +++++++++++++++++++ 5 files changed, 158 insertions(+), 19 deletions(-) create mode 100644 tests/dataparsers/test_nerfstudio_dataparser.py diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index d2da627b7..79490ae67 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/nerfstudio/cameras/camera_utils.py @@ -471,7 +471,7 @@ def auto_orient_and_center_poses( poses: TensorType["num_poses":..., 4, 4], method: Literal["pca", "up", "vertical", "none"] = "up", center_method: Literal["poses", "focus", "none"] = "poses", -) -> Tuple[TensorType["num_poses":..., 3, 4], TensorType[4, 4]]: +) -> Tuple[TensorType["num_poses":..., 3, 4], TensorType[3, 4]]: """Orients and centers the poses. We provide two methods for orientation: pca and up. pca: Orient the poses so that the principal directions of the camera centers are aligned diff --git a/nerfstudio/data/dataparsers/minimal_dataparser.py b/nerfstudio/data/dataparsers/minimal_dataparser.py index 2fa2218a4..6d8de4ec5 100644 --- a/nerfstudio/data/dataparsers/minimal_dataparser.py +++ b/nerfstudio/data/dataparsers/minimal_dataparser.py @@ -52,15 +52,17 @@ class MinimalDataParser(DataParser): filepath = self.config.data / f"{split}.npz" data = np.load(filepath, allow_pickle=True) - image_filenames = data["image_filenames"].tolist() - mask_filenames = data["mask_filenames"].tolist() if "mask_filenames" in data.keys() else None + image_filenames = [filepath.parent / path for path in data["image_filenames"].tolist()] + mask_filenames = None + if "mask_filenames" in data.keys(): + mask_filenames = [filepath.parent / path for path in data["mask_filenames"].tolist()] metadata = None if "semantics" in data.keys(): semantics = data["semantics"].item() metadata = { "semantics": Semantics( - filenames=semantics["filenames"].tolist(), + filenames=[filepath.parent / path for path in semantics["filenames"].tolist()], classes=semantics["classes"].tolist(), colors=torch.from_numpy(semantics["colors"]), mask_classes=semantics["mask_classes"].tolist(), @@ -87,11 +89,21 @@ class MinimalDataParser(DataParser): times=torch.from_numpy(camera_np["times"]) if "times" in camera_np.keys() else None, ) + applied_scale = 1.0 + applied_transform = np.eye(4, dtype=np.float32)[:3, :] + if "applied_scale" in data.keys(): + applied_scale = float(data["applied_scale"]) + if "applied_transform" in data.keys(): + applied_transform = data["applied_transform"].astype(np.float32) + assert applied_transform.shape == (3, 4) + dataparser_outputs = DataparserOutputs( image_filenames=image_filenames, cameras=cameras, scene_box=scene_box, mask_filenames=mask_filenames, + dataparser_transform=applied_transform, + dataparser_scale=applied_scale, metadata=metadata, ) return dataparser_outputs diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index ef25653fd..a712d5005 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py @@ -182,22 +182,36 @@ class Nerfstudio(DataParser): You should check that depth_file_path is specified for every frame (or zero frames) in transforms.json. """ - # 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_fraction) - num_eval_images = num_images - num_train_images - i_all = np.arange(num_images) - 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_eval = np.setdiff1d(i_all, i_train) # eval images are the remaining images - assert len(i_eval) == num_eval_images - if split == "train": - indices = i_train - elif split in ["val", "test"]: - indices = i_eval + has_split_files_spec = any(f"{split}_filenames" in meta for split in ("train", "val", "test")) + if f"{split}_filenames" in meta: + # Validate split first + split_filenames = set(self._get_fname(PurePath(x), data_dir) for x in meta[f"{split}_filenames"]) + unmatched_filenames = split_filenames.difference(image_filenames) + if unmatched_filenames: + raise RuntimeError(f"Some filenames for split {split} were not found: {unmatched_filenames}.") + + indices = [i for i, path in enumerate(image_filenames) if path in split_filenames] + CONSOLE.log(f"[yellow] Dataset is overriding {split}_indices to {indices}") + indices = np.array(indices, dtype=np.int32) + elif has_split_files_spec: + raise RuntimeError(f"The dataset's list of filenames for split {split} is missing.") else: - raise ValueError(f"Unknown dataparser split {split}") + # 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_fraction) + num_eval_images = num_images - num_train_images + i_all = np.arange(num_images) + 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_eval = np.setdiff1d(i_all, i_train) # eval images are the remaining images + assert len(i_eval) == num_eval_images + 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"] @@ -274,6 +288,15 @@ class Nerfstudio(DataParser): assert self.downscale_factor is not None cameras.rescale_output_resolution(scaling_factor=1.0 / self.downscale_factor) + if "applied_transform" in meta: + applied_transform = np.array(meta["applied_transform"], dtype=np.float32) + transform_matrix = transform_matrix @ np.concatenate( + (applied_transform, np.array([0, 0, 0, 1], dtype=np.float32)), 0 + ) + if "applied_scale" in meta: + applied_scale = float(meta["applied_scale"]) + scale_factor *= applied_scale + dataparser_outputs = DataparserOutputs( image_filenames=image_filenames, cameras=cameras, diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index f8254f0a8..24517a331 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/nerfstudio/process_data/colmap_utils.py @@ -455,6 +455,11 @@ def colmap_to_json( out = parse_colmap_camera_params(cam_id_to_camera[1]) out["frames"] = frames + applied_transform = np.eye(4)[:3, :] + applied_transform = applied_transform[np.array([1, 0, 2]), :] + applied_transform[2, :] *= -1 + out["applied_transform"] = applied_transform + with open(output_dir / "transforms.json", "w", encoding="utf-8") as f: json.dump(out, f, indent=4) diff --git a/tests/dataparsers/test_nerfstudio_dataparser.py b/tests/dataparsers/test_nerfstudio_dataparser.py new file mode 100644 index 000000000..c9dce5284 --- /dev/null +++ b/tests/dataparsers/test_nerfstudio_dataparser.py @@ -0,0 +1,99 @@ +""" +Nerfstudio dataparser +""" +# pylint: disable=all +import json +from pathlib import Path + +import numpy as np +import pytest +from PIL import Image +from pytest import fixture + + +@fixture +def mocked_dataset(tmp_path: Path): + """Mocked dataset with transforms""" + (tmp_path / "images_4").mkdir() + frames = [] + for i in range(10): + Image.new("RGB", (100, 150)).save(tmp_path / "images_4" / f"img_{i}.png") + frames.append( + { + "file_path": f"img_{i}.png", + "transform_matrix": np.eye(4).tolist(), + } + ) + with (tmp_path / "transforms.json").open("w+", encoding="utf8") as f: + json.dump({"fl_x": 2, "fl_y": 3, "cx": 4, "cy": 5, "h": 150, "w": 100, "frames": frames}, f) + return tmp_path + + +@pytest.mark.parametrize("orientation_method", ["up", "none", "pca"]) +def test_nerfstudio_dataparser_no_filelist(mocked_dataset, orientation_method): + """Tests basic load""" + assert (mocked_dataset / "images_4").exists() + from nerfstudio.data.dataparsers.nerfstudio_dataparser import ( + DataparserOutputs, + Nerfstudio, + NerfstudioDataParserConfig, + ) + + parser: Nerfstudio = NerfstudioDataParserConfig( + data=mocked_dataset, + downscale_factor=4, + orientation_method=orientation_method, + center_method="none", + auto_scale_poses=False, + ).setup() + + paths = set() + for split in ("train", "test", "val"): + out = parser.get_dataparser_outputs(split) + assert isinstance(out, DataparserOutputs) + assert len(out.image_filenames) > 0 + paths.update(out.image_filenames) + train_files = set(parser.get_dataparser_outputs("train").image_filenames) + assert len(train_files.intersection(parser.get_dataparser_outputs("val").image_filenames)) == 0 + assert len(train_files.intersection(parser.get_dataparser_outputs("test").image_filenames)) == 0 + assert len(paths) == 10 + + +def test_nerfstudio_dataparser_split_filelist(mocked_dataset): + """Tests basic load""" + assert (mocked_dataset / "images_4").exists() + with open(mocked_dataset / "transforms.json", "r+") as f: + data = json.load(f) + data["train_filenames"] = ["img_0.png", "img_1.png"] + data["val_filenames"] = ["img_2.png", "img_3.png"] + data["test_filenames"] = ["img_4.png", "img_5.png"] + f.seek(0) + f.truncate(0) + json.dump(data, f) + + from nerfstudio.data.dataparsers.nerfstudio_dataparser import ( + DataparserOutputs, + Nerfstudio, + NerfstudioDataParserConfig, + ) + + parser: Nerfstudio = NerfstudioDataParserConfig( + data=mocked_dataset, + downscale_factor=4, + orientation_method="none", + center_method="none", + auto_scale_poses=False, + ).setup() + + assert parser.get_dataparser_outputs("train").image_filenames == [ + mocked_dataset / "images_4/img_0.png", + mocked_dataset / "images_4/img_1.png", + ] + assert parser.get_dataparser_outputs("val").image_filenames == [ + mocked_dataset / "images_4/img_2.png", + mocked_dataset / "images_4/img_3.png", + ] + assert parser.get_dataparser_outputs("test").image_filenames == [ + mocked_dataset / "images_4/img_4.png", + mocked_dataset / "images_4/img_5.png", + ]