From 3d2e62a480ee15757433f9d96e67af26a5a553f9 Mon Sep 17 00:00:00 2001 From: Malcolm Mielle Date: Wed, 26 Apr 2023 19:30:01 +0200 Subject: [PATCH] Refactor of image and video dataset conversion (#1690) * Refactor of image and video dataset conversion The classes to convert images and video to a NerfStudio dataset are moved to the main repo and refactor: * Code duplications are removed * Import of the class is now done from nerfstudio package instead of scripts * Doc formatting nits * Doc formatting --------- Co-authored-by: Brent Yi --- .../base_converter_to_nerfstudio_dataset.py | 46 ++ .../colmap_converter_to_nerfstudio_dataset.py | 228 ++++++++ .../images_to_nerstudio_dataset.py | 106 ++++ .../video_to_nerfstudio_dataset.py | 122 ++++ scripts/process_data.py | 522 ++---------------- tests/process_data/test_process_images.py | 10 +- 6 files changed, 544 insertions(+), 490 deletions(-) create mode 100644 nerfstudio/process_data/base_converter_to_nerfstudio_dataset.py create mode 100644 nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py create mode 100644 nerfstudio/process_data/images_to_nerstudio_dataset.py create mode 100644 nerfstudio/process_data/video_to_nerfstudio_dataset.py diff --git a/nerfstudio/process_data/base_converter_to_nerfstudio_dataset.py b/nerfstudio/process_data/base_converter_to_nerfstudio_dataset.py new file mode 100644 index 000000000..a1578cc05 --- /dev/null +++ b/nerfstudio/process_data/base_converter_to_nerfstudio_dataset.py @@ -0,0 +1,46 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Base class to process images or video into a nerfstudio dataset +""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from pathlib import Path + + +@dataclass +class BaseConverterToNerfstudioDataset(ABC): + """Base class to process images or video into a nerfstudio dataset""" + + data: Path + """Path the data, either a video file or a directory of images.""" + output_dir: Path + """Path to the output directory.""" + verbose: bool = False + """If True, print extra logging.""" + + def __post_init__(self) -> None: + self.output_dir.mkdir(parents=True, exist_ok=True) + self.image_dir.mkdir(parents=True, exist_ok=True) + + @property + def image_dir(self) -> Path: # pylint: disable=missing-function-docstring + return self.output_dir / "images" + + @abstractmethod + def main(self) -> None: + """This method implements the conversion logic for each type of data""" + raise NotImplementedError("the main method for conversion needs to be implemented") diff --git a/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py b/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py new file mode 100644 index 000000000..4effdc377 --- /dev/null +++ b/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py @@ -0,0 +1,228 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Base class to processes a video or image sequence to a nerfstudio compatible dataset.""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +from rich.console import Console +from typing_extensions import Literal + +from nerfstudio.process_data import colmap_utils, hloc_utils, process_data_utils +from nerfstudio.process_data.base_converter_to_nerfstudio_dataset import ( + BaseConverterToNerfstudioDataset, +) +from nerfstudio.process_data.process_data_utils import CAMERA_MODELS +from nerfstudio.utils import install_checks + +CONSOLE = Console(width=120) + + +@dataclass +class ColmapConverterToNerfstudioDataset(BaseConverterToNerfstudioDataset): + """Base class to process images or video into a nerfstudio dataset using colmap""" + + camera_type: Literal["perspective", "fisheye", "equirectangular"] = "perspective" + """Camera model to use.""" + matching_method: Literal["exhaustive", "sequential", "vocab_tree"] = "vocab_tree" + """Feature matching method to use. Vocab tree is recommended for a balance of speed + and accuracy. Exhaustive is slower but more accurate. Sequential is faster but + should only be used for videos.""" + sfm_tool: Literal["any", "colmap", "hloc"] = "any" + """Structure from motion tool to use. Colmap will use sift features, hloc can use + many modern methods such as superpoint features and superglue matcher""" + refine_pixsfm: bool = False + """If True, runs refinement using Pixel Perfect SFM. + Only works with hloc sfm_tool""" + feature_type: Literal[ + "any", + "sift", + "superpoint", + "superpoint_aachen", + "superpoint_max", + "superpoint_inloc", + "r2d2", + "d2net-ss", + "sosnet", + "disk", + ] = "any" + """Type of feature to use.""" + matcher_type: Literal[ + "any", + "NN", + "superglue", + "superglue-fast", + "NN-superpoint", + "NN-ratio", + "NN-mutual", + "adalam", + ] = "any" + """Matching algorithm.""" + num_downscales: int = 3 + """Number of times to downscale the images. Downscales by 2 each time. For example a value of 3 will downscale the + images by 2x, 4x, and 8x.""" + skip_colmap: bool = False + """If True, skips COLMAP and generates transforms.json if possible.""" + skip_image_processing: bool = False + """If True, skips copying and downscaling of images and only runs COLMAP if possible and enabled""" + colmap_model_path: Path = Path("colmap/sparse/0") + """Optionally sets the path of the colmap model. Used only when --skip-colmap is set to True. The path is relative + to the output directory. + """ + colmap_cmd: str = "colmap" + """How to call the COLMAP executable.""" + images_per_equirect: Literal[8, 14] = 8 + """Number of samples per image to take from each equirectangular image. + Used only when camera-type is equirectangular. + """ + crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0) + """Portion of the image to crop. All values should be in [0,1]. (top, bottom, left, right)""" + crop_bottom: float = 0.0 + """Portion of the image to crop from the bottom. + Can be used instead of `crop-factor 0.0 [num] 0.0 0.0` Should be in [0,1]. + """ + gpu: bool = True + """If True, use GPU.""" + use_sfm_depth: bool = False + """If True, export and use depth maps induced from SfM points.""" + include_depth_debug: bool = False + """If --use-sfm-depth and this flag is True, also export debug images showing Sf overlaid upon input images.""" + + @staticmethod + def default_colmap_path() -> Path: # pylint: disable=missing-function-docstring + return Path("colmap/sparse/0") + + @property + def absolute_colmap_model_path(self) -> Path: # pylint: disable=missing-function-docstring + return self.output_dir / self.colmap_model_path + + @property + def absolute_colmap_path(self) -> Path: # pylint: disable=missing-function-docstring + return self.output_dir / "colmap" + + def _save_transforms( + self, + num_frames: int, + image_id_to_depth_path: Optional[Dict[int, Path]] = None, + camera_mask_path: Optional[Path] = None, + image_rename_map: Optional[Dict[str, str]] = None, + ) -> List[str]: + """Save colmap transforms into the output folder + + Args: + image_id_to_depth_path: When including sfm-based depth, embed these depth file paths in the exported json + image_rename_map: Use these image names instead of the names embedded in the COLMAP db + """ + summary_log = [] + if (self.absolute_colmap_model_path / "cameras.bin").exists(): + with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): + num_matched_frames = colmap_utils.colmap_to_json( + recon_dir=self.absolute_colmap_model_path, + output_dir=self.output_dir, + image_id_to_depth_path=image_id_to_depth_path, + camera_mask_path=camera_mask_path, + 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)) + + else: + CONSOLE.log( + "[bold yellow]Warning: Could not find existing COLMAP results. " "Not generating transforms.json" + ) + return summary_log + + def _export_depth(self) -> Tuple[Optional[Dict[int, Path]], List[str]]: + """If SFM is used for creating depth image, this method will create the depth images from image in + `self.image_dir`. + + Returns: + Depth file paths indexed by COLMAP image id, logs + """ + summary_log = [] + if self.use_sfm_depth: + depth_dir = self.output_dir / "depth" + depth_dir.mkdir(parents=True, exist_ok=True) + image_id_to_depth_path = colmap_utils.create_sfm_depth( + recon_dir=self.output_dir / self.default_colmap_path(), + output_dir=depth_dir, + include_depth_debug=self.include_depth_debug, + input_images_dir=self.image_dir, + verbose=self.verbose, + ) + summary_log.append( + process_data_utils.downscale_images( + depth_dir, + self.num_downscales, + folder_name="depths", + nearest_neighbor=True, + verbose=self.verbose, + ) + ) + return image_id_to_depth_path, summary_log + return None, summary_log + + def _run_colmap(self, mask_path: Optional[Path] = None): + """ + Args: + mask_path: Path to the camera mask. Defaults to None. + """ + self.absolute_colmap_path.mkdir(parents=True, exist_ok=True) + + (sfm_tool, feature_type, matcher_type,) = process_data_utils.find_tool_feature_matcher_combination( + self.sfm_tool, self.feature_type, self.matcher_type + ) + # check that sfm_tool is hloc if using refine_pixsfm + if self.refine_pixsfm: + assert sfm_tool == "hloc", "refine_pixsfm only works with sfm_tool hloc" + + if sfm_tool == "colmap": + colmap_utils.run_colmap( + image_dir=self.image_dir, + colmap_dir=self.absolute_colmap_path, + camera_model=CAMERA_MODELS[self.camera_type], + camera_mask_path=mask_path, + gpu=self.gpu, + verbose=self.verbose, + matching_method=self.matching_method, + colmap_cmd=self.colmap_cmd, + ) + elif sfm_tool == "hloc": + if mask_path is not None: + raise RuntimeError("Cannot use a mask with hloc. Please remove the cropping options " "and try again.") + hloc_utils.run_hloc( + image_dir=self.image_dir, + colmap_dir=self.absolute_colmap_path, + camera_model=CAMERA_MODELS[self.camera_type], + verbose=self.verbose, + matching_method=self.matching_method, + feature_type=feature_type, + matcher_type=matcher_type, + refine_pixsfm=self.refine_pixsfm, + ) + else: + raise RuntimeError("Invalid combination of sfm_tool, feature_type, and matcher_type, " "exiting") + + def __post_init__(self) -> None: + super().__post_init__() + install_checks.check_ffmpeg_installed() + install_checks.check_colmap_installed() + + if self.crop_bottom < 0.0 or self.crop_bottom > 1: + raise RuntimeError("crop_bottom must be set between 0 and 1.") + + if self.crop_bottom > 0.0: + self.crop_factor = (0.0, self.crop_bottom, 0.0, 0.0) diff --git a/nerfstudio/process_data/images_to_nerstudio_dataset.py b/nerfstudio/process_data/images_to_nerstudio_dataset.py new file mode 100644 index 000000000..477c94c20 --- /dev/null +++ b/nerfstudio/process_data/images_to_nerstudio_dataset.py @@ -0,0 +1,106 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Processes an image sequence to a nerfstudio compatible dataset.""" + +from dataclasses import dataclass +from typing import Optional + +from rich.console import Console + +from nerfstudio.process_data import equirect_utils, process_data_utils +from nerfstudio.process_data.colmap_converter_to_nerfstudio_dataset import ( + ColmapConverterToNerfstudioDataset, +) + +CONSOLE = Console(width=120) + + +@dataclass +class ImagesToNerfstudioDataset(ColmapConverterToNerfstudioDataset): + """Process images into a nerfstudio dataset. + + 1. Scales images to a specified size. + 2. Calculates the camera poses for each image using `COLMAP `_. + """ + + def main(self) -> None: # pylint: disable=R0915 + """Process images into a nerfstudio dataset.""" + + # pylint: disable=too-many-statements + require_cameras_exist = False + if self.colmap_model_path != ColmapConverterToNerfstudioDataset.default_colmap_path(): + if not self.skip_colmap: + raise RuntimeError("The --colmap-model-path can only be used when --skip-colmap is not set.") + if not (self.output_dir / self.colmap_model_path).exists(): + raise RuntimeError(f"The colmap-model-path {self.output_dir / self.colmap_model_path} does not exist.") + require_cameras_exist = True + + image_rename_map: Optional[dict[str, str]] = None + + # Generate planar projections if equirectangular + if self.camera_type == "equirectangular": + pers_size = equirect_utils.compute_resolution_from_equirect(self.data, self.images_per_equirect) + CONSOLE.log(f"Generating {self.images_per_equirect} {pers_size} sized images per equirectangular image") + self.data = equirect_utils.generate_planar_projections_from_equirectangular( + self.data, pers_size, self.images_per_equirect, crop_factor=self.crop_factor + ) + + summary_log = [] + + # Copy and downscale images + if not self.skip_image_processing: + # Copy images to output directory + image_rename_map_paths = process_data_utils.copy_images( + self.data, image_dir=self.image_dir, crop_factor=self.crop_factor, 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 + summary_log.append( + process_data_utils.downscale_images(self.image_dir, self.num_downscales, verbose=self.verbose) + ) + else: + num_frames = len(process_data_utils.list_images(self.data)) + if num_frames == 0: + raise RuntimeError("No usable images in the data folder.") + summary_log.append(f"Starting with {num_frames} images") + + # Run COLMAP + if not self.skip_colmap: + require_cameras_exist = True + self._run_colmap() + # Colmap uses renamed images + image_rename_map = None + + # Export depth maps + image_id_to_depth_path, log_tmp = self._export_depth() + summary_log += log_tmp + + if require_cameras_exist and not (self.absolute_colmap_model_path / "cameras.bin").exists(): + raise RuntimeError(f"Could not find existing COLMAP results ({self.colmap_model_path / 'cameras.bin'}).") + + summary_log += self._save_transforms( + num_frames, + image_id_to_depth_path, + None, + image_rename_map, + ) + + CONSOLE.log("[bold green]:tada: :tada: :tada: All DONE :tada: :tada: :tada:") + + for summary in summary_log: + CONSOLE.log(summary) diff --git a/nerfstudio/process_data/video_to_nerfstudio_dataset.py b/nerfstudio/process_data/video_to_nerfstudio_dataset.py new file mode 100644 index 000000000..405231fcf --- /dev/null +++ b/nerfstudio/process_data/video_to_nerfstudio_dataset.py @@ -0,0 +1,122 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Processes a video to a nerfstudio compatible dataset.""" + +import shutil +from dataclasses import dataclass + +from rich.console import Console + +from nerfstudio.process_data import equirect_utils, process_data_utils +from nerfstudio.process_data.colmap_converter_to_nerfstudio_dataset import ( + ColmapConverterToNerfstudioDataset, +) + +CONSOLE = Console(width=120) + + +@dataclass +class VideoToNerfstudioDataset(ColmapConverterToNerfstudioDataset): + """Process videos into a nerfstudio dataset. + + This script does the following: + + 1. Converts the video into images. + 2. Scales images to a specified size. + 3. Calculates the camera poses for each image using `COLMAP `_. + """ + + num_frames_target: int = 300 + """Target number of frames to use for the dataset, results may not be exact.""" + percent_radius_crop: float = 1.0 + """Create circle crop mask. The radius is the percent of the image diagonal.""" + + def main(self) -> None: # pylint: disable=R0915 + """Process video into a nerfstudio dataset.""" + + summary_log = [] + # Convert video to images + if self.camera_type == "equirectangular": + # create temp images folder to store the equirect and perspective images + temp_image_dir = self.output_dir / "temp_images" + temp_image_dir.mkdir(parents=True, exist_ok=True) + summary_log, num_extracted_frames = process_data_utils.convert_video_to_images( + self.data, + image_dir=temp_image_dir, + num_frames_target=self.num_frames_target, + crop_factor=(0.0, 0.0, 0.0, 0.0), + verbose=self.verbose, + ) + else: + summary_log, num_extracted_frames = process_data_utils.convert_video_to_images( + self.data, + image_dir=self.image_dir, + num_frames_target=self.num_frames_target, + crop_factor=self.crop_factor, + verbose=self.verbose, + ) + + # Generate planar projections if equirectangular + if self.camera_type == "equirectangular": + perspective_image_size = equirect_utils.compute_resolution_from_equirect( + self.output_dir / "temp_images", self.images_per_equirect + ) + + equirect_utils.generate_planar_projections_from_equirectangular( + self.output_dir / "temp_images", + perspective_image_size, + self.images_per_equirect, + crop_factor=self.crop_factor, + ) + + # copy the perspective images to the image directory + process_data_utils.copy_images( + self.output_dir / "temp_images" / "planar_projections", + image_dir=self.output_dir / "images", + verbose=False, + ) + + # remove the temp_images folder + shutil.rmtree(self.output_dir / "temp_images", ignore_errors=True) + + # # Create mask + mask_path = process_data_utils.save_mask( + image_dir=self.image_dir, + num_downscales=self.num_downscales, + crop_factor=(0.0, 0.0, 0.0, 0.0), + percent_radius=self.percent_radius_crop, + ) + if mask_path is not None: + summary_log.append(f"Saved mask to {mask_path}") + + # # Downscale images + summary_log.append( + process_data_utils.downscale_images(self.image_dir, self.num_downscales, verbose=self.verbose) + ) + + # Run Colmap + if not self.skip_colmap: + self._run_colmap(mask_path) + + # Export depth maps + image_id_to_depth_path, log_tmp = self._export_depth() + summary_log += log_tmp + + summary_log += self._save_transforms(num_extracted_frames, image_id_to_depth_path, mask_path) + + CONSOLE.log("[bold green]:tada: :tada: :tada: All DONE :tada: :tada: :tada:") + + for summary in summary_log: + CONSOLE.log(summary) diff --git a/scripts/process_data.py b/scripts/process_data.py index 98c896f28..f1f873ae3 100644 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -2,479 +2,37 @@ """Processes a video or image sequence to a nerfstudio compatible dataset.""" -import shutil import sys import zipfile from dataclasses import dataclass from pathlib import Path -from typing import Dict, Optional, Tuple, Union +from typing import Union import numpy as np import tyro from rich.console import Console -from typing_extensions import Annotated, Literal +from typing_extensions import Annotated from nerfstudio.process_data import ( - colmap_utils, - equirect_utils, - hloc_utils, metashape_utils, polycam_utils, process_data_utils, realitycapture_utils, record3d_utils, ) -from nerfstudio.process_data.process_data_utils import CAMERA_MODELS -from nerfstudio.utils import install_checks +from nerfstudio.process_data.colmap_converter_to_nerfstudio_dataset import ( + BaseConverterToNerfstudioDataset, +) +from nerfstudio.process_data.images_to_nerstudio_dataset import ( + ImagesToNerfstudioDataset, +) +from nerfstudio.process_data.video_to_nerfstudio_dataset import VideoToNerfstudioDataset CONSOLE = Console(width=120) -DEFAULT_COLMAP_PATH = Path("colmap/sparse/0") @dataclass -class ProcessImages: - """Process images into a nerfstudio dataset. - - This script does the following: - - 1. Scales images to a specified size. - 2. Calculates the camera poses for each image using `COLMAP `_. - """ - - data: Path - """Path the data, either a video file or a directory of images.""" - output_dir: Path - """Path to the output directory.""" - camera_type: Literal["perspective", "fisheye", "equirectangular"] = "perspective" - """Camera model to use.""" - matching_method: Literal["exhaustive", "sequential", "vocab_tree"] = "vocab_tree" - """Feature matching method to use. Vocab tree is recommended for a balance of speed and - accuracy. Exhaustive is slower but more accurate. Sequential is faster but should only be used for videos.""" - sfm_tool: Literal["any", "colmap", "hloc"] = "any" - """Structure from motion tool to use. Colmap will use sift features, hloc can use many modern methods - such as superpoint features and superglue matcher""" - refine_pixsfm: bool = False - """If True, runs refinement using Pixel Perfect SFM. Only works with hloc sfm_tool""" - feature_type: Literal[ - "any", - "sift", - "superpoint", - "superpoint_aachen", - "superpoint_max", - "superpoint_inloc", - "r2d2", - "d2net-ss", - "sosnet", - "disk", - ] = "any" - """Type of feature to use.""" - matcher_type: Literal[ - "any", "NN", "superglue", "superglue-fast", "NN-superpoint", "NN-ratio", "NN-mutual", "adalam" - ] = "any" - """Matching algorithm.""" - num_downscales: int = 3 - """Number of times to downscale the images. Downscales by 2 each time. For example a value of 3 - will downscale the images by 2x, 4x, and 8x.""" - skip_colmap: bool = False - """If True, skips COLMAP and generates transforms.json if possible.""" - skip_image_processing: bool = False - """If True, skips copying and downscaling of images and only runs COLMAP if possible and enabled""" - colmap_model_path: Path = DEFAULT_COLMAP_PATH - """Optionally sets the path of the colmap model. Used only when --skip-colmap is set to True. - The path is relative to the output directory. - """ - colmap_cmd: str = "colmap" - """How to call the COLMAP executable.""" - images_per_equirect: Literal[8, 14] = 8 - """Number of samples per image to take from each equirectangular image. - Used only when camera-type is equirectangular. - """ - crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0) - """Portion of the image to crop. All values should be in [0,1]. (top, bottom, left, right)""" - crop_bottom: float = 0.0 - """Portion of the image to crop from the bottom. - Can be used instead of `crop-factor 0.0 [num] 0.0 0.0` Should be in [0,1]. - """ - gpu: bool = True - """If True, use GPU.""" - use_sfm_depth: bool = False - """If True, export and use depth maps induced from SfM points.""" - include_depth_debug: bool = False - """If --use-sfm-depth and this flag is True, also export debug images showing SfM overlaid upon input images.""" - verbose: bool = False - """If True, print extra logging.""" - - def main(self) -> None: # pylint: disable=R0915 - """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: - if not self.skip_colmap: - CONSOLE.log("[bold red]The --colmap-model-path can only be used when --skip-colmap is not set.") - sys.exit(1) - elif not (self.output_dir / self.colmap_model_path).exists(): - CONSOLE.log( - f"[bold red]The colmap-model-path {self.output_dir / self.colmap_model_path} does not exist." - ) - sys.exit(1) - require_cameras_exist = True - - 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) - - if self.crop_bottom > 0.0: - self.crop_factor = (0.0, self.crop_bottom, 0.0, 0.0) - - # Generate planar projections if equirectangular - if self.camera_type == "equirectangular": - pers_size = equirect_utils.compute_resolution_from_equirect(self.data, self.images_per_equirect) - CONSOLE.log(f"Generating {self.images_per_equirect} {pers_size} sized images per equirectangular image") - self.data = equirect_utils.generate_planar_projections_from_equirectangular( - self.data, pers_size, self.images_per_equirect, crop_factor=self.crop_factor - ) - - summary_log = [] - - # Copy and downscale images - if not self.skip_image_processing: - # Copy images to output directory - image_rename_map_paths = process_data_utils.copy_images( - self.data, image_dir=image_dir, crop_factor=self.crop_factor, 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 - summary_log.append( - process_data_utils.downscale_images(image_dir, self.num_downscales, verbose=self.verbose) - ) - else: - num_frames = len(process_data_utils.list_images(self.data)) - if num_frames == 0: - CONSOLE.log("[bold red]:skull: No usable images in the data folder.") - sys.exit(1) - summary_log.append(f"Starting with {num_frames} images") - - # Run COLMAP - colmap_dir = self.output_dir / "colmap" - if not self.skip_colmap: - colmap_dir.mkdir(parents=True, exist_ok=True) - colmap_model_path = colmap_dir / "sparse" / "0" - require_cameras_exist = True - - self._run_colmap(image_dir, colmap_dir) - - # Colmap uses renamed images - image_rename_map = None - - # Export depth maps - if self.use_sfm_depth: - depth_dir = self.output_dir / "depth" - depth_dir.mkdir(parents=True, exist_ok=True) - image_id_to_depth_path = colmap_utils.create_sfm_depth( - recon_dir=colmap_dir / "sparse" / "0", - output_dir=depth_dir, - include_depth_debug=self.include_depth_debug, - input_images_dir=image_dir, - verbose=self.verbose, - ) - summary_log.append( - process_data_utils.downscale_images( - depth_dir, self.num_downscales, folder_name="depths", nearest_neighbor=True, verbose=self.verbose - ) - ) - else: - image_id_to_depth_path = None - - # Save transforms.json - if (colmap_model_path / "cameras.bin").exists(): - with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): - num_matched_frames = colmap_utils.colmap_to_json( - recon_dir=colmap_model_path, - output_dir=self.output_dir, - image_id_to_depth_path=image_id_to_depth_path, - 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)) - elif require_cameras_exist: - CONSOLE.log(f"[bold red]Could not find existing COLMAP results ({colmap_model_path / 'cameras.bin'}).") - sys.exit(1) - else: - CONSOLE.log("[bold yellow]Warning: could not find existing COLMAP results. Not generating transforms.json") - - CONSOLE.rule("[bold green]:tada: :tada: :tada: All DONE :tada: :tada: :tada:") - - for summary in summary_log: - CONSOLE.print(summary, justify="center") - CONSOLE.rule() - - def _run_colmap(self, image_dir, colmap_dir): - (sfm_tool, feature_type, matcher_type) = process_data_utils.find_tool_feature_matcher_combination( - self.sfm_tool, self.feature_type, self.matcher_type - ) - # check that sfm_tool is hloc if using refine_pixsfm - if self.refine_pixsfm: - assert sfm_tool == "hloc", "refine_pixsfm only works with sfm_tool hloc" - - if sfm_tool == "colmap": - colmap_utils.run_colmap( - image_dir=image_dir, - colmap_dir=colmap_dir, - camera_model=CAMERA_MODELS[self.camera_type], - gpu=self.gpu, - verbose=self.verbose, - matching_method=self.matching_method, - colmap_cmd=self.colmap_cmd, - ) - elif sfm_tool == "hloc": - hloc_utils.run_hloc( - image_dir=image_dir, - colmap_dir=colmap_dir, - camera_model=CAMERA_MODELS[self.camera_type], - verbose=self.verbose, - matching_method=self.matching_method, - feature_type=feature_type, - matcher_type=matcher_type, - refine_pixsfm=self.refine_pixsfm, - ) - else: - CONSOLE.log("[bold red]Invalid combination of sfm_tool, feature_type, and matcher_type, exiting") - sys.exit(1) - - -@dataclass -class ProcessVideo: - """Process videos into a nerfstudio dataset. - - This script does the following: - - 1. Converts the video into images. - 2. Scales images to a specified size. - 3. Calculates the camera poses for each image using `COLMAP `_. - """ - - data: Path - """Path the data, either a video file or a directory of images.""" - output_dir: Path - """Path to the output directory.""" - num_frames_target: int = 300 - """Target number of frames to use for the dataset, results may not be exact.""" - camera_type: Literal["perspective", "fisheye", "equirectangular"] = "perspective" - """Camera model to use.""" - matching_method: Literal["exhaustive", "sequential", "vocab_tree"] = "vocab_tree" - """Feature matching method to use. Vocab tree is recommended for a balance of speed and - accuracy. Exhaustive is slower but more accurate. Sequential is faster but should only be used for videos.""" - sfm_tool: Literal["any", "colmap", "hloc"] = "any" - """Structure from motion tool to use. Colmap will use sift features, hloc can use many modern methods - such as superpoint features and superglue matcher""" - refine_pixsfm: bool = False - """If True, runs refinement using Pixel Perfect SFM. Only has an effect if sfm_tool is hloc""" - feature_type: Literal[ - "any", - "sift", - "superpoint", - "superpoint_aachen", - "superpoint_max", - "superpoint_inloc", - "r2d2", - "d2net-ss", - "sosnet", - "disk", - ] = "any" - """Type of feature to use.""" - matcher_type: Literal[ - "any", "NN", "superglue", "superglue-fast", "NN-superpoint", "NN-ratio", "NN-mutual", "adalam" - ] = "any" - """Matching algorithm.""" - num_downscales: int = 3 - """Number of times to downscale the images. Downscales by 2 each time. For example a value of 3 - will downscale the images by 2x, 4x, and 8x.""" - skip_colmap: bool = False - """If True, skips COLMAP and generates transforms.json if possible.""" - colmap_cmd: str = "colmap" - """How to call the COLMAP executable.""" - images_per_equirect: Literal[8, 14] = 8 - """Number of samples per image to take from each equirectangular image. - Used only when camera-type is equirectangular. - """ - percent_radius_crop: float = 1.0 - """Create circle crop mask. The radius is the percent of the image diagonal.""" - crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0) - """Portion of the image to crop. All values should be in [0,1]. (top, bottom, left, right)""" - crop_bottom: float = 0.0 - """Portion of the image to crop from the bottom. - Can be used instead of `crop-factor 0.0 [num] 0.0 0.0` Should be in [0,1]. - """ - use_sfm_depth: bool = False - """If True, export and use depth maps induced from SfM points.""" - include_depth_debug: bool = False - """If --use-sfm-depth and this flag is True, also export debug images showing SfM overlaid upon input images.""" - gpu: bool = True - """If True, use GPU.""" - verbose: bool = False - """If True, print extra logging.""" - - def main(self) -> None: # pylint: disable=R0915 - """Process video into a nerfstudio dataset.""" - install_checks.check_ffmpeg_installed() - install_checks.check_colmap_installed() - - self.output_dir.mkdir(parents=True, exist_ok=True) - image_dir = self.output_dir / "images" - image_dir.mkdir(parents=True, exist_ok=True) - - if self.crop_bottom > 0.0: - self.crop_factor = (0.0, self.crop_bottom, 0.0, 0.0) - - summary_log = [] - # Convert video to images - if self.camera_type == "equirectangular": - # create temp images folder to store the equirect and perspective images - temp_image_dir = self.output_dir / "temp_images" - temp_image_dir.mkdir(parents=True, exist_ok=True) - summary_log, num_extracted_frames = process_data_utils.convert_video_to_images( - self.data, - image_dir=temp_image_dir, - num_frames_target=self.num_frames_target, - crop_factor=(0.0, 0.0, 0.0, 0.0), - verbose=self.verbose, - ) - else: - summary_log, num_extracted_frames = process_data_utils.convert_video_to_images( - self.data, - image_dir=image_dir, - num_frames_target=self.num_frames_target, - crop_factor=self.crop_factor, - verbose=self.verbose, - ) - - # Generate planar projections if equirectangular - if self.camera_type == "equirectangular": - perspective_image_size = equirect_utils.compute_resolution_from_equirect( - self.output_dir / "temp_images", self.images_per_equirect - ) - image_dir = equirect_utils.generate_planar_projections_from_equirectangular( - self.output_dir / "temp_images", - perspective_image_size, - self.images_per_equirect, - crop_factor=self.crop_factor, - ) - - # copy the perspective images to the image directory - process_data_utils.copy_images( - self.output_dir / "temp_images" / "planar_projections", - image_dir=self.output_dir / "images", - verbose=False, - ) - image_dir = self.output_dir / "images" - - # remove the temp_images folder - shutil.rmtree(self.output_dir / "temp_images", ignore_errors=True) - - # # Create mask - mask_path = process_data_utils.save_mask( - image_dir=image_dir, - num_downscales=self.num_downscales, - crop_factor=(0.0, 0.0, 0.0, 0.0), - percent_radius=self.percent_radius_crop, - ) - if mask_path is not None: - summary_log.append(f"Saved mask to {mask_path}") - - # # Downscale images - summary_log.append(process_data_utils.downscale_images(image_dir, self.num_downscales, verbose=self.verbose)) - - # Run Colmap - colmap_dir = self.output_dir / "colmap" - if not self.skip_colmap: - colmap_dir.mkdir(parents=True, exist_ok=True) - - (sfm_tool, feature_type, matcher_type) = process_data_utils.find_tool_feature_matcher_combination( - self.sfm_tool, self.feature_type, self.matcher_type - ) - - if sfm_tool == "colmap": - colmap_utils.run_colmap( - image_dir=image_dir, - colmap_dir=colmap_dir, - camera_model=CAMERA_MODELS[self.camera_type], - camera_mask_path=mask_path, - gpu=self.gpu, - verbose=self.verbose, - matching_method=self.matching_method, - colmap_cmd=self.colmap_cmd, - ) - elif sfm_tool == "hloc": - if mask_path is not None: - CONSOLE.log( - "[bold red]Cannot use a mask with hloc. Please remove the cropping options and try again." - ) - sys.exit(1) - hloc_utils.run_hloc( - image_dir=image_dir, - colmap_dir=colmap_dir, - camera_model=CAMERA_MODELS[self.camera_type], - verbose=self.verbose, - matching_method=self.matching_method, - feature_type=feature_type, - matcher_type=matcher_type, - refine_pixsfm=self.refine_pixsfm, - ) - else: - CONSOLE.log("[bold red]Invalid combination of sfm_tool, feature_type, and matcher_type, exiting") - sys.exit(1) - - # Export depth maps - if self.use_sfm_depth: - depth_dir = self.output_dir / "depth" - depth_dir.mkdir(parents=True, exist_ok=True) - image_id_to_depth_path = colmap_utils.create_sfm_depth( - recon_dir=colmap_dir / "sparse" / "0", - output_dir=depth_dir, - include_depth_debug=self.include_depth_debug, - input_images_dir=image_dir, - verbose=self.verbose, - ) - summary_log.append( - process_data_utils.downscale_images( - depth_dir, self.num_downscales, folder_name="depths", nearest_neighbor=True, verbose=self.verbose - ) - ) - else: - image_id_to_depth_path = None - - # Save transforms.json - if (colmap_dir / "sparse" / "0" / "cameras.bin").exists(): - with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): - num_matched_frames = colmap_utils.colmap_to_json( - recon_dir=colmap_dir / "sparse" / "0", - output_dir=self.output_dir, - image_id_to_depth_path=image_id_to_depth_path, - 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)) - else: - CONSOLE.log("[bold yellow]Warning: could not find existing COLMAP results. Not generating transforms.json") - - CONSOLE.rule("[bold green]:tada: :tada: :tada: All DONE :tada: :tada: :tada:") - - for summary in summary_log: - CONSOLE.print(summary, justify="center") - CONSOLE.rule() - - -@dataclass -class ProcessRecord3D: +class ProcessRecord3D(BaseConverterToNerfstudioDataset): """Process Record3D data into a nerfstudio dataset. This script does the following: @@ -483,10 +41,6 @@ class ProcessRecord3D: 2. Converts Record3D poses into the nerfstudio format. """ - data: Path - """Path to the record3D data.""" - output_dir: Path - """Path to the output directory.""" num_downscales: int = 3 """Number of times to downscale the images. Downscales by 2 each time. For example a value of 3 will downscale the images by 2x, 4x, and 8x.""" @@ -494,9 +48,6 @@ class ProcessRecord3D: """Max number of images to train on. If the dataset has more, images will be sampled approximately evenly. If -1, use all images.""" - verbose: bool = False - """If True, print extra logging.""" - def main(self) -> None: """Process images into a nerfstudio dataset.""" @@ -551,7 +102,7 @@ class ProcessRecord3D: @dataclass -class ProcessPolycam: +class ProcessPolycam(BaseConverterToNerfstudioDataset): """Process Polycam data into a nerfstudio dataset. To capture data, use the Polycam app on an iPhone or iPad with LiDAR. The capture must be in LiDAR or ROOM mode. @@ -564,10 +115,6 @@ class ProcessPolycam: 2. Converts Polycam poses into the nerfstudio format. """ - data: Path - """Path the polycam export data folder. Can be .zip file or folder.""" - output_dir: Path - """Path to the output directory.""" num_downscales: int = 3 """Number of times to downscale the images. Downscales by 2 each time. For example a value of 3 will downscale the images by 2x, 4x, and 8x.""" @@ -582,8 +129,6 @@ class ProcessPolycam: """Number of pixels to crop from each border of the image. Useful as borders may be black due to undistortion.""" use_depth: bool = False """If True, processes the generated depth maps from Polycam""" - verbose: bool = False - """If True, print extra logging.""" def main(self) -> None: """Process images into a nerfstudio dataset.""" @@ -662,7 +207,15 @@ class ProcessPolycam: @dataclass -class ProcessMetashape: +class _NoDefaultProcessMetashape: + """Private class to order the parameters of ProcessMetashape in the right order for default values.""" + + xml: Path + """Path to the Metashape xml file.""" + + +@dataclass +class ProcessMetashape(BaseConverterToNerfstudioDataset, _NoDefaultProcessMetashape): """Process Metashape data into a nerfstudio dataset. This script assumes that cameras have been aligned using Metashape. After alignment, it is necessary to export the @@ -674,20 +227,12 @@ class ProcessMetashape: 2. Converts Metashape poses into the nerfstudio format. """ - data: Path - """Path to a folder of images.""" - xml: Path - """Path to the Metashape xml file.""" - output_dir: Path - """Path to the output directory.""" num_downscales: int = 3 """Number of times to downscale the images. Downscales by 2 each time. For example a value of 3 will downscale the images by 2x, 4x, and 8x.""" max_dataset_size: int = 600 """Max number of images to train on. If the dataset has more, images will be sampled approximately evenly. If -1, use all images.""" - verbose: bool = False - """If True, print extra logging.""" def main(self) -> None: """Process images into a nerfstudio dataset.""" @@ -749,7 +294,15 @@ class ProcessMetashape: @dataclass -class ProcessRealityCapture: +class _NoDefaultProcessRealityCapture: + """Private class to order the parameters of ProcessRealityCapture in the right order for default values.""" + + csv: Path + """Path to the RealityCapture cameras CSV file.""" + + +@dataclass +class ProcessRealityCapture(BaseConverterToNerfstudioDataset, _NoDefaultProcessRealityCapture): """Process RealityCapture data into a nerfstudio dataset. This script assumes that cameras have been aligned using RealityCapture. After alignment, it is necessary to @@ -761,20 +314,12 @@ class ProcessRealityCapture: 2. Converts RealityCapture poses into the nerfstudio format. """ - data: Path - """Path to a folder of images.""" - csv: Path - """Path to the RealityCapture cameras CSV file.""" - output_dir: Path - """Path to the output directory.""" num_downscales: int = 3 """Number of times to downscale the images. Downscales by 2 each time. For example a value of 3 will downscale the images by 2x, 4x, and 8x.""" max_dataset_size: int = 600 """Max number of images to train on. If the dataset has more, images will be sampled approximately evenly. If -1, use all images.""" - verbose: bool = False - """If True, print extra logging.""" def main(self) -> None: """Process images into a nerfstudio dataset.""" @@ -835,8 +380,8 @@ class ProcessRealityCapture: Commands = Union[ - Annotated[ProcessImages, tyro.conf.subcommand(name="images")], - Annotated[ProcessVideo, tyro.conf.subcommand(name="video")], + Annotated[ImagesToNerfstudioDataset, tyro.conf.subcommand(name="images")], + Annotated[VideoToNerfstudioDataset, tyro.conf.subcommand(name="video")], Annotated[ProcessPolycam, tyro.conf.subcommand(name="polycam")], Annotated[ProcessMetashape, tyro.conf.subcommand(name="metashape")], Annotated[ProcessRealityCapture, tyro.conf.subcommand(name="realitycapture")], @@ -847,7 +392,10 @@ Commands = Union[ def entrypoint(): """Entrypoint for use with pyproject scripts.""" tyro.extras.set_accent_color("bright_yellow") - tyro.cli(Commands).main() + try: + tyro.cli(Commands).main() + except RuntimeError as e: + CONSOLE.log("[bold red]" + e.args[0]) if __name__ == "__main__": diff --git a/tests/process_data/test_process_images.py b/tests/process_data/test_process_images.py index 1b86a5adb..7f27e5cf5 100644 --- a/tests/process_data/test_process_images.py +++ b/tests/process_data/test_process_images.py @@ -16,7 +16,9 @@ from nerfstudio.data.utils.colmap_parsing_utils import ( write_cameras_binary, write_images_binary, ) -from scripts.process_data import ProcessImages +from nerfstudio.process_data.images_to_nerstudio_dataset import ( + ImagesToNerfstudioDataset, +) def random_quaternion(num_poses: int): @@ -79,8 +81,10 @@ def test_process_images_skip_colmap(tmp_path: Path): (tmp_path / "mocked_bin" / "colmap").touch(mode=0o777) (tmp_path / "mocked_bin" / "ffmpeg").touch(mode=0o777) - # Run ProcessImages - cmd = ProcessImages(tmp_path / "images", tmp_path / "nerfstudio", colmap_model_path=sparse_path, skip_colmap=True) + # 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