diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index 9db20647f..f9613dbc1 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/nerfstudio/process_data/colmap_utils.py @@ -587,6 +587,7 @@ def colmap_to_json( output_dir: Path, camera_model: CameraModel, camera_mask_path: Optional[Path] = None, + image_rename_map: Optional[Dict[str, str]] = None, ) -> int: """Converts COLMAP's cameras.bin and images.bin to a JSON file. @@ -619,7 +620,10 @@ def colmap_to_json( c2w = c2w[np.array([1, 0, 2, 3]), :] c2w[2, :] *= -1 - name = Path(f"./images/{im_data.name}") + name = im_data.name + if image_rename_map is not None: + name = image_rename_map[name] + name = Path(f"./images/{name}") frame = { "file_path": name.as_posix(), diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index edd205598..b1929deaa 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -24,7 +24,7 @@ from typing import List, Optional, Tuple import cv2 import numpy as np from rich.console import Console -from typing_extensions import Literal +from typing_extensions import Literal, OrderedDict from nerfstudio.utils.rich_utils import status from nerfstudio.utils.scripts import run_command @@ -242,7 +242,7 @@ def copy_and_upscale_polycam_depth_maps_list( return copied_depth_map_paths -def copy_images(data: Path, image_dir: Path, verbose) -> int: +def copy_images(data: Path, image_dir: Path, verbose) -> OrderedDict[Path, Path]: """Copy images from a directory to a new directory. Args: @@ -250,7 +250,7 @@ def copy_images(data: Path, image_dir: Path, verbose) -> int: image_dir: Path to the output directory. verbose: If True, print extra logging. Returns: - The number of images copied. + The mapping from the original filenames to the new ones. """ with status(msg="[bold yellow]Copying images...", spinner="bouncingBall", verbose=verbose): image_paths = list_images(data) @@ -259,9 +259,8 @@ def copy_images(data: Path, image_dir: Path, verbose) -> int: CONSOLE.log("[bold red]:skull: No usable images in the data folder.") sys.exit(1) - num_frames = len(copy_images_list(image_paths=image_paths, image_dir=image_dir, verbose=verbose)) - - return num_frames + copied_images = copy_images_list(image_paths=image_paths, image_dir=image_dir, verbose=verbose) + return OrderedDict((original_path, new_path) for original_path, new_path in zip(image_paths, copied_images)) def downscale_images(image_dir: Path, num_downscales: int, folder_name: str = "images", verbose: bool = False) -> str: diff --git a/scripts/process_data.py b/scripts/process_data.py index 7b6bb7922..b95146d99 100644 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -8,7 +8,7 @@ import sys import zipfile from dataclasses import dataclass from pathlib import Path -from typing import Tuple, Union +from typing import Dict, Optional, Tuple, Union import numpy as np import tyro @@ -96,6 +96,7 @@ class ProcessImages: def main(self) -> None: """Process images into a nerfstudio dataset.""" + # pylint: disable=too-many-statements require_cameras_exist = False colmap_model_path = self.output_dir / Path(self.colmap_model_path) if self.colmap_model_path != DEFAULT_COLMAP_PATH: @@ -112,6 +113,7 @@ class ProcessImages: install_checks.check_ffmpeg_installed() install_checks.check_colmap_installed() + image_rename_map: Optional[Dict[str, str]] = None self.output_dir.mkdir(parents=True, exist_ok=True) image_dir = self.output_dir / "images" image_dir.mkdir(parents=True, exist_ok=True) @@ -129,7 +131,11 @@ class ProcessImages: # Copy and downscale images if not self.skip_image_processing: # Copy images to output directory - num_frames = process_data_utils.copy_images(self.data, image_dir=image_dir, verbose=self.verbose) + image_rename_map_paths = process_data_utils.copy_images( + self.data, image_dir=image_dir, verbose=self.verbose + ) + image_rename_map = dict((a.name, 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") # Downscale images @@ -152,6 +158,9 @@ class ProcessImages: self._run_colmap(image_dir, colmap_dir) + # Colmap uses renamed images + image_rename_map = None + # Save transforms.json if (colmap_model_path / "cameras.bin").exists(): with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): @@ -160,6 +169,7 @@ class ProcessImages: images_path=colmap_model_path / "images.bin", output_dir=self.output_dir, camera_model=CAMERA_MODELS[self.camera_type], + image_rename_map=image_rename_map, ) summary_log.append(f"Colmap matched {num_matched_frames} images") summary_log.append(colmap_utils.get_matching_summary(num_frames, num_matched_frames)) @@ -371,6 +381,7 @@ class ProcessVideo: output_dir=self.output_dir, camera_model=CAMERA_MODELS[self.camera_type], camera_mask_path=mask_path, + image_rename_map=None, ) summary_log.append(f"Colmap matched {num_matched_frames} images") summary_log.append(colmap_utils.get_matching_summary(num_extracted_frames, num_matched_frames)) @@ -490,6 +501,7 @@ class ProcessInsta360: images_path=colmap_dir / "sparse" / "0" / "images.bin", output_dir=self.output_dir, camera_model=CAMERA_MODELS["fisheye"], + image_rename_map=None, ) summary_log.append(f"Colmap matched {num_matched_frames} images") summary_log.append(colmap_utils.get_matching_summary(num_extracted_frames, num_matched_frames))