diff --git a/docs/quickstart/custom_dataset.md b/docs/quickstart/custom_dataset.md index 01706b1ce..09e4d1930 100644 --- a/docs/quickstart/custom_dataset.md +++ b/docs/quickstart/custom_dataset.md @@ -5,7 +5,7 @@ Training model on existing datasets is only so fun. If you would like to train o To process your own data run: ```bash -ns-process-data {video,images,polycam,insta360,record3d} --data {DATA_PATH} --output-dir {PROCESSED_DATA_DIR} +ns-process-data {video,images,polycam,record3d} --data {DATA_PATH} --output-dir {PROCESSED_DATA_DIR} ``` A full set of arguments can be found {doc}`here`. @@ -297,10 +297,10 @@ Equirectangular data is data that has been taken by a 360 camera such as Insta36 For a set of equirectangular images, process the data using the following command: ```bash -ns-process-data images --camera-type equirectangular --images-per-equirect {8, or 14} --data {data directory} --output-dir {output directory} +ns-process-data images --camera-type equirectangular --images-per-equirect {8, or 14} --crop-factor {top bottom left right} --data {data directory} --output-dir {output directory} ``` -The images-per-equirect argument is the number of images that will be sampled from each equirectangular image. We have found that 8 images per equirectangular image is sufficient for most use cases so it defaults to that. However, if you find that there isn't enough detail in the nerf or that colmap is having trouble aligning the images, you can try increasing the number of images per equirectangular image to 14. +The images-per-equirect argument is the number of images that will be sampled from each equirectangular image. We have found that 8 images per equirectangular image is sufficient for most use cases so it defaults to that. However, if you find that there isn't enough detail in the nerf or that colmap is having trouble aligning the images, you can try increasing the number of images per equirectangular image to 14. See the video section below for details on cropping. ### Videos diff --git a/nerfstudio/process_data/insta360_utils.py b/nerfstudio/process_data/insta360_utils.py deleted file mode 100644 index 56e554e01..000000000 --- a/nerfstudio/process_data/insta360_utils.py +++ /dev/null @@ -1,190 +0,0 @@ -# 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. - -"""Helper utils for processing insta360 data.""" - -import sys -from pathlib import Path -from typing import List, Tuple - -from rich.console import Console - -from nerfstudio.process_data.process_data_utils import get_num_frames_in_video -from nerfstudio.utils.rich_utils import status -from nerfstudio.utils.scripts import run_command - -CONSOLE = Console(width=120) - - -def get_insta360_filenames(data: Path) -> Tuple[Path, Path]: - """Returns the filenames of the Insta360 videos from a single video file. - - Example input name: VID_20220212_070353_00_003.insv - - Args: - data: Path to a Insta360 file. - - Returns: - The filenames of the Insta360 videios. - """ - if data.suffix != ".insv": - raise ValueError("The input file must be an .insv file.") - file_parts = data.stem.split("_") - - stem_back = f"VID_{file_parts[1]}_{file_parts[2]}_00_{file_parts[4]}.insv" - stem_front = f"VID_{file_parts[1]}_{file_parts[2]}_10_{file_parts[4]}.insv" - - filename_back = data.parent / stem_back - filename_front = data.parent / stem_front - - return filename_back, filename_front - - -def convert_insta360_to_images( - video_front: Path, - video_back: Path, - image_dir: Path, - num_frames_target: int, - crop_percentage: float = 0.7, - verbose: bool = False, -) -> Tuple[List[str], int]: - """Converts a video into a sequence of images. - - Args: - video_front: Path to the front video. - video_back: Path to the back video. - output_dir: Path to the output directory. - num_frames_target: Number of frames to extract. - crop_percentage: Percentage used to calculate the cropped dimensions of extracted frames. Currently used to crop - out the curved portions of the fish-eye lens. - verbose: If True, logs the output of the command. - Returns: - A tuple containing summary of the conversion and the number of extracted frames. - """ - - with status(msg="Converting video to images...", spinner="bouncingBall", verbose=verbose): - # delete existing images in folder - for img in image_dir.glob("*.png"): - if verbose: - CONSOLE.log(f"Deleting {img}") - img.unlink() - - num_frames_front = get_num_frames_in_video(video_front) - num_frames_back = get_num_frames_in_video(video_back) - if num_frames_front == 0: - CONSOLE.print(f"[bold red]Error: Video has no frames: {video_front}") - sys.exit(1) - if num_frames_back == 0: - CONSOLE.print(f"[bold red]Error: Video has no frames: {video_front}") - sys.exit(1) - - spacing = num_frames_front // (num_frames_target // 2) - vf_cmds = [] - if spacing > 1: - vf_cmds = [f"thumbnail={spacing}", "setpts=N/TB"] - else: - CONSOLE.print("[bold red]Can't satisfy requested number of frames. Extracting all frames.") - - vf_cmds.append(f"crop=iw*{crop_percentage}:ih*{crop_percentage}") - - front_vf_cmds = vf_cmds + ["transpose=2"] - back_vf_cmds = vf_cmds + ["transpose=1"] - - front_ffmpeg_cmd = f"ffmpeg -i {video_front} -vf {','.join(front_vf_cmds)} -r 1 {image_dir / 'frame_%05d.png'}" - back_ffmpeg_cmd = ( - f"ffmpeg -i {video_back} -vf {','.join(back_vf_cmds)} -r 1 {image_dir / 'back_frame_%05d.png'}" - ) - - run_command(front_ffmpeg_cmd, verbose=verbose) - run_command(back_ffmpeg_cmd, verbose=verbose) - - num_extracted_front_frames = len(list(image_dir.glob("frame*.png"))) - for i, img in enumerate(image_dir.glob("back_frame_*.png")): - img.rename(image_dir / f"frame_{i+1+num_extracted_front_frames:05d}.png") - - num_final_frames = len(list(image_dir.glob("*.png"))) - summary_log = [] - summary_log.append(f"Starting with {num_frames_front + num_frames_back} video frames") - summary_log.append(f"We extracted {num_final_frames} images") - CONSOLE.log("[bold green]:tada: Done converting insta360 to images.") - - return summary_log, num_final_frames - - -def convert_insta360_single_file_to_images( - video: Path, - image_dir: Path, - num_frames_target: int, - crop_percentage: float = 0.7, - verbose: bool = False, -) -> Tuple[List[str], int]: - """Converts a video into a sequence of images. - - Args: - video: Path to the video. - output_dir: Path to the output directory. - num_frames_target: Number of frames to extract. - crop_percentage: Percentage used to calculate the cropped dimensions of extracted frames. Currently used to crop - out the curved portions of the fish-eye lens. - verbose: If True, logs the output of the command. - Returns: - A tuple containing summary of the conversion and the number of extracted frames. - """ - - with status(msg="Converting video to images...", spinner="bouncingBall", verbose=verbose): - # delete existing images in folder - for img in image_dir.glob("*.png"): - if verbose: - CONSOLE.log(f"Deleting {img}") - img.unlink() - - num_frames = get_num_frames_in_video(video) - if num_frames == 0: - CONSOLE.print(f"[bold red]Error: Video has no frames: {video}") - sys.exit(1) - - spacing = num_frames // (num_frames_target // 2) - vf_cmds = [] - if spacing > 1: - vf_cmds = [f"thumbnail={spacing}", "setpts=N/TB"] - else: - CONSOLE.print("[bold red]Can't satisfy requested number of frames. Extracting all frames.") - - vf_cmds_back = vf_cmds.copy() - vf_cmds_front = vf_cmds.copy() - - vf_cmds_back.append( - f"crop=ih*{crop_percentage}:ih*{crop_percentage}:ih*{crop_percentage}/4:ih*{crop_percentage}/4" - ) - vf_cmds_front.append( - f"crop=ih*{crop_percentage}:ih*{crop_percentage}:iw/2+ih*{crop_percentage/4}:ih*{crop_percentage/4}" - ) - - front_ffmpeg_cmd = f"ffmpeg -i {video} -vf {','.join(vf_cmds_front)} -r 1 {image_dir / 'frame_%05d.png'}" - back_ffmpeg_cmd = f"ffmpeg -i {video} -vf {','.join(vf_cmds_back)} -r 1 {image_dir / 'back_frame_%05d.png'}" - - run_command(back_ffmpeg_cmd, verbose=verbose) - run_command(front_ffmpeg_cmd, verbose=verbose) - - num_extracted_frames = len(list(image_dir.glob("frame*.png"))) - for i, img in enumerate(image_dir.glob("back_frame_*.png")): - img.rename(image_dir / f"frame_{i+1+num_extracted_frames:05d}.png") - - num_final_frames = len(list(image_dir.glob("*.png"))) - summary_log = [] - summary_log.append(f"Starting with {num_frames} video frames") - summary_log.append(f"We extracted {num_final_frames} images") - CONSOLE.log("[bold green]:tada: Done converting insta360 to images.") - - return summary_log, num_final_frames diff --git a/scripts/process_data.py b/scripts/process_data.py index 34d3b2fb1..2eaa9e11c 100644 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -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, @@ -458,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 `_. - """ - - 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. @@ -967,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")], ]