Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import cv2 | |
| import numpy as np | |
| import tensorflow as tf | |
| from ultralytics import YOLO | |
| import tempfile | |
| import os | |
| import time | |
| # ────────────────────────────────────────── | |
| # LOAD MODELS | |
| # ────────────────────────────────────────── | |
| print("Loading YOLO...") | |
| yolo_model = YOLO("models/yolo.pt") | |
| print("Loading EfficientNet-B0...") | |
| # Loads your newly upgraded model architecture | |
| resnet_model = tf.keras.models.load_model("models/efficientnetb0.keras", compile=False) | |
| # Explicitly hardcoding your optimal benchmark evaluation threshold for EfficientNet-B0 | |
| RESNET_THRESHOLD = 0.35 | |
| CONF_THRESHOLD = 0.20 | |
| HEAD_CROP_RATIO = 0.80 | |
| # Warm up the model | |
| dummy = np.zeros((1, 224, 224, 3), dtype="float32") | |
| resnet_model.predict(dummy, verbose=0) | |
| print(f"[✓] YOLO classes : {yolo_model.names}") | |
| print(f"[✓] Optimal EffNet Threshold : {RESNET_THRESHOLD:.2f}") | |
| VEHICLE_CLASSES = {"motorcycle", "bike", "car", "truck", "bus", "vehicle", "auto", "scooter"} | |
| # ────────────────────────────────────────── | |
| # HELPERS | |
| # ────────────────────────────────────────── | |
| def preprocess_crop(crop_bgr): | |
| # Matches the exact preprocessing dimensions | |
| crop = cv2.resize(crop_bgr, (224, 224)) | |
| crop = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB) | |
| crop = crop.astype("float32") | |
| return np.expand_dims(crop, axis=0) | |
| def batch_predict(crops): | |
| if not crops: | |
| return [] | |
| batch = np.concatenate(crops, axis=0) | |
| preds = resnet_model.predict(batch, verbose=0) | |
| # Check if binary (sigmoid) vs categorical (softmax) | |
| if preds.shape[-1] == 1: | |
| return preds.flatten().tolist() | |
| else: | |
| # If model outputs 2 classes, extract probability belonging to 'no_helmet' (class 1) | |
| return preds[:, 1].tolist() | |
| def draw_box(frame, x1, y1, x2, y2, label, color): | |
| cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2) | |
| (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2) | |
| cv2.rectangle(frame, (x1, y1 - th - 10), (x1 + tw + 6, y1), color, -1) | |
| cv2.putText(frame, label, (x1 + 3, y1 - 4), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2) | |
| def draw_summary(frame, riders, helmets, no_helmets): | |
| pct = int(helmets / max(riders, 1) * 100) | |
| summary = f"Riders:{riders} Helmet:{helmets} NoHelmet:{no_helmets} {pct}%" | |
| bar_w = min(len(summary) * 13, frame.shape[1]) | |
| cv2.rectangle(frame, (0, 0), (bar_w, 36), (20, 20, 20), -1) | |
| cv2.putText(frame, summary, (8, 26), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) | |
| def classify_by_name(cls_name): | |
| n = cls_name.lower().strip() | |
| if any(x in n for x in ["without", "no helmet", "nohelmet", "no-helmet"]): | |
| return "no_helmet" | |
| if any(x in n for x in ["with helmet", "helmet"]): | |
| return "helmet" | |
| return "person" | |
| def process_frame(frame): | |
| results = yolo_model(frame, conf=CONF_THRESHOLD, verbose=False) | |
| annotated = frame.copy() | |
| boxes_info = [] | |
| crops_list = [] | |
| for r in results: | |
| for box in r.boxes: | |
| cls_id = int(box.cls[0]) | |
| cls_name = yolo_model.names[cls_id].lower().strip() | |
| if cls_name in VEHICLE_CLASSES: | |
| continue | |
| x1, y1, x2, y2 = map(int, box.xyxy[0]) | |
| x1 = max(0, x1); y1 = max(0, y1) | |
| x2 = min(frame.shape[1], x2) | |
| y2 = min(frame.shape[0], y2) | |
| if (x2 - x1) < 15 or (y2 - y1) < 15: | |
| continue | |
| mode = classify_by_name(cls_name) | |
| if mode in ("helmet", "no_helmet"): | |
| boxes_info.append((x1, y1, x2, y2, mode)) | |
| crops_list.append(None) | |
| else: | |
| head_y2 = min(y1 + int((y2 - y1) * HEAD_CROP_RATIO), frame.shape[0]) | |
| crop = frame[y1:head_y2, x1:x2] | |
| if crop.size == 0: | |
| continue | |
| boxes_info.append((x1, y1, x2, head_y2, "efficientnet")) | |
| crops_list.append(preprocess_crop(crop)) | |
| resnet_crops = [c for c in crops_list if c is not None] | |
| resnet_preds = batch_predict(resnet_crops) | |
| pred_iter = iter(resnet_preds) | |
| riders = helmets = no_helmets = 0 | |
| for (x1, y1, x2, y2, mode) in boxes_info: | |
| riders += 1 | |
| if mode == "helmet": | |
| label = "Helmet" | |
| conf = 0.94 | |
| elif mode == "no_helmet": | |
| label = "No Helmet" | |
| conf = 0.94 | |
| else: | |
| prob = next(pred_iter) | |
| label = "No Helmet" if prob >= RESNET_THRESHOLD else "Helmet" | |
| conf = prob if label == "No Helmet" else (1 - prob) | |
| color = (0, 200, 0) if label == "Helmet" else (0, 0, 220) | |
| if label == "Helmet": | |
| helmets += 1 | |
| else: | |
| no_helmets += 1 | |
| draw_box(annotated, x1, y1, x2, y2, f"{label} {conf:.2f}", color) | |
| draw_summary(annotated, riders, helmets, no_helmets) | |
| return annotated, riders, helmets, no_helmets | |
| # ────────────────────────────────────────── | |
| # INTERFACE IMPLEMENTATIONS | |
| # ────────────────────────────────────────── | |
| def predict_image(image_rgb): | |
| if image_rgb is None: | |
| return None, "No image uploaded." | |
| frame = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR) | |
| result, riders, helmets, no_helmets = process_frame(frame) | |
| result_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB) | |
| stats = ( | |
| f"Total Riders : {riders}\n" | |
| f"With Helmet : {helmets}\n" | |
| f"No Helmet : {no_helmets}\n" | |
| f"Compliance : {int(helmets / max(riders,1) * 100)}%" | |
| ) | |
| return result_rgb, stats | |
| def predict_video(video_path, skip_frames, max_minutes, progress=gr.Progress()): | |
| if video_path is None: | |
| return None, "No video uploaded." | |
| cap = cv2.VideoCapture(video_path) | |
| if not cap.isOpened(): | |
| return None, "Cannot open video." | |
| width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = cap.get(cv2.CAP_PROP_FPS) or 25.0 | |
| total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| scale = 1.0 | |
| if width > 640: | |
| scale = 640.0 / width | |
| width = 640 | |
| height = int(height * scale) | |
| max_frames = int(max_minutes * 60 * fps) | |
| total = min(total, max_frames) | |
| out_path = tempfile.mktemp(suffix=".mp4") | |
| writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height)) | |
| skip = max(1, int(skip_frames)) | |
| frame_idx = 0 | |
| proc_count = 0 | |
| last_annotated = None | |
| total_helmets = 0 | |
| total_no_helmets = 0 | |
| frames_with_riders = 0 | |
| start_time = time.time() | |
| progress(0, desc="Starting...") | |
| while frame_idx < total: | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| frame_idx += 1 | |
| if scale < 1.0: | |
| frame = cv2.resize(frame, (width, height)) | |
| if frame_idx % skip == 0: | |
| annotated, r, h, nh = process_frame(frame) | |
| last_annotated = annotated | |
| proc_count += 1 | |
| if r > 0: | |
| frames_with_riders += 1 | |
| total_helmets += h | |
| total_no_helmets += nh | |
| else: | |
| annotated = last_annotated if last_annotated is not None else frame | |
| writer.write(annotated) | |
| if frame_idx % 10 == 0: | |
| elapsed = time.time() - start_time | |
| spd = proc_count / max(elapsed, 1) | |
| progress(frame_idx / total, desc=f"Frame {frame_idx}/{total} | {spd:.1f} fps") | |
| cap.release() | |
| writer.release() | |
| elapsed = time.time() - start_time | |
| avg_h = total_helmets // max(frames_with_riders, 1) | |
| avg_nh = total_no_helmets // max(frames_with_riders, 1) | |
| avg_r = (avg_h + avg_nh) | |
| stats = ( | |
| f"Total frames : {frame_idx}\n" | |
| f"Frames processed : {proc_count}\n" | |
| f"Processing time : {elapsed:.1f}s\n" | |
| f"Speed : {proc_count/max(elapsed,1):.1f} fps\n" | |
| f"Avg Helmet/frame : {avg_h}\n" | |
| f"Avg No Helmet : {avg_nh}\n" | |
| f"Compliance rate : {int(avg_h/max(avg_r,1)*100)}%" | |
| ) | |
| return out_path, stats | |
| # ────────────────────────────────────────── | |
| # GRADIO BLOCK UI BUILD | |
| # ────────────────────────────────────────── | |
| with gr.Blocks(title="Smart Helmet Detection", theme=gr.themes.Soft(), css="footer{display:none!important} .mono{font-family:monospace}") as demo: | |
| gr.Markdown("# 🏍️ Smart Helmet Detection System") | |
| gr.Markdown("**YOLOv8 + EfficientNet-B0** — Upload an image or video") | |
| with gr.Tabs(): | |
| with gr.Tab("📷 Image Detection"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| img_input = gr.Image(label="Upload Image", type="numpy", sources=["upload","clipboard"]) | |
| img_button = gr.Button("🔍 Detect Helmets", variant="primary", size="lg") | |
| with gr.Column(): | |
| img_output = gr.Image(label="Detection Result") | |
| img_stats = gr.Textbox(label="Statistics", lines=5, elem_classes="mono") | |
| img_button.click(predict_image, inputs=img_input, outputs=[img_output, img_stats]) | |
| with gr.Tab("🎥 Video Detection"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| vid_input = gr.Video(label="Upload Video", sources=["upload"]) | |
| with gr.Row(): | |
| skip_slider = gr.Slider(minimum=1, maximum=15, value=5, step=1, label="Skip frames") | |
| max_min_slider = gr.Slider(minimum=0.5, maximum=3.0, value=0.5, step=0.5, label="Max minutes") | |
| vid_button = gr.Button("🎬 Process Video", variant="primary", size="lg") | |
| with gr.Column(): | |
| vid_output = gr.Video(label="Result Video") | |
| vid_stats = gr.Textbox(label="Statistics", lines=10, elem_classes="mono") | |
| vid_button.click(predict_video, inputs=[vid_input, skip_slider, max_min_slider], outputs=[vid_output, vid_stats]) | |
| demo.queue(max_size=5).launch() | |