mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
d542646c72
* `isort` -> `ruff` * `pylint` -> `ruff` * Remove line * More minor cleanup * Better documentation for rules we are enabling * Fix gitignore, run black * vscode settings --------- Co-authored-by: Matt Tancik <mattt4001@gmail.com>
99 lines
3.2 KiB
Python
99 lines
3.2 KiB
Python
"""
|
|
Nerfstudio dataparser
|
|
"""
|
|
|
|
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 (
|
|
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",
|
|
]
|