Merge branch 'main' of https://github.com/nerfstudio-project/nerfstudio into terrance/dreamembedding

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

tests/process_data/test_process_images.py

@@ -0,0 +1,80 @@
+"""
+Process images test
+"""
+import os
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
+from nerfstudio.data.utils.colmap_parsing_utils import Camera
+from nerfstudio.data.utils.colmap_parsing_utils import Image as ColmapImage
+from nerfstudio.data.utils.colmap_parsing_utils import (
+    write_cameras_binary,
+    write_images_binary,
+)
+from scripts.process_data import ProcessImages
+
+
+def test_process_images_skip_colmap(tmp_path: Path):
+    """
+    Test ns-process-data images
+    """
+    # Mock a colmap sparse model
+    width = 100
+    height = 150
+    sparse_path = tmp_path / "sparse" / "0"
+    sparse_path.mkdir(exist_ok=True, parents=True)
+    (tmp_path / "images").mkdir(exist_ok=True, parents=True)
+    write_cameras_binary(
+        {1: Camera(1, "OPENCV", width, height, [110, 110, 50, 75, 0, 0, 0, 0, 0, 0])},
+        sparse_path / "cameras.bin",
+    )
+    frames = {}
+    num_frames = 10
+    qvecs = np.random.uniform(size=(num_frames, 4))
+    tvecs = np.random.uniform(size=(num_frames, 3))
+    # original_poses = np.concatenate(
+    #     (
+    #         np.concatenate(
+    #             (
+    #                 np.stack(list(map(qvec2rotmat, qvecs))),
+    #                 tvecs[:, :, None],
+    #             ),
+    #             -1,
+    #         ),
+    #         np.array([[[0, 0, 0, 1]]], dtype=qvecs.dtype).repeat(num_frames, 0),
+    #     ),
+    #     -2,
+    # )
+    for i in range(num_frames):
+        frames[i + 1] = ColmapImage(i + 1, qvecs[i], tvecs[i], 1, f"image_{i}.png", [], [])
+        Image.new("RGB", (width, height)).save(tmp_path / "images" / f"image_{i}.png")
+    write_images_binary(frames, sparse_path / "images.bin")
+
+    # Mock missing COLMAP and ffmpeg in the dev env
+    old_path = os.environ.get("PATH", "")
+    os.environ["PATH"] = str(tmp_path / "mocked_bin") + f":{old_path}"
+    (tmp_path / "mocked_bin").mkdir()
+    (tmp_path / "mocked_bin" / "colmap").touch(mode=0o777)
+    (tmp_path / "mocked_bin" / "ffmpeg").touch(mode=0o777)
+
+    # Run ProcessImages
+    cmd = ProcessImages(tmp_path / "images", tmp_path / "nerfstudio", colmap_model_path=sparse_path, skip_colmap=True)
+    cmd.main()
+    os.environ["PATH"] = old_path
+
+    assert (tmp_path / "nerfstudio" / "transforms.json").exists()
+    parser = NerfstudioDataParserConfig(
+        data=tmp_path / "nerfstudio",
+        downscale_factor=None,
+        orientation_method="none",  # orientation_method,
+        center_method="none",
+        auto_scale_poses=False,
+    ).setup()
+    outputs = parser.get_dataparser_outputs("train")
+    assert len(outputs.image_filenames) == 9
+    assert torch.is_tensor(outputs.dataparser_transform)
+    # TODO @jkulhanek: Add tests if the loaded poses are the same

tests/test_train.py

@@ -27,6 +27,7 @@ BLACKLIST = [
     "nerfplayer-ngp",
     "nerfplayer-nerfacto",
     "dreamfusion",
+    "neus",
 ]