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

# Charger le modèle de segmentation
model = YOLO("yolov8n-seg.pt")

# 🔹 Fonction de détection par segmentation sur image
def detect_segmentation_image(img):
    results = model(img)
    annotated_img = results[0].plot()
    return annotated_img

# 🔹 Fonction de détection par segmentation sur vidéo
def detect_segmentation_video(video_path):
    cap = cv2.VideoCapture(video_path)
    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))

    # Créer un fichier vidéo temporaire pour la sortie
    temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    output_path = temp_output.name
    temp_output.close()

    # Définir le codec et le writer
    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

# 🔧 Interface Gradio combinée avec onglets
with gr.Blocks(title="YOLOv8n-seg - Détection par Segmentation") as app:
    gr.Markdown("## 🧠 YOLOv8n Segmentation - Image & Vidéo")
    gr.Markdown("Détection par segmentation automatique avec le modèle **YOLOv8n-seg**.")

    with gr.Tab("📷 Segmentation sur Image"):
        image_input = gr.Image(type="numpy", label="Importer une image")
        image_output = gr.Image(type="numpy", label="Image segmentée")
        image_button = gr.Button("Lancer la détection")
        image_button.click(fn=detect_segmentation_image, inputs=image_input, outputs=image_output)

    with gr.Tab("🎥 Segmentation sur Vidéo"):
        video_input = gr.Video(label="Importer une vidéo")
        video_output = gr.Video(label="Vidéo segmentée")
        video_button = gr.Button("Lancer la détection")
        video_button.click(fn=detect_segmentation_video, inputs=video_input, outputs=video_output)

# 🚀 Lancement de l'application Gradio
app.launch()