diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index 7b3712cd0..dae490626 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/nerfstudio/cameras/camera_utils.py @@ -431,7 +431,7 @@ def rotation_matrix(a: TensorType[3], b: TensorType[3]) -> TensorType[3, 3]: def auto_orient_and_center_poses( poses: TensorType["num_poses":..., 4, 4], method: Literal["pca", "up", "none"] = "up", center_poses: bool = True -) -> Tuple[TensorType["num_poses":..., 3, 4], TensorType[3, 4]]: +) -> Tuple[TensorType["num_poses":..., 3, 4], TensorType[4, 4]]: """Orients and centers the poses. We provide two methods for orientation: pca and up. pca: Orient the poses so that the principal component of the points is aligned with the axes. diff --git a/nerfstudio/data/dataparsers/minimal_dataparser.py b/nerfstudio/data/dataparsers/minimal_dataparser.py index 6d8de4ec5..2fa2218a4 100644 --- a/nerfstudio/data/dataparsers/minimal_dataparser.py +++ b/nerfstudio/data/dataparsers/minimal_dataparser.py @@ -52,17 +52,15 @@ class MinimalDataParser(DataParser): filepath = self.config.data / f"{split}.npz" data = np.load(filepath, allow_pickle=True) - 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()] + image_filenames = data["image_filenames"].tolist() + mask_filenames = data["mask_filenames"].tolist() if "mask_filenames" in data.keys() else None metadata = None if "semantics" in data.keys(): semantics = data["semantics"].item() metadata = { "semantics": Semantics( - filenames=[filepath.parent / path for path in semantics["filenames"].tolist()], + filenames=semantics["filenames"].tolist(), classes=semantics["classes"].tolist(), colors=torch.from_numpy(semantics["colors"]), mask_classes=semantics["mask_classes"].tolist(), @@ -89,21 +87,11 @@ 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 2d51b8be5..fd0d49830 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py @@ -182,36 +182,22 @@ class Nerfstudio(DataParser): You should check that depth_file_path is specified for every frame (or zero frames) in transforms.json. """ - 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.") + # 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: - # 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}") + raise ValueError(f"Unknown dataparser split {split}") if "orientation_override" in meta: orientation_method = meta["orientation_override"] @@ -288,15 +274,6 @@ 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 f34d6cd84..f8254f0a8 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/nerfstudio/process_data/colmap_utils.py @@ -455,11 +455,6 @@ 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, 3]), :] - 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 deleted file mode 100644 index c5dfd9518..000000000 --- a/tests/dataparsers/test_nerfstudio_dataparser.py +++ /dev/null @@ -1,99 +0,0 @@ -""" -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_poses=False, - 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_poses=False, - 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", - ]