File size: 1,603 Bytes
58ac08a 87f3303 58ac08a 87f3303 58ac08a |
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 |
import cv2
import os
from roboflow import Roboflow
from dotenv import load_dotenv
load_dotenv()
## When the ball is no longer detected, we start a new segment
def detect_event_segments(video_path, confidence=0.4):
rf = Roboflow(api_key=os.getenv("ROBOFLOW_API_KEY"))
project = rf.workspace().project("soccer-event-detection")
model = project.version(1).model
cap = cv2.VideoCapture(video_path)
fps = cap.get(cv2.CAP_PROP_FPS)
events = []
active_event = None
frame_data = []
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
frame_number = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
detections = model.predict(frame, confidence=confidence).json().get('predictions', [])
frame_data.append({"frame": frame_number, "objects": detections})
ball_detected = any(obj['class'] == 'ball' for obj in detections)
goal_area_activity = any(obj['class'] == 'goal' for obj in detections) and ball_detected
if goal_area_activity and active_event is None:
active_event = {"start_frame": frame_number, "frames": []}
if active_event:
active_event["frames"].append(frame_data[-1])
if active_event and not ball_detected:
active_event["end_frame"] = frame_number
events.append(active_event)
active_event = None
cap.release()
# Convert frames to timestamps
for event in events:
event['start_sec'] = event['start_frame'] / fps
event['end_sec'] = event['end_frame'] / fps
return events
|