import cv2
import gradio as gr
import tempfile
from ultralytics import YOLO
from deep_sort_realtime.deepsort_tracker import DeepSort
import os

# Load model and tracker
model = YOLO("yolov8n.pt")
tracker = DeepSort(max_age=30)

def track_players(video):
    # Create temp output path
    temp_out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    output_path = temp_out.name

    cap = cv2.VideoCapture(video)
    if not cap.isOpened():
        raise IOError("Video file could not be opened.")

    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0

    out = cv2.VideoWriter(
        output_path,
        cv2.VideoWriter_fourcc(*"avc1"),  # 'avc1' for Linux compatibility
        fps,
        (width, height)
    )

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        results = model(frame)[0]
        detections = []

        for result in results.boxes.data.tolist():
            x1, y1, x2, y2, score, class_id = result
            if int(class_id) == 0 and score > 0.4:
                detections.append(([x1, y1, x2 - x1, y2 - y1], score, 'player'))

        tracks = tracker.update_tracks(detections, frame=frame)

        for track in tracks:
            if not track.is_confirmed():
                continue
            track_id = track.track_id
            l, t, r, b = track.to_ltrb()
            cv2.rectangle(frame, (int(l), int(t)), (int(r), int(b)), (0, 255, 0), 2)
            cv2.putText(frame, f'Player {track_id}', (int(l), int(t - 10)),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)

        out.write(frame)

    cap.release()
    out.release()
    return output_path

demo = gr.Interface(
    fn=track_players,
    inputs=gr.Video(label="Upload Match Video"),
    outputs=gr.Video(label="Tracked Output"),
    title="🏟️ Football Player Tracker",
    description="Upload a match video and get players tracked with consistent IDs."
)

if __name__ == "__main__":
    demo.launch()