Update utils/ingest_video.py

utils/ingest_video.py

@@ -1,11 +1,10 @@
-import zipfile
 import os
+import cv2
 import chromadb
 from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction
 from chromadb.utils.data_loaders import ImageLoader
-import cv2
 
-# Initialize ChromaDB Persistent Client
+# Initialize ChromaDB client and collection
 path = "mm_vdb2"
 client = chromadb.PersistentClient(path=path)
 
@@ -17,65 +16,33 @@ video_collection = client.get_or_create_collection(
     data_loader=image_loader
 )
 
-def unzip_file_flat(zip_path, extract_to):
-    """
-    Unzips a zip file and extracts all files into a single directory, ignoring internal folders.
-
-    Args:
-        zip_path (str): Path to the zip file.
-        extract_to (str): Directory where the contents should be extracted.
-    """
-    try:
-        os.makedirs(extract_to, exist_ok=True)  # Ensure the output directory exists
-
-        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
-            for file in zip_ref.namelist():
-                if not file.endswith('/'):  # Ignore directories
-                    file_name = os.path.basename(file)
-                    if not file_name:  # Skip if it's a directory entry
-                        continue
-                    extracted_path = os.path.join(extract_to, file_name)
-
-                    # Handle filename conflicts by appending a number to duplicates
-                    base_name, ext = os.path.splitext(file_name)
-                    counter = 1
-                    while os.path.exists(extracted_path):
-                        extracted_path = os.path.join(extract_to, f"{base_name}_{counter}{ext}")
-                        counter += 1
-
-                    with zip_ref.open(file) as source, open(extracted_path, 'wb') as target:
-                        target.write(source.read())
-        print(f"Successfully extracted {zip_path} to {extract_to}")
-    except Exception as e:
-        print(f"An error occurred while unzipping {zip_path}: {e}")
-
 def extract_frames(video_folder, output_folder):
     """
-    Extracts frames from video files at regular intervals.
-
+    Extracts frames from all videos in the video_folder and saves them in the output_folder.
+    
     Args:
-        video_folder (str): Folder containing video files.
-        output_folder (str): Folder to save extracted frames.
+        video_folder (str): Path to the folder containing video files.
+        output_folder (str): Path to the folder where extracted frames will be saved.
     """
-    os.makedirs(output_folder, exist_ok=True)
+    if not os.path.exists(output_folder):
+        os.makedirs(output_folder)
 
     for video_filename in os.listdir(video_folder):
-        if video_filename.lower().endswith('.mp4'):
+        if video_filename.endswith('.mp4'):
             video_path = os.path.join(video_folder, video_filename)
             video_capture = cv2.VideoCapture(video_path)
             fps = video_capture.get(cv2.CAP_PROP_FPS)
-            if fps == 0:
-                print(f"Warning: FPS is zero for {video_filename}. Skipping.")
-                continue
             frame_count = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))
+            duration = frame_count / fps
 
             output_subfolder = os.path.join(output_folder, os.path.splitext(video_filename)[0])
-            os.makedirs(output_subfolder, exist_ok=True)
+            if not os.path.exists(output_subfolder):
+                os.makedirs(output_subfolder)
 
             success, image = video_capture.read()
             frame_number = 0
             while success:
-                # Extract frame at the start, every 5 seconds, and the last frame
+                # 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')
@@ -85,27 +52,33 @@ def extract_frames(video_folder, output_folder):
                 frame_number += 1
 
             video_capture.release()
-            print(f"Frames extracted from {video_filename} to {output_subfolder}")
+
 
 def add_frames_to_chromadb(video_dir, frames_dir):
     """
-    Adds metadata and URIs of video frames to ChromaDB collection.
-
+    Adds extracted frames from videos to the ChromaDB collection.
+    
     Args:
-        video_dir (str): Directory containing original video files.
-        frames_dir (str): Directory containing extracted frames.
+        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 = {}
 
-    # Since all video files are directly in video_dir, no subfolders
+    # Process each video and associate its frames
     for video_file in os.listdir(video_dir):
-        if video_file.lower().endswith('.mp4'):
-            video_title = os.path.splitext(video_file)[0]
+        if video_file.endswith('.mp4'):
+            video_title = video_file[:-4]
             frame_folder = os.path.join(frames_dir, video_title)
             if os.path.exists(frame_folder):
-                video_frames[video_title] = [f for f in os.listdir(frame_folder) if f.lower().endswith('.jpg')]
+                # 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
+    ids = []
+    uris = []
+    metadatas = []
 
-    ids, uris, metadatas = [], [], []
     for video_title, frames in video_frames.items():
         video_path = os.path.join(video_dir, f"{video_title}.mp4")
         for frame in frames:
@@ -115,41 +88,39 @@ def add_frames_to_chromadb(video_dir, frames_dir):
             uris.append(frame_path)
             metadatas.append({'video_uri': video_path})
 
-    if ids:
-        video_collection.add(ids=ids, uris=uris, metadatas=metadatas)
-        print(f"Added {len(ids)} frames to ChromaDB collection from {len(video_frames)} videos.")
-    else:
-        print("No frames to add to ChromaDB.")
+    # Add frames to the ChromaDB collection
+    video_collection.add(ids=ids, uris=uris, metadatas=metadatas)
+
 
-def process_uploaded_files(zip_files, extract_dir="extracted_videos", frame_output_dir="video"):
+def initiate_video(video_folder_path):
     """
-    Processes uploaded zip files by extracting their contents, processing video files, 
-    and adding their frames to ChromaDB.
+    Initiates the video processing pipeline: extracts frames from videos
+    and adds them to the ChromaDB collection.
 
     Args:
-        zip_files (list): List of paths to uploaded zip files.
-        extract_dir (str): Directory to extract zip contents.
-        frame_output_dir (str): Directory to store extracted frames.
+        video_folder_path (str): Path to the folder containing video files.
+
+    Returns:
+        The ChromaDB collection with the added frames.
     """
-    os.makedirs(extract_dir, exist_ok=True)
-    os.makedirs(frame_output_dir, exist_ok=True)
-
-    for zip_file in zip_files:
-        print(f"Processing {zip_file}...")
-        unzip_file_flat(zip_file, extract_dir)
-
-    # After extraction, all video files are directly in extract_dir
-    extract_frames(extract_dir, frame_output_dir)
-    add_frames_to_chromadb(extract_dir, frame_output_dir)
-
-# # Example Usage
-# if __name__ == "__main__":
-#     # Example list of uploaded zip file paths
-#     uploaded_files = [
-#         "uploaded/video_package1.zip",
-#         "uploaded/video_package2.zip"
-#     ]
-#     extract_dir = "extracted_videos"      # All videos extracted here directly
-#     frame_output_dir = "video_frames"     # All frames stored here
-
-#     process_uploaded_files(uploaded_files, extract_dir, frame_output_dir)
+    try:
+        print("Starting video processing pipeline...")
+
+        # Define output folder for extracted frames
+        output_folder_path = os.path.join(video_folder_path, 'extracted_frames')
+
+        # Extract frames from videos
+        print("Extracting frames...")
+        extract_frames(video_folder_path, output_folder_path)
+        print("Frames extracted successfully.")
+
+        # 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.")
+
+        return video_collection
+
+    except Exception as e:
+        print(f"An error occurred during video processing: {e}")
+        return None