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