mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-13 19:33:46 +08:00
merging from main
This commit is contained in:
terrance committed
commit
02b6a460aa
57 files changed
+35038
-744
No files matched your search
@@ -399,3 +399,6 @@ def entrypoint():
|
||||
|
||||
if __name__ == "__main__":
|
||||
entrypoint()
|
||||
|
||||
# For sphinx docs
|
||||
get_parser_fn = lambda: tyro.extras.get_parser(main) # noqa
|
||||
@@ -359,6 +359,65 @@ class SDFstudioDemoDownload(DatasetDownload):
|
||||
os.remove(download_path)
|
||||
|
||||
|
||||
# pylint: disable=line-too-long
|
||||
nerfosr_downloads = {
|
||||
"europa": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=europa&downloadStartSecret=0k2r95c1fdej",
|
||||
"lk2": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=lk2&downloadStartSecret=w8kuvjzmchc",
|
||||
"lwp": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=lwp&downloadStartSecret=gtnc4vmkcjq",
|
||||
"rathaus": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=rathaus&downloadStartSecret=7372aewy6rr",
|
||||
"schloss": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=schloss&downloadStartSecret=y8t00nqx0h",
|
||||
"st": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=st&downloadStartSecret=kl9ptuxe8v",
|
||||
"stjacob": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=stjacob&downloadStartSecret=sntsim6ebvm",
|
||||
"stjohann": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=stjohann&downloadStartSecret=g80ug1fsbmh",
|
||||
"trevi": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=trevi&downloadStartSecret=ot1483bigjm",
|
||||
"all": None,
|
||||
}
|
||||
|
||||
|
||||
NeRFOSRCaptureName = tyro.extras.literal_type_from_choices(nerfosr_downloads.keys())
|
||||
|
||||
|
||||
@dataclass
|
||||
class NeRFOSRDownload(DatasetDownload):
|
||||
"""Download the NeRF-OSR dataset."""
|
||||
|
||||
capture_name: NeRFOSRCaptureName = "europa"
|
||||
|
||||
def download(self, save_dir: Path):
|
||||
"""Download the NeRF-OSR dataset: https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk"""
|
||||
|
||||
if self.capture_name == "all":
|
||||
for capture_name in nerfosr_downloads:
|
||||
if capture_name != "all":
|
||||
NeRFOSRDownload(capture_name=capture_name).download(save_dir)
|
||||
return
|
||||
|
||||
assert (
|
||||
self.capture_name in nerfosr_downloads
|
||||
), f"Capture name {self.capture_name} not found in {nerfosr_downloads.keys()}"
|
||||
url = nerfosr_downloads[self.capture_name]
|
||||
target_path = str(save_dir / f"NeRF-OSR/Data/{self.capture_name}")
|
||||
os.makedirs(target_path, exist_ok=True)
|
||||
download_path = Path(f"{target_path}.zip")
|
||||
tmp_path = str(save_dir / ".temp")
|
||||
shutil.rmtree(tmp_path, ignore_errors=True)
|
||||
os.makedirs(tmp_path, exist_ok=True)
|
||||
|
||||
os.system(f"curl -L '{url}' > {download_path}")
|
||||
|
||||
# Extract the zip file
|
||||
with zipfile.ZipFile(download_path, "r") as zip_ref:
|
||||
zip_ref.extractall(tmp_path)
|
||||
|
||||
inner_folders = os.listdir(tmp_path)
|
||||
assert len(inner_folders) == 1, "There is more than one folder inside this zip file."
|
||||
folder = os.path.join(tmp_path, inner_folders[0])
|
||||
shutil.rmtree(target_path)
|
||||
shutil.move(folder, target_path)
|
||||
shutil.rmtree(tmp_path)
|
||||
os.remove(download_path)
|
||||
|
||||
|
||||
Commands = Union[
|
||||
Annotated[BlenderDownload, tyro.conf.subcommand(name="blender")],
|
||||
Annotated[Sitcoms3DDownload, tyro.conf.subcommand(name="sitcoms3d")],
|
||||
@@ -367,6 +426,7 @@ Commands = Union[
|
||||
Annotated[DNerfDownload, tyro.conf.subcommand(name="dnerf")],
|
||||
Annotated[PhototourismDownload, tyro.conf.subcommand(name="phototourism")],
|
||||
Annotated[SDFstudioDemoDownload, tyro.conf.subcommand(name="sdfstudio")],
|
||||
Annotated[NeRFOSRDownload, tyro.conf.subcommand(name="nerfosr")],
|
||||
]
|
||||
|
||||
|
||||
|
||||
+1
-145
@@ -2,7 +2,6 @@
|
||||
"""Processes a video or image sequence to a nerfstudio compatible dataset."""
|
||||
|
||||
|
||||
import json
|
||||
import shutil
|
||||
import sys
|
||||
import zipfile
|
||||
@@ -19,7 +18,6 @@ from nerfstudio.process_data import (
|
||||
colmap_utils,
|
||||
equirect_utils,
|
||||
hloc_utils,
|
||||
insta360_utils,
|
||||
metashape_utils,
|
||||
polycam_utils,
|
||||
process_data_utils,
|
||||
@@ -238,6 +236,7 @@ class ProcessImages:
|
||||
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")
|
||||
@@ -457,148 +456,6 @@ class ProcessVideo:
|
||||
CONSOLE.rule()
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProcessInsta360:
|
||||
"""Process Insta360 videos into a nerfstudio dataset. Currently this uses a center crop of the raw data
|
||||
so data at the extreme edges of the video will be lost.
|
||||
|
||||
Expects data from a 2 camera Insta360, single or >2 camera models will not work.
|
||||
(tested with Insta360 One X2)
|
||||
|
||||
This script does the following:
|
||||
|
||||
1. Converts the videos 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, It should be one of the 3 .insv files saved with each capture (Any work)."""
|
||||
output_dir: Path
|
||||
"""Path to the output directory."""
|
||||
num_frames_target: int = 400
|
||||
"""Target number of frames to use for the dataset, results may not be exact."""
|
||||
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."""
|
||||
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."""
|
||||
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:
|
||||
"""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)
|
||||
|
||||
filename_back, filename_front = insta360_utils.get_insta360_filenames(self.data)
|
||||
|
||||
if not filename_back.exists():
|
||||
raise FileNotFoundError(f"Could not find {filename_back}")
|
||||
|
||||
ffprobe_cmd = f'ffprobe -v quiet -print_format json -show_streams -select_streams v:0 "{filename_back}"'
|
||||
|
||||
ffprobe_output = process_data_utils.run_command(ffprobe_cmd)
|
||||
|
||||
assert ffprobe_output is not None
|
||||
ffprobe_decoded = json.loads(ffprobe_output)
|
||||
|
||||
width, height = ffprobe_decoded["streams"][0]["width"], ffprobe_decoded["streams"][0]["height"]
|
||||
|
||||
summary_log, num_extracted_frames = [], 0
|
||||
|
||||
if width / height == 1:
|
||||
if not filename_front.exists():
|
||||
raise FileNotFoundError(f"Could not find {filename_front}")
|
||||
# Convert video to images
|
||||
summary_log, num_extracted_frames = insta360_utils.convert_insta360_to_images(
|
||||
video_front=filename_front,
|
||||
video_back=filename_back,
|
||||
image_dir=image_dir,
|
||||
num_frames_target=self.num_frames_target,
|
||||
verbose=self.verbose,
|
||||
)
|
||||
else:
|
||||
summary_log, num_extracted_frames = insta360_utils.convert_insta360_single_file_to_images(
|
||||
video=filename_back,
|
||||
image_dir=image_dir,
|
||||
num_frames_target=self.num_frames_target,
|
||||
verbose=self.verbose,
|
||||
)
|
||||
|
||||
# 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)
|
||||
|
||||
colmap_utils.run_colmap(
|
||||
image_dir=image_dir,
|
||||
colmap_dir=colmap_dir,
|
||||
camera_model=CAMERA_MODELS["fisheye"],
|
||||
gpu=self.gpu,
|
||||
verbose=self.verbose,
|
||||
matching_method=self.matching_method,
|
||||
colmap_cmd=self.colmap_cmd,
|
||||
)
|
||||
|
||||
# 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,
|
||||
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:
|
||||
"""Process Record3D data into a nerfstudio dataset.
|
||||
@@ -966,7 +823,6 @@ Commands = Union[
|
||||
Annotated[ProcessPolycam, tyro.conf.subcommand(name="polycam")],
|
||||
Annotated[ProcessMetashape, tyro.conf.subcommand(name="metashape")],
|
||||
Annotated[ProcessRealityCapture, tyro.conf.subcommand(name="realitycapture")],
|
||||
Annotated[ProcessInsta360, tyro.conf.subcommand(name="insta360")],
|
||||
Annotated[ProcessRecord3D, tyro.conf.subcommand(name="record3d")],
|
||||
]
|
||||
|
||||
|
||||
Reference in new issue
Block a user