From 1b9dfe8bb378750457553882d25658ea845c72b1 Mon Sep 17 00:00:00 2001 From: Zac Stewart Date: Wed, 3 Apr 2024 02:17:17 +0800 Subject: [PATCH] Support COLMAP reconstructions with non-flat image dirs (#3006) * 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 --- .../images_to_nerfstudio_dataset.py | 4 +- nerfstudio/process_data/process_data_utils.py | 2 +- tests/process_data/test_process_images.py | 87 +++++++++++++++++++ 3 files changed, 91 insertions(+), 2 deletions(-) diff --git a/nerfstudio/process_data/images_to_nerfstudio_dataset.py b/nerfstudio/process_data/images_to_nerfstudio_dataset.py index 7ae7d5bd4..fa2d8b2bc 100644 --- a/nerfstudio/process_data/images_to_nerfstudio_dataset.py +++ b/nerfstudio/process_data/images_to_nerfstudio_dataset.py @@ -87,7 +87,9 @@ class ImagesToNerfstudioDataset(ColmapConverterToNerfstudioDataset): ) image_rename_map_paths.update(eval_image_rename_map_paths) - image_rename_map = dict((a.name, b.name) for a, b in image_rename_map_paths.items()) + image_rename_map = dict( + (a.relative_to(self.data).as_posix(), b.name) for a, b in image_rename_map_paths.items() + ) num_frames = len(image_rename_map) summary_log.append(f"Starting with {num_frames} images") diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index d754525d3..b5b2391a0 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -62,7 +62,7 @@ CAMERA_MODELS = { } -def list_images(data: Path, recursive: bool = False) -> List[Path]: +def list_images(data: Path, recursive: bool = True) -> List[Path]: """Lists all supported images in a directory Args: diff --git a/tests/process_data/test_process_images.py b/tests/process_data/test_process_images.py index 8482676c5..fd506b1fb 100644 --- a/tests/process_data/test_process_images.py +++ b/tests/process_data/test_process_images.py @@ -121,3 +121,90 @@ def test_process_images_skip_colmap(tmp_path: Path): ) 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)