Iralion's picture
Update app.py
6e1bc04 verified
import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
import tempfile
import os
import shutil
model = YOLO("best.pt")
bagages_classes = ['suitcase', 'backpack', 'handbag']
# Fonction de détection sur image
def detect_objects_image(image):
results = model(image)[0]
annotated_image = image.copy()
count_by_class = {cls: 0 for cls in bagages_classes}
colors = {
'suitcase': (0, 255, 0),
'backpack': (255, 0, 0),
'handbag': (0, 128, 255)
}
for box in results.boxes:
class_id = int(box.cls)
class_name = model.names[class_id]
if class_name in bagages_classes:
x1, y1, x2, y2 = map(int, box.xyxy[0])
conf = float(box.conf[0])
color = colors.get(class_name, (255, 255, 255))
label = f"{class_name} ({conf:.2f})"
(tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
cv2.rectangle(annotated_image, (x1, y1), (x2, y2), color, 2)
cv2.rectangle(annotated_image, (x1, y1 - th - 10), (x1 + tw + 4, y1), color, -1)
cv2.putText(annotated_image, label, (x1 + 2, y1 - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 2)
count_by_class[class_name] += 1
resume = "📷 Résumé détection image:\n"
total = sum(count_by_class.values())
resume += f"Total bagages détectés : {total}\n"
for cls, count in count_by_class.items():
resume += f"🧳{cls}: {count}\n"
return annotated_image, resume
# Fonction de détection de vidéo
def detect_objects_video(video):
import time
import cv2
import os
import shutil
import tempfile
video_path = video if isinstance(video, str) else video.name
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
print("Impossible d'ouvrir la vidéo.")
return None
fps = cap.get(cv2.CAP_PROP_FPS)
W = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
H = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
midline = W // 2
with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp:
temp_output_path = tmp.name
out = cv2.VideoWriter(temp_output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (W, H))
counted_ids_direction = set()
count_by_class = {cls: 0 for cls in bagages_classes}
previous_positions = {}
entrants = sortants = 0
class_colors = {
'suitcase': (0, 255, 0),
'backpack': (255, 0, 0),
'handbag': (0, 0, 255)
}
results = model.track(source=video_path, persist=True, stream=True, tracker="botsort.yaml", device='cpu')
for result in results:
frame = getattr(result, 'orig_img', None)
boxes = getattr(result, 'boxes', None)
if frame is None:
continue
cv2.line(frame, (midline, 0), (midline, H), (0, 0, 255), 2)
if boxes is not None and boxes.id is not None:
ids = boxes.id.cpu().numpy().astype(int)
clss = boxes.cls.cpu().numpy().astype(int)
confs = boxes.conf.cpu().numpy()
coords = boxes.xyxy.cpu().numpy()
for obj_id, cls_id, box, conf in zip(ids, clss, coords, confs):
class_name = model.names[int(cls_id)]
if class_name not in bagages_classes:
continue
x1, y1, x2, y2 = map(int, box)
cx = int((x1 + x2) / 2)
prev_cx = previous_positions.get(obj_id, cx)
direction = None
if prev_cx < midline <= cx:
direction = 'in'
elif prev_cx > midline >= cx:
direction = 'out'
previous_positions[obj_id] = cx
if direction:
key = (obj_id, direction)
if key not in counted_ids_direction:
counted_ids_direction.add(key)
if direction == 'in':
entrants += 1
else:
sortants += 1
count_by_class[class_name] += 1
color = class_colors.get(class_name, (255, 255, 255))
label = f"{class_name} #{obj_id} {conf*100:.1f}%"
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
(tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
cv2.rectangle(frame, (x1, y1 - 20), (x1 + tw + 6, y1), color, -1)
cv2.putText(frame, label, (x1 + 3, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
# Résumé dans cadre
total = len(counted_ids_direction)
line1 = f"Total: {total} Entrants: {entrants} Sortants: {sortants}"
line2 = " ".join([f"{cls}: {count_by_class[cls]}" for cls in bagages_classes])
(w1, h1), _ = cv2.getTextSize(line1, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
(w2, h2), _ = cv2.getTextSize(line2, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
pad_x, pad_y = 10, 10
text_height = h1 + h2 + pad_y * 3
text_width = max(w1, w2) + pad_x * 2
top_left = (10, 10)
bottom_right = (top_left[0] + text_width, top_left[1] + text_height)
cv2.rectangle(frame, top_left, bottom_right, (50, 50, 50), -1)
org1 = (top_left[0] + pad_x, top_left[1] + h1 + pad_y)
cv2.putText(frame, line1, org1, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 3)
cv2.putText(frame, line1, org1, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)
org2 = (top_left[0] + pad_x, org1[1] + h2 + pad_y)
cv2.putText(frame, line2, org2, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 3)
cv2.putText(frame, line2, org2, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)
out.write(frame)
cap.release()
out.release()
os.makedirs("videos_sortie", exist_ok=True)
final_output_path = os.path.join("videos_sortie", "bagages_detectes.mp4")
shutil.copy(temp_output_path, final_output_path)
os.remove(temp_output_path)
resume = f"🎥 Résumé vidéo :\n"
resume += f"🎒 Total bagages détectés (franchissement ligne rouge): {len(counted_ids_direction)}\n"
for cls in bagages_classes:
resume += f"🧳 {cls}: {count_by_class[cls]}\n"
resume += f"\n✅ Entrants: {entrants}\n🚪 Sortants: {sortants}"
return final_output_path, resume, final_output_path
# Interface Gradio
with gr.Blocks(theme="SebastianBravo/simci_css") as demo:
gr.Markdown(
"""
<h1 style="text-align: center;">🎯 Détection et traçage de bagages</h1>
<p style="text-align: center;">Détecte et compte automatiquement les bagages avec YOLOv8 : <strong>suitcase</strong>, <strong>backpack</strong>, <strong>handbag</strong>.</p>
<p style="text-align: center;">Téléverse une image ou une vidéo, et laisse l'IA faire le travail 🧠</p>
"""
)
with gr.Tab("🖼️Détection sur Image"):
with gr.Row():
with gr.Column(scale=1):
image_input = gr.Image(type="numpy", label="📥 Charger une image")
image_btn = gr.Button("🔍 Lancer la détection")
gr.Examples(
examples=[
"Image1.jpg",
"Image2.jpg",
"Image3.jpg"
],
inputs=image_input,
label="📂 Exemples d’images"
)
with gr.Column(scale=1):
image_output = gr.Image(type="numpy", label="📸 Image détectée")
image_summary = gr.Textbox(label="📋 Résumé", lines=4, interactive=False)
image_btn.click(fn=detect_objects_image, inputs=image_input, outputs=[image_output, image_summary])
with gr.Tab("🎞️ Détection sur Vidéo"):
with gr.Row():
with gr.Column(scale=1):
video_input = gr.Video(label="📥 Charger une vidéo (.mp4)")
video_btn = gr.Button("🔎 Lancer la détection vidéo")
with gr.Column(scale=1):
video_output = gr.Video(label="🎬 Vidéo traitée")
video_info = gr.Textbox(label="📊 Résumé", lines=6, interactive=False)
video_download = gr.File(label="⬇️ Télécharger la vidéo")
def run_video_pipeline(video):
output_path, resume, file_path = detect_objects_video(video)
return output_path, resume, file_path
video_btn.click(fn=run_video_pipeline, inputs=video_input, outputs=[video_output, video_info, video_download])
with gr.Accordion("ℹ️ Astuce", open=False):
gr.Markdown("""
**🔎 Bagages détectés :**
- Valise (`suitcase`)
- Sac à dos (`backpack`)
- Sac à main (`handbag`)
**Fonctionnalités de l'application :**
✅ Suivi des bagages en temps réel
✅ Comptage unique après franchissement de la ligne rouge
✅ Résumé en texte + vidéo annotée
✅ Téléchargement de la vidéo traitée
**Cas d'usage sécurité :**
🧳 Surveillance de bagages :
- Compter les bagages entrants/sortants
- Détection d’abandon ou de dépôt suspect
🚫 Zones sensibles :
- Détection de franchissement interdit
- Alerte en cas d’objet abandonné ou non autorisé
""")
demo.launch()