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

# 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 (frame par frame)
def detect_objects_video(video):
    cap = cv2.VideoCapture(video)
    frames = []

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

    cap.release()

    # Convertir la liste de frames en vidéo
    output_path = "annotated_video.mp4"
    height, width, _ = frames[0].shape
    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), 20, (width, height))
    for frame in frames:
        out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
    out.release()

    return output_path

# Interfaces Gradio
image_interface = gr.Interface(
    fn=detect_objects_image,
    inputs=gr.Image(type="numpy", label="Image à analyser"),
    outputs=gr.Image(type="numpy", label="Image annotée"),
    title="Détection sur Image"
)

video_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 sur Vidéo"
)

# Interface avec onglets
demo = gr.TabbedInterface(
    interface_list=[image_interface, video_interface],
    tab_names=["Image", "Vidéo"],
    theme='JohnSmith9982/small_and_pretty'
)

demo.launch()