diff --git a/scripts/process_data.py b/scripts/process_data.py index 1b114a782..2bcd6a3d4 100755 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -134,7 +134,7 @@ def get_num_frames_in_video(video: Path) -> int: def convert_video_to_images( video_path: Path, image_dir: Path, num_frames_target: int, verbose: bool = False -) -> List[str]: +) -> Tuple[List[str], int]: """Converts a video into a sequence of images. Args: @@ -143,7 +143,7 @@ def convert_video_to_images( num_frames_target: Number of frames to extract. verbose: If True, logs the output of the command. Returns: - A summary of the conversion. + A tuple containing summary of the conversion and the number of extracted frames. """ with status(msg="Converting video to images...", spinner="bouncingBall", verbose=verbose): @@ -172,12 +172,13 @@ def convert_video_to_images( run_command(ffmpeg_cmd, verbose=verbose) + 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 {len(list(image_dir.glob('*.png')))} images") + summary_log.append(f"We extracted {num_final_frames} images") CONSOLE.log("[bold green]:tada: Done converting video to images.") - return summary_log + return summary_log, num_final_frames def convert_insta360_to_images( @@ -187,7 +188,7 @@ def convert_insta360_to_images( num_frames_target: int, crop_percentage: float = 0.7, verbose: bool = False, -) -> List[str]: +) -> Tuple[List[str], int]: """Converts a video into a sequence of images. Args: @@ -197,7 +198,7 @@ def convert_insta360_to_images( num_frames_target: Number of frames to extract. verbose: If True, logs the output of the command. Returns: - A summary of the conversion. + A tuple containing summary of the conversion and the number of extracted frames. """ with status(msg="Converting video to images...", spinner="bouncingBall", verbose=verbose): @@ -240,12 +241,13 @@ def convert_insta360_to_images( 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 {len(list(image_dir.glob('*.png')))} images") + summary_log.append(f"We extracted {num_final_frames} images") CONSOLE.log("[bold green]:tada: Done converting insta360 to images.") - return summary_log + return summary_log, num_final_frames def copy_images(data: Path, image_dir: Path, verbose) -> int: @@ -514,6 +516,33 @@ def get_insta360_filenames(data: Path) -> Tuple[Path, Path]: return filename_back, filename_front +def get_matching_summary(num_intial_frames: int, num_matched_frames: int) -> str: + """Returns a summary of the matching results. + + Args: + num_intial_frames: The number of initial frames. + num_matched_frames: The number of matched frames. + + Returns: + A summary of the matching results. + """ + match_ratio = num_matched_frames / num_intial_frames + if match_ratio == 1: + return "[bold green]COLAMP found poses for all images, CONGRATS!" + if match_ratio < 0.4: + result = f"[bold red]COLMAP only found poses for {num_matched_frames / num_intial_frames * 100:.2f}%" + result += " of the images. This is low.\nThis can be caused by a variety of reasons," + result += " such poor scene coverage, blurry images, or large exposure changes." + return result + if match_ratio < 0.8: + result = f"[bold yellow]COLMAP only found poses for {num_matched_frames / num_intial_frames * 100:.2f}%" + result += " of the images.\nThis isn't great, but may be ok." + result += "\nMissing poses can be caused by a variety of reasons, such poor scene coverage, blurry images," + result += " or large exposure changes." + return result + return f"[bold green]COLMAP found poses for {num_matched_frames / num_intial_frames * 100:.2f}% of the images." + + @dataclass class ProcessImages: """Process images into a nerfstudio dataset. @@ -585,6 +614,7 @@ class ProcessImages: camera_model=CAMERA_MODELS[self.camera_type], ) summary_log.append(f"Colmap matched {num_matched_frames} images") + summary_log.append(get_matching_summary(num_frames, num_matched_frames)) else: CONSOLE.log("[bold yellow]Warning: could not find existing COLMAP results. Not generating transforms.json") @@ -637,7 +667,7 @@ class ProcessVideo: image_dir.mkdir(parents=True, exist_ok=True) # Convert video to images - summary_log = convert_video_to_images( + summary_log, num_extracted_frames = convert_video_to_images( self.data, image_dir=image_dir, num_frames_target=self.num_frames_target, verbose=self.verbose ) @@ -668,6 +698,7 @@ class ProcessVideo: camera_model=CAMERA_MODELS[self.camera_type], ) summary_log.append(f"Colmap matched {num_matched_frames} images") + summary_log.append(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") @@ -724,7 +755,7 @@ class ProcessInsta360: filename_back, filename_front = get_insta360_filenames(self.data) # Convert video to images - summary_log = convert_insta360_to_images( + summary_log, num_extracted_frames = convert_insta360_to_images( video_front=filename_front, video_back=filename_back, image_dir=image_dir, @@ -759,6 +790,7 @@ class ProcessInsta360: camera_model=CAMERA_MODELS["fisheye"], ) summary_log.append(f"Colmap matched {num_matched_frames} images") + summary_log.append(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")