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

# Chargement du modèle YOLOv8 personnalisé
detector = YOLO("My_best_model.pt")

# Noms de classes et couleurs pour l'affichage
CLASSES = {0: "valise", 1: "sac à dos", 2: "sac à main"}
COLORS = {
    "valise": (255, 0, 0),
    "sac à dos": (0, 255, 0),
    "sac à main": (0, 0, 255)
}

def analyser_video(video_path):
    cap = cv2.VideoCapture(video_path)
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    width, height = int(cap.get(3)), int(cap.get(4))

    output_path = "annotated_output.mp4"
    writer = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))

    suivi_ids = set()
    comptage = {nom: 0 for nom in CLASSES.values()}

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        resultat = detector.track(frame, persist=True, tracker="bytetrack.yaml")

        if resultat[0].boxes.id is not None:
            for box, track_id, classe_id, conf in zip(
                resultat[0].boxes.xyxy.cpu().numpy().astype(int),
                resultat[0].boxes.id.cpu().numpy().astype(int),
                resultat[0].boxes.cls.cpu().numpy().astype(int),
                resultat[0].boxes.conf.cpu().numpy()
            ):
                nom_classe = CLASSES.get(classe_id)
                if nom_classe is None:
                    continue

                if track_id not in suivi_ids:
                    suivi_ids.add(track_id)
                    comptage[nom_classe] += 1

                x1, y1, x2, y2 = box
                couleur = COLORS.get(nom_classe, (255, 255, 255))
                etiquette = f"{nom_classe} ID:{track_id} ({conf:.2f})"

                cv2.rectangle(frame, (x1, y1), (x2, y2), couleur, 2)
                cv2.putText(frame, etiquette, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, couleur, 2)

        # Overlay des comptes
        y = 25
        cv2.rectangle(frame, (10, 10), (300, 150), (0, 0, 0), -1)
        for nom, nb in comptage.items():
            cv2.putText(frame, f"{nom}: {nb}", (20, y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, COLORS[nom], 2)
            y += 30
        total = sum(comptage.values())
        cv2.putText(frame, f"Total: {total}", (20, y), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)

        writer.write(frame)

    cap.release()
    writer.release()

    # Générer le résumé texte directement
    resume = "\n".join([f"{k} : {v}" for k, v in comptage.items()])
    resume += f"\nTotal : {sum(comptage.values())}"
    return output_path, resume

# Gradio : Interface
gr.Interface(
    fn=analyser_video,
    inputs=gr.Video(label="Téléverser une vidéo"),
    outputs=[
        gr.Video(label="Vidéo annotée"),
        gr.Textbox(label="Résumé du comptage")
    ],
    title="Analyse intelligente de bagages avec YOLOv8",
    description="Comptez automatiquement les valises, sacs à dos et sacs à main dans une vidéo. YOLO + ByteTrack intégré."
).launch()