"""PPE compliance detection API — Gradio app served on Hugging Face Spaces. Two endpoints: /detect image -> annotated image + text summary + structured report /detect_video video -> annotated MP4 + text summary + structured report """ import json import os import tempfile import cv2 import gradio as gr import imageio.v2 as imageio from PIL import Image from ultralytics import YOLO BASE_DIR = os.path.dirname(__file__) with open(os.path.join(BASE_DIR, "model_config.json"), encoding="utf-8") as f: _config = json.load(f) CLASS_NAMES = _config["class_names"] VIOLATION_CLASSES = set(_config["violation_classes"]) COMPLIANT_CLASSES = set(_config["compliant_classes"]) model = YOLO(os.path.join(BASE_DIR, _config["model"]["file"]), task="detect") EMPTY_REPORT = {"violations": [], "compliant": [], "workers": 0, "total_detections": 0} # Video limits. Inference on the free CPU tier costs ~2-3 s per frame, so the # model runs on a sampled subset of frames while every frame is still written # out with the most recent boxes drawn on it — the result plays back smoothly # at the source frame rate instead of looking like a slideshow. MAX_VIDEO_SECONDS = 15 SAMPLE_FPS = 2 MAX_OUTPUT_WIDTH = 1280 # --------------------------------------------------------------------------- image def detect_ppe(image, conf_threshold=0.25): """Detect PPE in an image and report compliance. Returns the annotated image, a human-readable summary, and a structured report for programmatic consumers (the web frontend). """ if image is None: return None, "No image provided.", dict(EMPTY_REPORT) # PIL input lets Ultralytics handle the RGB->BGR conversion itself; feeding # a raw RGB ndarray would be misread as BGR and skew color-sensitive # classes (hi-vis vests especially). result = model.predict(image, imgsz=640, conf=conf_threshold, verbose=False)[0] # result.plot() returns a BGR array (OpenCV convention) — flip to RGB. annotated = Image.fromarray(result.plot()[..., ::-1]) report = summarize(result) return annotated, build_summary(report), report def summarize(result): """Turn one Ultralytics result into our structured report.""" violations, compliant = [], [] workers = 0 for box in result.boxes: label = CLASS_NAMES[int(box.cls)] confidence = round(float(box.conf), 3) if label in VIOLATION_CLASSES: violations.append({"label": label, "confidence": confidence}) elif label in COMPLIANT_CLASSES: compliant.append({"label": label, "confidence": confidence}) elif label == "person": workers += 1 return { "violations": violations, "compliant": compliant, "workers": workers, "total_detections": len(result.boxes), } def build_summary(report): """Render the structured report as text for the Gradio demo page.""" lines = [] if report["violations"]: lines.append(f"⚠️ VIOLATIONS DETECTED ({len(report['violations'])}):") lines += [f" ❌ {v['label']} ({v['confidence']:.0%})" for v in report["violations"]] lines.append("") else: lines.append("✅ No PPE violations detected.") lines.append("") if report["compliant"]: lines.append(f"PPE compliant items ({len(report['compliant'])}):") lines += [f" ✅ {c['label']} ({c['confidence']:.0%})" for c in report["compliant"]] lines.append("") lines.append(f"Workers detected: {report['workers']}") lines.append(f"Total detections: {report['total_detections']}") return "\n".join(lines) # --------------------------------------------------------------------------- video def detect_ppe_video(video_path, conf_threshold=0.25): """Annotate a short video and report aggregate compliance. Returns the annotated MP4 path, a text summary, and a structured report. """ empty = { "frames_analyzed": 0, "frames_with_violations": 0, "violation_counts": {}, "compliant_counts": {}, "peak_workers": 0, "duration_seconds": 0.0, } if not video_path: return None, "No video provided.", empty capture = cv2.VideoCapture(video_path) if not capture.isOpened(): return None, "Could not read that video file.", empty fps = capture.get(cv2.CAP_PROP_FPS) or 25.0 frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT) or 0) duration = frame_count / fps if fps else 0.0 if duration > MAX_VIDEO_SECONDS + 0.5: capture.release() return (None, f"Video is {duration:.0f}s — the maximum is {MAX_VIDEO_SECONDS} seconds.", empty) stride = max(1, round(fps / SAMPLE_FPS)) out_path = os.path.join(tempfile.mkdtemp(), "ppe_annotated.mp4") # libx264 so the result plays in browsers; OpenCV's default mp4v does not. writer = imageio.get_writer(out_path, fps=fps, codec="libx264", quality=7, macro_block_size=None) analyzed = 0 frames_with_violations = 0 violation_counts, compliant_counts = {}, {} peak_workers = 0 last_result = None index = 0 try: while True: ok, frame = capture.read() if not ok: break if frame.shape[1] > MAX_OUTPUT_WIDTH: scale = MAX_OUTPUT_WIDTH / frame.shape[1] frame = cv2.resize(frame, (MAX_OUTPUT_WIDTH, int(frame.shape[0] * scale))) if index % stride == 0: # cv2 gives BGR, which is exactly what Ultralytics expects from # an ndarray — no conversion here (unlike the PIL image path). last_result = model.predict(frame, imgsz=640, conf=conf_threshold, verbose=False)[0] report = summarize(last_result) analyzed += 1 if report["violations"]: frames_with_violations += 1 # Count each label once per frame, not once per box: these are # reported as "seen in N of M frames", so a frame containing # three hard hats must not push the share above 100%. for label in {v["label"] for v in report["violations"]}: violation_counts[label] = violation_counts.get(label, 0) + 1 for label in {c["label"] for c in report["compliant"]}: compliant_counts[label] = compliant_counts.get(label, 0) + 1 peak_workers = max(peak_workers, report["workers"]) # Draw the most recent detections onto every frame so playback is # smooth rather than only updating twice a second. annotated = last_result.plot(img=frame) if last_result is not None else frame writer.append_data(annotated[..., ::-1]) # BGR -> RGB for the encoder index += 1 finally: capture.release() writer.close() report = { "frames_analyzed": analyzed, "frames_with_violations": frames_with_violations, "violation_counts": violation_counts, "compliant_counts": compliant_counts, "peak_workers": peak_workers, "duration_seconds": round(duration, 2), } return out_path, build_video_summary(report), report def build_video_summary(report): lines = [] total = report["frames_analyzed"] flagged = report["frames_with_violations"] if flagged: lines.append(f"⚠️ VIOLATIONS DETECTED in {flagged} of {total} analyzed frames") lines.append("") for label, count in sorted(report["violation_counts"].items(), key=lambda kv: -kv[1]): lines.append(f" ❌ {label} — seen in {count} frames") else: lines.append(f"✅ No PPE violations detected across {total} analyzed frames.") lines.append("") if report["compliant_counts"]: lines.append("PPE compliant items:") for label, count in sorted(report["compliant_counts"].items(), key=lambda kv: -kv[1]): lines.append(f" ✅ {label} — seen in {count} frames") lines.append("") lines.append(f"Peak workers in frame: {report['peak_workers']}") lines.append(f"Clip length: {report['duration_seconds']:.1f}s") return "\n".join(lines) # --------------------------------------------------------------------------- ui example_dir = os.path.join(BASE_DIR, "examples") examples = [] if os.path.isdir(example_dir): examples = [ [os.path.join(example_dir, f)] for f in sorted(os.listdir(example_dir)) if f.lower().endswith((".jpg", ".jpeg", ".png")) ] DESCRIPTION = ( "Detects hard hats and high-visibility vests on construction sites and flags " "missing equipment as violations.\n\n" "**Model:** YOLO11s (ONNX) · **Classes:** hardhat, no-hardhat, vest, no-vest, person" ) with gr.Blocks(title="PPE Compliance Detector") as demo: gr.Markdown("# 🦺 PPE Compliance Detector") gr.Markdown(DESCRIPTION) with gr.Tab("Image"): with gr.Row(): with gr.Column(): image_in = gr.Image(type="pil", label="Construction Site Image") image_conf = gr.Slider(0.1, 0.9, value=0.25, step=0.05, label="Confidence Threshold") image_btn = gr.Button("Detect PPE", variant="primary") with gr.Column(): image_out = gr.Image(label="Detection Result") image_text = gr.Textbox(label="Compliance Summary", lines=10) image_json = gr.JSON(label="Structured Report") if examples: gr.Examples(examples=examples, inputs=image_in) image_btn.click(detect_ppe, [image_in, image_conf], [image_out, image_text, image_json], api_name="detect") with gr.Tab("Video"): gr.Markdown( f"Upload a clip up to **{MAX_VIDEO_SECONDS} seconds**. Detection runs at " f"{SAMPLE_FPS} fps and the boxes are drawn onto every frame, so the result " "plays back smoothly. Processing takes roughly a minute on the free CPU tier." ) with gr.Row(): with gr.Column(): video_in = gr.Video(label="Construction Site Video") video_conf = gr.Slider(0.1, 0.9, value=0.25, step=0.05, label="Confidence Threshold") video_btn = gr.Button("Detect PPE in Video", variant="primary") with gr.Column(): video_out = gr.Video(label="Annotated Result") video_text = gr.Textbox(label="Compliance Summary", lines=10) video_json = gr.JSON(label="Structured Report") video_btn.click(detect_ppe_video, [video_in, video_conf], [video_out, video_text, video_json], api_name="detect_video") demo.queue(max_size=8).launch(ssr_mode=False)