blip / src /utils.py
amezi's picture
minor bux fixed
87f3303
raw
history blame
1.77 kB
import cv2
import os
from dotenv import load_dotenv
load_dotenv()
def extract_key_frames(video_path, start_sec, end_sec):
cap = cv2.VideoCapture(video_path)
fps = cap.get(cv2.CAP_PROP_FPS)
start_frame = int(start_sec * fps)
end_frame = int(end_sec * fps)
mid_frame = (start_frame + end_frame) // 2
frames = []
for frame_number in [start_frame, mid_frame, end_frame]:
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
ret, frame = cap.read()
if ret:
frames.append(frame)
cap.release()
return frames
def save_frames_locally(frames, event_id):
os.makedirs("/data", exist_ok=True)
frame_paths = []
for idx, frame in enumerate(frames):
path = f"/data/frame_{event_id}_{idx}.jpg"
cv2.imwrite(path, frame)
frame_paths.append(path)
return frame_paths
def generate_frame_urls(frame_paths):
base_url = os.getenv("SPACE_URL", "http://localhost:8000")
return [f"{base_url}/data/{os.path.basename(path)}" for path in frame_paths]
def match_transcript_to_events(events, transcript):
for event in events:
matched_lines = [
line["text"] for line in transcript
if line["start_sec"] <= event["end_sec"] and line["end_sec"] >= event["start_sec"]
]
event["transcript"] = "\n".join(matched_lines) or "(No matching commentary)"
return events
def clip_video_segment(video_path, start_sec, end_sec, event_id):
output_path = f"/data/clip_{event_id}.mp4"
duration = end_sec - start_sec
command = [
"ffmpeg", "-y",
"-ss", str(start_sec),
"-i", video_path,
"-t", str(duration),
"-c", "copy", output_path
]
os.system(" ".join(command))
return output_path