mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
1b9dfe8bb3
* Support COLMAP reconstructions with nested image dirs The code was previously looking for the file.name, which drops all directory information from the path. Instead, we want to get the file path relative to the data dir, which will only drop parts of the path above the data dir, allowing us to recurse into subdirectories within the images dir * Adhere to style guidelines * Add testcase to cover image dirs with subdirs
211 lines
7.2 KiB
Python
211 lines
7.2 KiB
Python
"""
|
|
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,
|
|
Image as ColmapImage,
|
|
Point3D,
|
|
qvec2rotmat,
|
|
write_cameras_binary,
|
|
write_images_binary,
|
|
write_points3D_binary,
|
|
)
|
|
from nerfstudio.process_data.images_to_nerfstudio_dataset import ImagesToNerfstudioDataset
|
|
|
|
|
|
def random_quaternion(num_poses: int):
|
|
"""
|
|
Generates random rotation quaternion.
|
|
"""
|
|
u, v, w = np.random.uniform(size=(3, num_poses))
|
|
quaternion = np.stack(
|
|
(
|
|
np.sqrt(1 - u) * np.sin(2 * np.pi * v),
|
|
np.sqrt(1 - u) * np.cos(2 * np.pi * v),
|
|
np.sqrt(u) * np.sin(2 * np.pi * w),
|
|
np.sqrt(u) * np.cos(2 * np.pi * w),
|
|
),
|
|
-1,
|
|
)
|
|
return quaternion
|
|
|
|
|
|
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",
|
|
)
|
|
write_points3D_binary(
|
|
{
|
|
1: Point3D(
|
|
id=1,
|
|
xyz=np.array([0, 0, 0]),
|
|
rgb=np.array([0, 0, 0]),
|
|
error=np.array([0]),
|
|
image_ids=np.array([1]),
|
|
point2D_idxs=np.array([0]),
|
|
),
|
|
},
|
|
sparse_path / "points3D.bin",
|
|
)
|
|
frames = {}
|
|
num_frames = 10
|
|
qvecs = random_quaternion(num_frames)
|
|
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)
|
|
|
|
# Convert images into a NerfStudio dataset
|
|
cmd = ImagesToNerfstudioDataset(
|
|
data=tmp_path / "images", output_dir=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)
|
|
|
|
# Test if the original poses can be obtained back
|
|
gt_poses = original_poses[[0, 1, 2, 3, 4, 5, 6, 7, 9]]
|
|
dataparser_poses = outputs.transform_poses_to_original_space(outputs.cameras.camera_to_worlds, "opencv").numpy()
|
|
dataparser_poses = np.concatenate(
|
|
(dataparser_poses, np.array([[[0, 0, 0, 1]]]).repeat(len(dataparser_poses), 0)), 1
|
|
)
|
|
dataparser_poses = np.linalg.inv(dataparser_poses)
|
|
np.testing.assert_allclose(gt_poses, dataparser_poses, rtol=0, atol=1e-5)
|
|
|
|
|
|
def test_process_images_recursively_skip_colmap(tmp_path: Path):
|
|
"""
|
|
Test ns-process-data images when images contains subdirectories"
|
|
"""
|
|
# 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",
|
|
)
|
|
write_points3D_binary(
|
|
{
|
|
1: Point3D(
|
|
id=1,
|
|
xyz=np.array([0, 0, 0]),
|
|
rgb=np.array([0, 0, 0]),
|
|
error=np.array([0]),
|
|
image_ids=np.array([1]),
|
|
point2D_idxs=np.array([0]),
|
|
),
|
|
},
|
|
sparse_path / "points3D.bin",
|
|
)
|
|
frames = {}
|
|
num_frames = 9
|
|
num_subdirs = 3
|
|
qvecs = random_quaternion(num_frames)
|
|
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):
|
|
subdir = f"subdir_{num_frames // num_subdirs}"
|
|
frames[i + 1] = ColmapImage(i + 1, qvecs[i], tvecs[i], 1, f"{subdir}/image_{i}.png", [], [])
|
|
(tmp_path / "images" / subdir).mkdir(parents=True, exist_ok=True)
|
|
Image.new("RGB", (width, height)).save(tmp_path / "images" / subdir / 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)
|
|
|
|
# Convert images into a NerfStudio dataset
|
|
cmd = ImagesToNerfstudioDataset(
|
|
data=tmp_path / "images", output_dir=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)
|
|
|
|
# Test if the original poses can be obtained back
|
|
dataparser_poses = outputs.transform_poses_to_original_space(outputs.cameras.camera_to_worlds, "opencv").numpy()
|
|
dataparser_poses = np.concatenate(
|
|
(dataparser_poses, np.array([[[0, 0, 0, 1]]]).repeat(len(dataparser_poses), 0)), 1
|
|
)
|
|
dataparser_poses = np.linalg.inv(dataparser_poses)
|
|
np.testing.assert_allclose(original_poses, dataparser_poses, rtol=0, atol=1e-5)
|