Update utils/ingest_video.py

utils/ingest_video.py

@@ -1,10 +1,12 @@
+import gdown
+import zipfile
 import os
-import cv2
 import chromadb
 from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction
 from chromadb.utils.data_loaders import ImageLoader
 
-# Initialize ChromaDB client and collection
+import cv2
+
 path = "mm_vdb2"
 client = chromadb.PersistentClient(path=path)
 
@@ -16,14 +18,8 @@ video_collection = client.get_or_create_collection(
     data_loader=image_loader
 )
 
+
 def extract_frames(video_folder, output_folder):
-    """
-    Extracts frames from all videos in the video_folder and saves them in the output_folder.
-    
-    Args:
-        video_folder (str): Path to the folder containing video files.
-        output_folder (str): Path to the folder where extracted frames will be saved.
-    """
     if not os.path.exists(output_folder):
         os.makedirs(output_folder)
 
@@ -42,7 +38,6 @@ def extract_frames(video_folder, output_folder):
             success, image = video_capture.read()
             frame_number = 0
             while success:
-                # Save frames at 0 seconds, every 5 seconds, and the last frame
                 if frame_number == 0 or frame_number % int(fps * 5) == 0 or frame_number == frame_count - 1:
                     frame_time = frame_number / fps
                     output_frame_filename = os.path.join(output_subfolder, f'frame_{int(frame_time)}.jpg')
@@ -53,15 +48,7 @@ def extract_frames(video_folder, output_folder):
 
             video_capture.release()
 
-
 def add_frames_to_chromadb(video_dir, frames_dir):
-    """
-    Adds extracted frames from videos to the ChromaDB collection.
-    
-    Args:
-        video_dir (str): Path to the folder containing video files.
-        frames_dir (str): Path to the folder containing the extracted frames.
-    """
     # Dictionary to hold video titles and their corresponding frames
     video_frames = {}
 
@@ -74,7 +61,7 @@ def add_frames_to_chromadb(video_dir, frames_dir):
                 # List all jpg files in the folder
                 video_frames[video_title] = [f for f in os.listdir(frame_folder) if f.endswith('.jpg')]
 
-    # Prepare ids, uris, and metadatas for ChromaDB
+    # Prepare ids, uris and metadatas
     ids = []
     uris = []
     metadatas = []
@@ -88,39 +75,61 @@ def add_frames_to_chromadb(video_dir, frames_dir):
             uris.append(frame_path)
             metadatas.append({'video_uri': video_path})
 
-    # Add frames to the ChromaDB collection
     video_collection.add(ids=ids, uris=uris, metadatas=metadatas)
 
+# Running it
 
-def initiate_video(video_folder_path):
-    """
-    Initiates the video processing pipeline: extracts frames from videos
-    and adds them to the ChromaDB collection.
 
+def unzip_file(zip_path, extract_to):
+    """
+    Unzips a zip file to the specified directory and flattens the folder structure.
+    
     Args:
-        video_folder_path (str): Path to the folder containing video files.
-
-    Returns:
-        The ChromaDB collection with the added frames.
+        zip_path (str): Path to the zip file.
+        extract_to (str): Directory where the contents should be extracted.
     """
     try:
-        print("Starting video processing pipeline...")
+        # Ensure the destination directory exists
+        os.makedirs(extract_to, exist_ok=True)
+
+        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+            for file in zip_ref.namelist():
+                # Extract only files (not directories)
+                if not file.endswith('/'):
+                    # Flatten by extracting all files to the root of extract_to
+                    file_name = os.path.basename(file)
+                    if file_name:  # Ensure it is not an empty string
+                        extracted_path = os.path.join(extract_to, file_name)
+                        with zip_ref.open(file) as source, open(extracted_path, 'wb') as target:
+                            target.write(source.read())
+        
+        print(f"Successfully extracted and flattened {zip_path} to {extract_to}")
+    except Exception as e:
+        print(f"An error occurred during extraction: {e}")
 
-        # Define output folder for extracted frames
-        output_folder_path = os.path.join(video_folder_path, 'extracted_frames')
+def initiate_video():
+    file_id = "1Nzy-ep9Zn15_mLAq8rUi-iKPGOpN8g9Q"
+    output_file = r"StockVideos-CC01.zip"
+    # Download the ZIP file
+    gdown.download(f"https://drive.google.com/uc?id={file_id}", output_file, quiet=False)
+    print(f"File downloaded successfully: {output_file}")
 
-        # Extract frames from videos
-        print("Extracting frames...")
-        extract_frames(video_folder_path, output_folder_path)
-        print("Frames extracted successfully.")
+    # Define paths
+    zip_file_path = output_file
+    flattened_video_folder = r"videos_flattened"
+    frames_output_folder = r"extracted_frames"
 
-        # Add frames to ChromaDB collection
-        print("Adding frames to ChromaDB...")
-        add_frames_to_chromadb(video_folder_path, output_folder_path)
-        print("Frames added to ChromaDB successfully.")
+    # Ensure directories exist
+    os.makedirs(flattened_video_folder, exist_ok=True)
+    os.makedirs(frames_output_folder, exist_ok=True)
 
-        return video_collection
+    # Unzip and flatten the videos
+    unzip_file(zip_file_path, flattened_video_folder)
 
-    except Exception as e:
-        print(f"An error occurred during video processing: {e}")
-        return None
+    # Process the videos and extract frames
+    extract_frames(flattened_video_folder, frames_output_folder)
+
+    # Add frames to ChromaDB
+    add_frames_to_chromadb(flattened_video_folder, frames_output_folder)
+
+    return video_collection