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)