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

# Charger le modèle YOLOv8 pré-entraîné
model = YOLO("yolov8n.pt")

# Fonction pour la détection sur image
def detect_objects_image(img):
    results = model(img)  # Détection
    annotated_frame = results[0].plot()  # Annoter les résultats
    return annotated_frame

# Fonction pour la détection sur vidéo
def detect_objects_video(video):
    # Lire la vidéo uploadée
    cap = cv2.VideoCapture(video)
    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 un fichier temporaire pour enregistrer la sortie
    temp_output = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    output_path = temp_output.name

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

    while cap.isOpened():
        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

# Interface Gradio
with gr.Blocks() as demo:
    gr.Markdown("## Détection d'objets avec YOLOv8")

    with gr.Tab("Image"):
        image_input = gr.Image(type="numpy", label="Image à analyser")
        image_output = gr.Image(type="numpy", label="Image annotée")
        image_btn = gr.Button("Analyser l'image")
        image_btn.click(fn=detect_objects_image, inputs=image_input, outputs=image_output)

    with gr.Tab("Vidéo"):
        video_input = gr.Video(label="Vidéo à analyser")
        video_output = gr.Video(label="Vidéo annotée")
        video_btn = gr.Button("Analyser la vidéo")
        video_btn.click(fn=detect_objects_video, inputs=video_input, outputs=video_output)

demo.launch()