Add more debug logs to data processing (#867)

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 @@ def main(self) -> None:
                     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 @@ def main(self) -> None:
         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 @@ def main(self) -> None:
                     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 @@ def main(self) -> None:
         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 @@ def main(self) -> None:
                     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")