Spaces:
Running
Running
File size: 3,007 Bytes
6cc042b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
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()
|