MinimalDataParser add options, allow Nerfstudio data format to specify test indices (#1541)

* MinimalDataParser add options, allow Nerfstudio data format to specify test indices

* Switch from indices to filenames and add tests

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[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 component of the points is aligned with the axes.

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

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,

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

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_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",
+    ]