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 video
def detect_video(source=0):
    """
    Détection d'objets YOLOv8 sur chaque frame d'une vidéo.
    :param source: 0 pour webcam, ou chemin d'une vidéo (.mp4, etc.)
    """
    # Charger le modèle YOLOv8 (n = nano, rapide)
    model = YOLO("yolov8n.pt")

    # Ouvrir la vidéo ou la webcam
    cap = cv2.VideoCapture(source)
    if not cap.isOpened():
        print("Erreur : Impossible d'ouvrir la vidéo ou webcam.")
        return

    while True:
        ret, frame = cap.read()
        if not ret:
            break  # Fin de vidéo

        # Appliquer YOLO sur la frame
        results = model(frame)

        # Annoter la frame avec les résultats
        annotated_frame = results[0].plot()

        # Afficher la frame annotée
        cv2.imshow("Détection vidéo - YOLOv8", annotated_frame)

        # Quitter avec la touche 'q'
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # Libérer les ressources
    cap.release()
    cv2.destroyAllWindows()



demo = gr.Blocks()

#Interface Gradio pour image
image_input = gr.Image(type="numpy", label="Image à analyser")

image_output = gr.Image(type="numpy", label="Image annotée")

interface = gr.Interface(fn=detect_objects_image, inputs=image_input, outputs=image_output, title="Détection sur Image")

# Interface Gradio pour video
interface_video= gr.Interface(
    fn=detect_video,
    inputs=gr.Video(),
    outputs=gr.Video(),
    title="Détection d'objets sur Vidéo (YOLOv8)",
    description="Charge une vidéo pour détecter les bagages avec YOLOv8."
)
# Lancer l'application avec les 2 onglets
gr.TabbedInterface([interface, interface_video],
 ["Image", 'Vidéo']).launch(debug=True)

# interface.launch() # This line is not needed as the TabbedInterface launches the interfaces.