File size: 6,321 Bytes
d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 27dc37c d2ebea5 29d5834 27dc37c d2ebea5 27dc37c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
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="extracted_videos", frame_output_dir="video"):
"""
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)
|