Spaces:
Sleeping
Sleeping
| """ | |
| Streamlit Web Interface – Traffic Detection | |
| Supports: YOLOv11 (tracking) | SSD MobileNetV3 | |
| GPU forced automatically. | |
| Upload and processing are separated to avoid 403 errors. | |
| CSV logs follow the shared schema. | |
| """ | |
| import streamlit as st | |
| import cv2 as cv | |
| import tempfile | |
| import os | |
| import csv | |
| import torch | |
| import pandas as pd | |
| import plotly.express as px | |
| from ultralytics import YOLO | |
| from datetime import datetime | |
| from utils.yolo_tracker import normalize_class_name | |
| # ── Session state init ───────────────────────────────────────────────────────── | |
| if 'global_unique_ids' not in st.session_state: | |
| st.session_state.global_unique_ids = {} | |
| if 'video_ready' not in st.session_state: | |
| st.session_state.video_ready = False | |
| if 'tmp_path' not in st.session_state: | |
| st.session_state.tmp_path = None | |
| if 'processing_done' not in st.session_state: | |
| st.session_state.processing_done = False | |
| def refine_class_by_shape(cls_name, x1, y1, x2, y2, ratio=1.2): | |
| w = max(1, x2 - x1) | |
| h = max(1, y2 - y1) | |
| if cls_name == "car" and (h / w) > ratio: | |
| return "person" | |
| return cls_name | |
| # ── GPU Detection ────────────────────────────────────────────────────────────── | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| GPU_NAME = torch.cuda.get_device_name(0) if DEVICE == "cuda" else "CPU only" | |
| TRAFFIC_CLASSES = ['car', 'truck', 'bus', 'motorbike', 'bicycle', | |
| 'person', 'traffic sign', 'traffic light'] | |
| YOLO_BASE_PATH = "models/yolo11n.pt" | |
| YOLO_FINETUNED_PATH = "models/best.pt" | |
| os.makedirs("logs", exist_ok=True) | |
| st.set_page_config(page_title="Traffic Monitor", layout="wide") | |
| st.title("🚦 Traffic Object Detection & Tracking") | |
| if DEVICE == "cuda": | |
| st.success(f"⚡ GPU Active : {GPU_NAME}") | |
| else: | |
| st.warning("⚠️ GPU not available – running on CPU") | |
| # ── Sidebar ──────────────────────────────────────────────────────────────────── | |
| st.sidebar.header("⚙️ Configuration") | |
| model_choice = st.sidebar.radio( | |
| "Model", | |
| ["YOLOv11 + ByteTrack (Base)", | |
| "YOLOv11 + ByteTrack (Fine-tuned)", | |
| "SSD MobileNetV3"], | |
| ) | |
| if model_choice == "YOLOv11 + ByteTrack (Base)": | |
| selected_yolo_path = YOLO_BASE_PATH | |
| st.sidebar.info("Weights: `models/yolo11n.pt`") | |
| elif model_choice == "YOLOv11 + ByteTrack (Fine-tuned)": | |
| selected_yolo_path = YOLO_FINETUNED_PATH | |
| if os.path.exists(selected_yolo_path): | |
| st.sidebar.success("Weights: `models/best.pt` ✅") | |
| else: | |
| st.sidebar.error("`models/best.pt` not found.") | |
| else: | |
| selected_yolo_path = None | |
| selected_classes = st.sidebar.multiselect( | |
| "Classes to detect", TRAFFIC_CLASSES, | |
| default=['car', 'truck', 'bus', 'person', 'traffic light'] | |
| ) | |
| confidence_threshold = st.sidebar.slider("Confidence threshold", 0.1, 1.0, 0.4) | |
| min_box_area = st.sidebar.number_input("Min box area (px²)", min_value=0, value=1600, step=100) | |
| min_track_hits = st.sidebar.number_input("Min track hits (YOLO only)", min_value=1, value=5, step=1) | |
| shape_refine = st.sidebar.checkbox("Reclassify tall 'car' as person", value=True) | |
| frame_skip = st.sidebar.slider("Process every N frames", 1, 10, 2) | |
| person_car_ratio = st.sidebar.slider("Person/Car ratio (H/W)", 0.8, 2.0, 1.2, 0.05) | |
| # ── Métadonnées pour le schéma partagé ──────────────────────────────────────── | |
| st.sidebar.divider() | |
| st.sidebar.subheader("📋 Log Metadata") | |
| scene_name = st.sidebar.text_input("Scene name", value="scene_01", | |
| placeholder="intersection_A") | |
| group_id = st.sidebar.text_input("Group ID", value="group_01", | |
| placeholder="group_01") | |
| if not selected_classes: | |
| st.warning("Select at least one class in the sidebar.") | |
| st.stop() | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| # ÉTAPE 1 — UPLOAD | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| st.subheader("📤 Step 1 — Upload your video") | |
| uploaded_file = st.file_uploader( | |
| "Choose a video file (mp4, avi, mov)", | |
| type=["mp4", "avi", "mov"] | |
| ) | |
| if uploaded_file is not None: | |
| save_path = "temp_video.mp4" | |
| if not st.session_state.video_ready or st.session_state.tmp_path != save_path: | |
| with st.spinner("Saving video to server..."): | |
| with open(save_path, "wb") as f: | |
| uploaded_file.seek(0) | |
| while True: | |
| chunk = uploaded_file.read(1024 * 1024) # 1MB | |
| if not chunk: | |
| break | |
| f.write(chunk) | |
| st.session_state.tmp_path = save_path | |
| st.session_state.video_ready = True | |
| st.session_state.global_unique_ids = {} | |
| st.session_state.processing_done = False | |
| st.success(f"✅ Video ready: **{uploaded_file.name}**") | |
| cap_info = cv.VideoCapture(st.session_state.tmp_path) | |
| if not cap_info.isOpened(): | |
| st.error("Error: Could not open video file.") | |
| else: | |
| fps_info = cap_info.get(cv.CAP_PROP_FPS) | |
| w_info = int(cap_info.get(cv.CAP_PROP_FRAME_WIDTH)) | |
| h_info = int(cap_info.get(cv.CAP_PROP_FRAME_HEIGHT)) | |
| total_info = int(cap_info.get(cv.CAP_PROP_FRAME_COUNT)) | |
| cap_info.release() | |
| c1, c2, c3, c4 = st.columns(4) | |
| c1.metric("Resolution", f"{w_info}x{h_info}") | |
| c2.metric("FPS", f"{fps_info:.1f}") | |
| c3.metric("Frames", total_info) | |
| c4.metric("Duration", f"{total_info/fps_info:.1f}s" if fps_info > 0 else "N/A") | |
| st.divider() | |
| # ══════════════════════════════════════════════════════════════════════════ | |
| # ÉTAPE 2 — TRAITEMENT | |
| # ══════════════════════════════════════════════════════════════════════════ | |
| st.subheader("▶️ Step 2 — Run Detection") | |
| start = st.button("▶️ Start Detection", type="primary", | |
| disabled=st.session_state.processing_done) | |
| if start: | |
| st.session_state.processing_done = False | |
| st.session_state.global_unique_ids = {} | |
| col_video, col_stats = st.columns([2, 1]) | |
| with col_video: | |
| st.subheader("📹 Annotated video") | |
| frame_display = st.empty() | |
| status_display = st.empty() | |
| with col_stats: | |
| st.subheader("🔢 Unique counters") | |
| counter_display = st.empty() | |
| progress_bar = st.progress(0) | |
| tmp_path = st.session_state.tmp_path | |
| video_name = uploaded_file.name | |
| logs = [] | |
| # ── YOLOv11 ─────────────────────────────────────────────────────────── | |
| if model_choice in ("YOLOv11 + ByteTrack (Base)", | |
| "YOLOv11 + ByteTrack (Fine-tuned)"): | |
| if not os.path.exists(selected_yolo_path): | |
| st.error(f"Model not found: `{selected_yolo_path}`") | |
| st.stop() | |
| model = YOLO(selected_yolo_path) | |
| cap = cv.VideoCapture(tmp_path) | |
| fps = cap.get(cv.CAP_PROP_FPS) or 25 | |
| total = int(cap.get(cv.CAP_PROP_FRAME_COUNT)) | |
| w_vid = int(cap.get(cv.CAP_PROP_FRAME_WIDTH)) | |
| h_vid = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT)) | |
| line_y = h_vid // 2 | |
| track_hits = {} | |
| prev_pos = {} | |
| frame_idx = 0 | |
| log_path = f"logs/log_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv" | |
| name_to_id = {v: k for k, v in model.names.items()} | |
| selected_ids = [name_to_id[c] for c in selected_classes if c in name_to_id] | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| if frame_idx % frame_skip != 0: | |
| frame_idx += 1 | |
| continue | |
| timestamp = round(frame_idx / fps, 3) | |
| results = model.track( | |
| frame, persist=True, tracker="bytetrack.yaml", | |
| conf=confidence_threshold, | |
| classes=selected_ids if selected_ids else None, | |
| device=DEVICE, verbose=False | |
| ) | |
| no_object = True | |
| accepted = [] | |
| if results[0].boxes is not None and results[0].boxes.id is not None: | |
| for box in results[0].boxes: | |
| track_id = int(box.id[0]) | |
| cls_name = normalize_class_name(model.names[int(box.cls[0])]) | |
| conf = round(float(box.conf[0]), 3) | |
| x1, y1, x2, y2 = map(int, box.xyxy[0]) | |
| if shape_refine: | |
| cls_name = refine_class_by_shape( | |
| cls_name, x1, y1, x2, y2, person_car_ratio) | |
| area = max(0, x2-x1) * max(0, y2-y1) | |
| if cls_name not in selected_classes: continue | |
| if conf < confidence_threshold: continue | |
| if area < min_box_area: continue | |
| no_object = False | |
| track_hits[track_id] = track_hits.get(track_id, 0) + 1 | |
| if track_hits[track_id] >= min_track_hits: | |
| st.session_state.global_unique_ids.setdefault( | |
| track_id, cls_name) | |
| # ── Schéma partagé ───────────────────────────────── | |
| cx = (x1 + x2) // 2 | |
| cy = (y1 + y2) // 2 | |
| if track_id in prev_pos: | |
| px_, py_ = prev_pos[track_id] | |
| dist = ((cx-px_)**2 + (cy-py_)**2) ** 0.5 | |
| speed = round(dist * fps, 2) | |
| if abs(cx-px_) > abs(cy-py_): | |
| direction = "right" if cx > px_ else "left" | |
| else: | |
| direction = "down" if cy > py_ else "up" | |
| else: | |
| speed, direction = 0.0, "" | |
| prev_pos[track_id] = (cx, cy) | |
| crossed = "true" if abs(cy - line_y) < 10 else "false" | |
| accepted.append((track_id, cls_name, conf, x1, y1, x2, y2)) | |
| logs.append([ | |
| frame_idx, timestamp, | |
| scene_name, group_id, video_name, | |
| track_id, cls_name, conf, | |
| x1, y1, x2, y2, | |
| cx, cy, w_vid, h_vid, | |
| crossed, direction, speed | |
| ]) | |
| annotated = frame.copy() | |
| # Ligne de comptage | |
| cv.line(annotated, (0, line_y), (w_vid, line_y), (0, 255, 0), 2) | |
| cv.putText(annotated, "Counting line", (10, line_y - 10), | |
| cv.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) | |
| for track_id, cls_name, conf, x1, y1, x2, y2 in accepted: | |
| color = (0, 0, 255) if cls_name == "stop sign" else (0, 255, 0) | |
| thickness = 3 if cls_name == "stop sign" else 2 | |
| label = f"{cls_name} ID:{track_id} {conf:.2f}" | |
| cv.rectangle(annotated, (x1, y1), (x2, y2), color, thickness) | |
| cv.putText(annotated, label, (x1, max(y1-8, 0)), | |
| cv.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) | |
| if no_object: | |
| cv.putText(annotated, "No selected object detected", | |
| (20, 50), cv.FONT_HERSHEY_SIMPLEX, | |
| 1.0, (0, 0, 255), 2) | |
| status_display.warning("⚠️ No selected object in this frame") | |
| else: | |
| status_display.empty() | |
| frame_display.image( | |
| cv.cvtColor(annotated, cv.COLOR_BGR2RGB), | |
| channels="RGB", use_column_width=True | |
| ) | |
| stats = {} | |
| for cls in st.session_state.global_unique_ids.values(): | |
| stats[cls] = stats.get(cls, 0) + 1 | |
| counter_display.markdown( | |
| "\n".join([f"**{c}** : {n} unique" | |
| for c, n in sorted(stats.items())]) | |
| or "_Waiting for objects..._" | |
| ) | |
| if total > 0: | |
| progress_bar.progress(min(frame_idx / total, 1.0)) | |
| frame_idx += 1 | |
| cap.release() | |
| progress_bar.progress(1.0) | |
| with open(log_path, 'w', newline='') as f: | |
| writer = csv.writer(f) | |
| writer.writerow([ | |
| 'frame', 'timestamp_sec', 'scene_name', 'group_id', 'video_name', | |
| 'track_id', 'class_name', 'confidence', | |
| 'bbox_x1', 'bbox_y1', 'bbox_x2', 'bbox_y2', | |
| 'cx', 'cy', 'frame_width', 'frame_height', | |
| 'crossed_line', 'direction', 'speed_px_s' | |
| ]) | |
| writer.writerows(logs) | |
| st.success(f"✅ Done — {frame_idx} frames analyzed") | |
| # ── SSD ─────────────────────────────────────────────────────────────── | |
| else: | |
| import torchvision | |
| from torchvision.transforms import functional as TF | |
| import numpy as np | |
| COCO_CLASSES = [ | |
| "__background__","person","bicycle","car","motorcycle","airplane", | |
| "bus","train","truck","boat","traffic light","fire hydrant","N/A", | |
| "stop sign","parking meter","bench","bird","cat","dog","horse", | |
| "sheep","cow","elephant","bear","zebra","giraffe","N/A","backpack", | |
| "umbrella","N/A","N/A","handbag","tie","suitcase","frisbee","skis", | |
| "snowboard","sports ball","kite","baseball bat","baseball glove", | |
| "skateboard","surfboard","tennis racket","bottle","N/A","wine glass", | |
| "cup","fork","knife","spoon","bowl","banana","apple","sandwich", | |
| "orange","broccoli","carrot","hot dog","pizza","donut","cake", | |
| "chair","couch","potted plant","bed","N/A","dining table","N/A", | |
| "N/A","toilet","N/A","tv","laptop","mouse","remote","keyboard", | |
| "cell phone","microwave","oven","toaster","sink","refrigerator", | |
| "N/A","book","clock","vase","scissors","teddy bear","hair drier", | |
| "toothbrush" | |
| ] | |
| device = torch.device(DEVICE) | |
| weights = torchvision.models.detection\ | |
| .SSDLite320_MobileNet_V3_Large_Weights.COCO_V1 | |
| ssd_model = torchvision.models.detection\ | |
| .ssdlite320_mobilenet_v3_large(weights=weights) | |
| ssd_model.to(device) | |
| ssd_model.eval() | |
| cap = cv.VideoCapture(tmp_path) | |
| fps = cap.get(cv.CAP_PROP_FPS) or 25 | |
| total = int(cap.get(cv.CAP_PROP_FRAME_COUNT)) | |
| w_vid = int(cap.get(cv.CAP_PROP_FRAME_WIDTH)) | |
| h_vid = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT)) | |
| line_y = h_vid // 2 | |
| frame_idx = 0 | |
| log_path = f"logs/ssd_log_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv" | |
| np.random.seed(42) | |
| COLORS = np.random.randint(0, 255, size=(len(COCO_CLASSES), 3), | |
| dtype="uint8") | |
| seen_classes = set() | |
| prev_pos = {} | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| if frame_idx % frame_skip != 0: | |
| frame_idx += 1 | |
| continue | |
| timestamp = round(frame_idx / fps, 3) | |
| img_rgb = cv.cvtColor(frame, cv.COLOR_BGR2RGB) | |
| tensor = TF.to_tensor(img_rgb).unsqueeze(0).to(device) | |
| with torch.no_grad(): | |
| outputs = ssd_model(tensor)[0] | |
| boxes_arr = outputs["boxes"].cpu().numpy() | |
| labels_arr = outputs["labels"].cpu().numpy() | |
| scores_arr = outputs["scores"].cpu().numpy() | |
| no_object = True | |
| for box, label, score in zip(boxes_arr, labels_arr, scores_arr): | |
| if score < confidence_threshold: | |
| continue | |
| cls_name = COCO_CLASSES[label] if label < len(COCO_CLASSES) \ | |
| else "unknown" | |
| cls_name = normalize_class_name(cls_name) | |
| x1, y1, x2, y2 = map(int, box) | |
| if shape_refine: | |
| cls_name = refine_class_by_shape( | |
| cls_name, x1, y1, x2, y2, person_car_ratio) | |
| area = max(0, x2-x1) * max(0, y2-y1) | |
| if cls_name not in selected_classes: continue | |
| if area < min_box_area: continue | |
| no_object = False | |
| seen_classes.add(cls_name) | |
| # ── Schéma partagé ───────────────────────────────────── | |
| cx = (x1 + x2) // 2 | |
| cy = (y1 + y2) // 2 | |
| key = int(label) | |
| if key in prev_pos: | |
| px_, py_ = prev_pos[key] | |
| dist = ((cx-px_)**2 + (cy-py_)**2) ** 0.5 | |
| speed = round(dist * fps, 2) | |
| if abs(cx-px_) > abs(cy-py_): | |
| direction = "right" if cx > px_ else "left" | |
| else: | |
| direction = "down" if cy > py_ else "up" | |
| else: | |
| speed, direction = 0.0, "" | |
| prev_pos[key] = (cx, cy) | |
| crossed = "true" if abs(cy - line_y) < 10 else "false" | |
| color = (0,0,255) if cls_name == "stop sign" \ | |
| else [int(c) for c in COLORS[label]] | |
| cv.rectangle(frame, (x1,y1),(x2,y2), color, | |
| 3 if cls_name=="stop sign" else 2) | |
| label_text = "STOP SIGN" if cls_name=="stop sign" \ | |
| else f"{cls_name}: {score:.2f}" | |
| cv.putText(frame, label_text, (x1, max(y1-8,0)), | |
| cv.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) | |
| logs.append([ | |
| frame_idx, timestamp, | |
| scene_name, group_id, video_name, | |
| -1, cls_name, round(float(score), 3), | |
| x1, y1, x2, y2, | |
| cx, cy, w_vid, h_vid, | |
| crossed, direction, speed | |
| ]) | |
| # Ligne de comptage | |
| cv.line(frame, (0, line_y), (w_vid, line_y), (0, 255, 0), 2) | |
| cv.putText(frame, "Counting line", (10, line_y - 10), | |
| cv.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) | |
| if no_object: | |
| cv.putText(frame, "No selected object detected", | |
| (20,50), cv.FONT_HERSHEY_SIMPLEX, | |
| 1.0, (0,0,255), 2) | |
| status_display.warning("⚠️ No selected object in this frame") | |
| else: | |
| status_display.empty() | |
| frame_display.image( | |
| cv.cvtColor(frame, cv.COLOR_BGR2RGB), | |
| channels="RGB", use_column_width=True | |
| ) | |
| counter_display.markdown( | |
| "\n".join([f"**{c}** : 1 unique" for c in sorted(seen_classes)]) | |
| or "_Waiting..._" | |
| ) | |
| if total > 0: | |
| progress_bar.progress(min(frame_idx / total, 1.0)) | |
| frame_idx += 1 | |
| cap.release() | |
| progress_bar.progress(1.0) | |
| with open(log_path, 'w', newline='') as f: | |
| writer = csv.writer(f) | |
| writer.writerow([ | |
| 'frame', 'timestamp_sec', 'scene_name', 'group_id', 'video_name', | |
| 'track_id', 'class_name', 'confidence', | |
| 'bbox_x1', 'bbox_y1', 'bbox_x2', 'bbox_y2', | |
| 'cx', 'cy', 'frame_width', 'frame_height', | |
| 'crossed_line', 'direction', 'speed_px_s' | |
| ]) | |
| writer.writerows(logs) | |
| st.success(f"✅ Done — {frame_idx} frames analyzed") | |
| st.session_state.processing_done = True | |
| # ── Final Stats ──────────────────────────────────────────────────────── | |
| if logs: | |
| cols = [ | |
| 'frame', 'timestamp_sec', 'scene_name', 'group_id', 'video_name', | |
| 'track_id', 'class_name', 'confidence', | |
| 'bbox_x1', 'bbox_y1', 'bbox_x2', 'bbox_y2', | |
| 'cx', 'cy', 'frame_width', 'frame_height', | |
| 'crossed_line', 'direction', 'speed_px_s' | |
| ] | |
| df = pd.DataFrame(logs, columns=cols) | |
| st.subheader("📊 Final Statistics") | |
| ca, cb = st.columns(2) | |
| with ca: | |
| counts = df['class_name'].value_counts().reset_index() | |
| counts.columns = ['Class', 'Count'] | |
| st.plotly_chart( | |
| px.bar(counts, x='Class', y='Count', | |
| title="Detections per class", color='Class'), | |
| use_container_width=True | |
| ) | |
| with cb: | |
| df['time_bin'] = (df['timestamp_sec'] // 5) * 5 | |
| intensity = df.groupby('time_bin')['class_name'].count().reset_index() | |
| intensity.columns = ['Time (s)', 'Detections'] | |
| st.plotly_chart( | |
| px.line(intensity, x='Time (s)', y='Detections', | |
| title="Detection intensity over time"), | |
| use_container_width=True | |
| ) | |
| st.download_button( | |
| "📥 Download CSV logs", | |
| data=open(log_path).read(), | |
| file_name=os.path.basename(log_path), | |
| mime="text/csv" | |
| ) | |
| if st.button("🔄 Process another video"): | |
| st.session_state.video_ready = False | |
| st.session_state.tmp_path = None | |
| st.session_state.processing_done = False | |
| st.session_state.global_unique_ids = {} | |
| st.rerun() | |
| else: | |
| st.session_state.video_ready = False | |
| st.session_state.tmp_path = None | |
| st.session_state.processing_done = False | |
| st.info("👆 Upload a video to start.") |