Spaces:
Running
Running
File size: 2,515 Bytes
887f5f0 | 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 | import cv2
import os
import numpy as np
class FrameExtractor:
"""
Extracts frames from video files at regular intervals using OpenCV.
"""
def __init__(self, max_frames: int = 5):
self.max_frames = max_frames
def extract_frames(self, video_path: str) -> list[dict]:
"""
Extracts up to max_frames from the video at even intervals.
Returns a list of dicts containing the frame index, timestamp, and RGB image array.
"""
if not os.path.exists(video_path):
raise FileNotFoundError(f"Video file not found: {video_path}")
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
raise ValueError(f"Could not open video file: {video_path}")
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
fps = cap.get(cv2.CAP_PROP_FPS)
if total_frames <= 0 or fps <= 0:
# Fallback for video streams/files without index information
fps = 25.0
total_frames = 100
duration = total_frames / fps
# Determine step size to extract max_frames
num_to_extract = min(self.max_frames, total_frames)
if num_to_extract <= 1:
indices = [0]
else:
indices = np.linspace(0, total_frames - 1, num_to_extract, dtype=int).tolist()
frames = []
for i, idx in enumerate(indices):
cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
ret, frame = cap.read()
if not ret or frame is None:
continue
# Convert to RGB (OpenCV uses BGR by default)
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# Scale down large frames to prevent excessive memory/CPU usage
height, width = frame_rgb.shape[:2]
max_dim = 960
if max(height, width) > max_dim:
scale = max_dim / max(height, width)
new_width = int(width * scale)
new_height = int(height * scale)
frame_rgb = cv2.resize(frame_rgb, (new_width, new_height), interpolation=cv2.INTER_AREA)
width, height = new_width, new_height
timestamp = idx / fps
frames.append({
"frame_idx": int(idx),
"timestamp": round(timestamp, 2),
"image": frame_rgb,
"width": width,
"height": height
})
cap.release()
return frames
|