src/segmenter.py

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