import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
import tempfile
import os

# Charger le modèle YOLOv8
model = YOLO("yolov8n.pt")

# Fonction pour traiter la vidéo
def detect_objects_video(video_path):
    cap = cv2.VideoCapture(video_path)

    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)

    # Créer une vidéo temporaire de sortie
    temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    out = cv2.VideoWriter(temp_output.name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        results = model(frame)
        annotated = results[0].plot()
        out.write(annotated)

    cap.release()
    out.release()
    return temp_output.name

# Interface Gradio pour vidéo
interface = gr.Interface(
    fn=detect_objects_video,
    inputs=gr.Video(label="Vidéo à analyser"),
    outputs=gr.Video(label="Vidéo annotée"),
    title="Détection d'objets dans une vidéo avec YOLOv8"
)

interface.launch()