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

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

# 🔹 Fonction pour détecter sur image
def detect_objects_image(img):
    results = model(img)
    annotated_frame = results[0].plot()
    return annotated_frame

# 🔹 Fonction pour détecter sur vidéo
def detect_objects_video(video_input):
    cap = cv2.VideoCapture(video_input)
    fps = cap.get(cv2.CAP_PROP_FPS)
    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    output_path = temp_output.name
    temp_output.close()

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

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

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

# Créer l'application avec Gradio Blocks
with gr.Blocks(title="YOLOv8 - Détection d'objets sur Image et Vidéo") as app:
    gr.Markdown("## 🧠 Détection d'objets avec YOLOv8 (Image & Vidéo)")
    gr.Markdown("Choisissez une option ci-dessous pour détecter les objets dans une image ou une vidéo.")

    with gr.Tab("📷 Détection sur Image"):
        img_input = gr.Image(type="numpy", label="Image à analyser")
        img_output = gr.Image(type="numpy", label="Image annotée")
        img_button = gr.Button("Lancer la détection")
        img_button.click(fn=detect_objects_image, inputs=img_input, outputs=img_output)

    with gr.Tab("🎥 Détection sur Vidéo"):
        vid_input = gr.Video(label="Vidéo à analyser")
        vid_output = gr.Video(label="Vidéo annotée")
        vid_button = gr.Button("Lancer la détection")
        vid_button.click(fn=detect_objects_video, inputs=vid_input, outputs=vid_output)

# Lancer l'application
app.launch()