Debug mode

functions.py

@@ -55,7 +55,7 @@ def save_pixel_values_as_video(pixel_values, output_path):
 
     logger.info(f"Video saved as {output_path}")
 
-def store_chunk(video, start_time, end_time, video_metadata, embedding_model, index):
+def store_chunk(video, start_time, end_time, video_metadata, embedding_model, index, debug=False):
     """
     Processes a video chunk, extracts and stores its embedding, and updates the FAISS index and SQLite database.
     """
@@ -80,11 +80,14 @@ def store_chunk(video, start_time, end_time, video_metadata, embedding_model, in
     logger.debug(f"Generated embedding with shape {embedding.shape}")
 
     # Save the processed frame image for reference
-    path = Path(video_metadata['url'])
-    output_image_path = Path('debug') / f'track_{path.stem}_{start_time}.png'
-    output_image_path.parent.mkdir(parents=True, exist_ok=True)  # Ensure the directory exists
-    vutils.save_image(frame, str(output_image_path))
-    logger.debug(f"Saved processed frame to {output_image_path}")
+    if debug:
+        path = Path(video_metadata['url'])
+        output_image_path = Path('debug') / f'track_{path.stem}_{start_time}.png'
+        output_image_path.parent.mkdir(parents=True, exist_ok=True)
+        vutils.save_image(frame, str(output_image_path))
+        logger.debug(f"Saved processed frame to {output_image_path}")
+    else:
+        logger.debug("Debug mode is OFF, skipping saving debug frame.")
 
     # Insert embedding into FAISS index and SQLite
     video_id = video_metadata['db_id']
@@ -117,7 +120,7 @@ def log_gpu_memory(stage=""):
     else:
         logger.info("CUDA is not available.")
 
-def store_video(video_metadata, embedding_model, index, max_time=1.0):
+def store_video(video_metadata, embedding_model, index, max_time=1.0, debug=False):
     """
     Processes and stores video embeddings up to a specified maximum time.
     """
@@ -209,13 +212,16 @@ def store_video(video_metadata, embedding_model, index, max_time=1.0):
                 pixel_values = remove_watermark_batch(pixel_values, final_x, final_y)
             
             # Save the processed video chunk
-            output_video_path = Path('debug') / f'test_vid_{db_video_metadata["id"]}_{start_time}.mp4'
-            save_pixel_values_as_video(pixel_values, output_video_path)
-            logger.info(f"Saved pixel values as video: {output_video_path}")
+            if debug:
+                output_video_path = Path('debug') / f'test_vid_{db_video_metadata['id']}_{start_time}.mp4'
+                save_pixel_values_as_video(pixel_values, output_video_path)
+                logger.info(f"Saved pixel values as video: {output_video_path}")
+            else:
+                logger.debug("Debug mode is OFF, skipping saving video chunk.")
             
             # Store the embedding chunk
             end_time = start_time + sampling_interval
-            store_chunk(pixel_values, start_time, end_time, video_metadata, embedding_model, index)
+            store_chunk(pixel_values, start_time, end_time, video_metadata, embedding_model, index, debug=debug)
 
             # Explicitly delete frames and pixel_values to free GPU memory
             del frames, pixel_values

main.py

@@ -130,6 +130,13 @@ def store(
             help="Integer subdirectory for concurrent mode. If not specified, a new one will be created automatically.",
         ),
     ] = None,
+    debug: Annotated[
+        bool,
+        typer.Option(
+            "--debug",
+            help="Enable debug mode to save frames and videos to the debug folder.",
+        ),
+    ] = False,
 ):
     """
     Store videos from a directory and/or a CSV file into the vector DB.
@@ -137,6 +144,11 @@ def store(
     # 1. Handle concurrency: Decide if we override DB path and index path
     global INDEX_PATH  # We'll reassign if needed
 
+    if debug:
+        logger.info("Debug mode is ON. Images and videos will be saved in the debug folder.")
+    else:
+        logger.info("Debug mode is OFF. Skipping saving debug frames and videos.")
+
     if concurrent_store:
         logger.info("[CONCURRENT] concurrent_store mode is ON.")
 
@@ -202,7 +214,7 @@ def store(
                         'db_id': None,  # Will be set in store_video
                         'saved_up_to': 0.0,
                     }
-                    store_video(video_metadata, dinov2_vitb14_reg, index, max_time)
+                    store_video(video_metadata, dinov2_vitb14_reg, index, max_time, debug=debug)
                     # Save updated FAISS index
                     faiss.write_index(index, INDEX_PATH)
         logger.info("Storage complete. FAISS index saved.")
@@ -243,7 +255,7 @@ def store(
                     'saved_up_to': 0.0,
                 }
 
-                success = store_video(video_metadata, dinov2_vitb14_reg, index, max_time)
+                success = store_video(video_metadata, dinov2_vitb14_reg, index, max_time, debug=debug)
                 if success:
                     processed_entries += 1
                     # Save updated FAISS index periodically or after each successful insertion