utils/ingest_video.py

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)

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_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.

    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.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))

            output_subfolder = os.path.join(output_folder, os.path.splitext(video_filename)[0])
            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')
                    cv2.imwrite(output_frame_filename, image)

                success, image = video_capture.read()
                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.

    Args:
        video_dir (str): Directory containing original video files.
        frames_dir (str): Directory containing extracted frames.
    """
    video_frames = {}

    # Since all video files are directly in video_dir, no subfolders
    for video_file in os.listdir(video_dir):
        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):
                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:
            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})

    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 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.

    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)