Update utils/ingest_video.py

utils/ingest_video.py

@@ -1,12 +1,11 @@
-import gdown
 import zipfile
 import os
 import chromadb
 from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction
 from chromadb.utils.data_loaders import ImageLoader
-
 import cv2
 
+# Initialize ChromaDB Persistent Client
 path = "mm_vdb2"
 client = chromadb.PersistentClient(path=path)
 
@@ -18,46 +17,65 @@ video_collection = client.get_or_create_collection(
     data_loader=image_loader
 )
 
-def unzip_file(zip_path, extract_to):
+def unzip_file_flat(zip_path, extract_to):
     """
-    Unzips a zip file to the specified directory.
-    
+    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:
-        # Ensure the destination directory exists
-        os.makedirs(extract_to, exist_ok=True)
+        os.makedirs(extract_to, exist_ok=True)  # Ensure the output directory exists
 
-        # Open the zip file
         with zipfile.ZipFile(zip_path, 'r') as zip_ref:
-            # Extract all the contents
-            zip_ref.extractall(extract_to)
+            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: {e}")
-
+        print(f"An error occurred while unzipping {zip_path}: {e}")
 
 def extract_frames(video_folder, output_folder):
-    if not os.path.exists(output_folder):
-        os.makedirs(output_folder)
+    """
+    Extracts frames from video files at regular intervals.
+
+    Args:
+        video_folder (str): Folder containing video files.
+        output_folder (str): Folder to save extracted frames.
+    """
+    os.makedirs(output_folder, exist_ok=True)
 
     for video_filename in os.listdir(video_folder):
-        if video_filename.endswith('.mp4'):
+        if video_filename.lower().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])
-            if not os.path.exists(output_subfolder):
-                os.makedirs(output_subfolder)
+            os.makedirs(output_subfolder, exist_ok=True)
 
             success, image = video_capture.read()
             frame_number = 0
             while success:
+                # Extract frame at the start, 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')
@@ -67,25 +85,27 @@ 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):
-    # Dictionary to hold video titles and their corresponding frames
+    """
+    Adds metadata and URIs of video frames to ChromaDB collection.
+
+    Args:
+        video_dir (str): Directory containing original video files.
+        frames_dir (str): Directory containing extracted frames.
+    """
     video_frames = {}
 
-    # Process each video and associate its frames
+    # Since all video files are directly in video_dir, no subfolders
     for video_file in os.listdir(video_dir):
-        if video_file.endswith('.mp4'):
-            video_title = video_file[:-4]
+        if video_file.lower().endswith('.mp4'):
+            video_title = os.path.splitext(video_file)[0]
             frame_folder = os.path.join(frames_dir, video_title)
             if os.path.exists(frame_folder):
-                # 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
-    ids = []
-    uris = []
-    metadatas = []
+                video_frames[video_title] = [f for f in os.listdir(frame_folder) if f.lower().endswith('.jpg')]
 
+    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:
@@ -95,28 +115,41 @@ def add_frames_to_chromadb(video_dir, frames_dir):
             uris.append(frame_path)
             metadatas.append({'video_uri': video_path})
 
-    video_collection.add(ids=ids, uris=uris, metadatas=metadatas)
-
-# Running it
-
-
+    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.")
 
-def intiate_video():
-    file_id = "1Fm8Cge1VM4w8fmE0cZfRKhIQV0UgBXzp"
-    output_file = r"video/StockVideos-CC01.zip"
-    gdown.download(f"https://drive.google.com/uc?id={file_id}", output_file, quiet=False)
-
-    print(f"File downloaded successfully: {output_file}")
-    # Example Usage
-    zip_file_path = r"video/StockVideos-CC01.zip"
-    destination_folder = r"video"
-    unzip_file(zip_file_path, destination_folder)
-    print("Unzipped")
-    video_folder_path = r'video/StockVideos-CC0'
-    output_folder_path = r'video/StockVideos-CC0-frames'
-
-    extract_frames(video_folder_path, output_folder_path)
+def process_uploaded_files(zip_files, extract_dir, frame_output_dir):
+    """
+    Processes uploaded zip files by extracting their contents, processing video files, 
+    and adding their frames to ChromaDB.
 
-    add_frames_to_chromadb(video_folder_path, output_folder_path)
-    return video_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.
+    """
+    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)