|
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'] |
|
|
|
|
|
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 |
|
|
|
|
|
|
|
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) |
|
|
|
|
|
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 |
|
|
|
|
|
|
|
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() |
|
|