Update utils/ingest_video.py

utils/ingest_video.py

@@ -1,122 +1,122 @@
-import gdown
-import zipfile
-import os
-import chromadb
-from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction
-from chromadb.utils.data_loaders import ImageLoader
-
-import cv2
-
-path = "mm_vdb2"
-client = chromadb.PersistentClient(path=path)
-
-image_loader = ImageLoader()
-CLIP = OpenCLIPEmbeddingFunction()
-video_collection = client.get_or_create_collection(
-    name='video_collection',
-    embedding_function=CLIP,
-    data_loader=image_loader
-)
-
-def unzip_file(zip_path, extract_to):
-    """
-    Unzips a zip file to the specified directory.
-    
-    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)
-
-        # Open the zip file
-        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
-            # Extract all the contents
-            zip_ref.extractall(extract_to)
-        print(f"Successfully extracted {zip_path} to {extract_to}")
-    except Exception as e:
-        print(f"An error occurred: {e}")
-
-
-def extract_frames(video_folder, output_folder):
-    if not os.path.exists(output_folder):
-        os.makedirs(output_folder)
-
-    for video_filename in os.listdir(video_folder):
-        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)
-            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)
-
-            success, image = video_capture.read()
-            frame_number = 0
-            while success:
-                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')
-                    cv2.imwrite(output_frame_filename, image)
-
-                success, image = video_capture.read()
-                frame_number += 1
-
-            video_capture.release()
-
-def add_frames_to_chromadb(video_dir, frames_dir):
-    # Dictionary to hold video titles and their corresponding frames
-    video_frames = {}
-
-    # Process each video and associate its frames
-    for video_file in os.listdir(video_dir):
-        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):
-                # 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 = []
-
-    for video_title, frames in video_frames.items():
-        video_path = os.path.join(video_dir, f"{video_title}.mp4")
-        for frame in frames:
-            frame_id = f"{frame[:-4]}_{video_title}"
-            frame_path = os.path.join(frames_dir, video_title, frame)
-            ids.append(frame_id)
-            uris.append(frame_path)
-            metadatas.append({'video_uri': video_path})
-
-    video_collection.add(ids=ids, uris=uris, metadatas=metadatas)
-
-# Running it
-
-
-
-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)
-
-    add_frames_to_chromadb(video_folder_path, output_folder_path)
-    return video_collection
+import gdown
+import zipfile
+import os
+import chromadb
+from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction
+from chromadb.utils.data_loaders import ImageLoader
+
+import cv2
+
+path = "mm_vdb2"
+client = chromadb.PersistentClient(path=path)
+
+image_loader = ImageLoader()
+CLIP = OpenCLIPEmbeddingFunction()
+video_collection = client.get_or_create_collection(
+    name='video_collection',
+    embedding_function=CLIP,
+    data_loader=image_loader
+)
+
+def unzip_file(zip_path, extract_to):
+    """
+    Unzips a zip file to the specified directory.
+    
+    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)
+
+        # Open the zip file
+        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+            # Extract all the contents
+            zip_ref.extractall(extract_to)
+        print(f"Successfully extracted {zip_path} to {extract_to}")
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+
+def extract_frames(video_folder, output_folder):
+    if not os.path.exists(output_folder):
+        os.makedirs(output_folder)
+
+    for video_filename in os.listdir(video_folder):
+        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)
+            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)
+
+            success, image = video_capture.read()
+            frame_number = 0
+            while success:
+                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')
+                    cv2.imwrite(output_frame_filename, image)
+
+                success, image = video_capture.read()
+                frame_number += 1
+
+            video_capture.release()
+
+def add_frames_to_chromadb(video_dir, frames_dir):
+    # Dictionary to hold video titles and their corresponding frames
+    video_frames = {}
+
+    # Process each video and associate its frames
+    for video_file in os.listdir(video_dir):
+        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):
+                # 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 = []
+
+    for video_title, frames in video_frames.items():
+        video_path = os.path.join(video_dir, f"{video_title}.mp4")
+        for frame in frames:
+            frame_id = f"{frame[:-4]}_{video_title}"
+            frame_path = os.path.join(frames_dir, video_title, frame)
+            ids.append(frame_id)
+            uris.append(frame_path)
+            metadatas.append({'video_uri': video_path})
+
+    video_collection.add(ids=ids, uris=uris, metadatas=metadatas)
+
+# Running it
+
+
+
+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)
+
+    add_frames_to_chromadb(video_folder_path, output_folder_path)
+    return video_collection