tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-13 03:14:00 +08:00
Code Activity

Refactor of image and video dataset conversion (#1690)

Browse Source 
* 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 <yibrenth@gmail.com>
This commit is contained in:
Malcolm MielleandBrent Yi authored and GitHub committed 2023-04-26 10:30:01 -07:00
1 parent 281615e638
commit 3d2e62a480
6 files changed
+544 -490
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

scripts/process_data.py
+35 -487
View File
@@ -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 <https://colmap.github.io/>`_.
"""
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 <https://colmap.github.io/>`_.
"""
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__":