mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
288e1496f1
* ns-process-data updated to output a .ply of sparse points from colmap * docs for loading 3D points * format fixing * ruff formatting * added read_points3D_text * formatting * added 3D points to test * prompted user bool fix * small fixes to pass test cases * fix colors * fixed coord system * Tweak frames * Comment wording * ruff * Comment intentional default change for load_3D_points * Wording * More wording * Fix double/float issue --------- Co-authored-by: Ethan Weber <ethanweber@berkeley.edu> Co-authored-by: Brent Yi <yibrenth@gmail.com>
124 lines
4.0 KiB
Python
124 lines
4.0 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)
|