Update app.py

app.py

@@ -1,47 +1,25 @@
 """
-=============================================================
-  HuggingFace Spaces — Gradio Demo
-  Mudflap / Black-sheet Detection
-  Models: YOLOv11x  +  YOLOv5x
-=============================================================
-File layout in your HF Space repo:
-  app.py                ← this file
-  requirements.txt
-  models/
-      yolov11x_mudflap_best.pt
-      yolov5x_mudflap_best.pt
-  examples/
-      example1.jpg
-      example2.jpg
-=============================================================
+HuggingFace Spaces — Gradio Demo
+Mudflap / Black-sheet Detection  |  Models: YOLOv11x + YOLOv5x
 """
 
 import os
 import gradio as gr
 import torch
 import numpy as np
-import cv2
 from pathlib import Path
 from PIL import Image, ImageDraw, ImageFont
 import time
 
-# Allow Ultralytics to write to /tmp (needed on HF Spaces)
 os.environ.setdefault("YOLO_CONFIG_DIR", "/tmp/Ultralytics")
 
-# ─────────────────────────────────────────────────────────────
-# 1. MODEL LOADING
-# ─────────────────────────────────────────────────────────────
-DEVICE       = "cuda" if torch.cuda.is_available() else "cpu"
-MODEL_DIR    = Path("models")
-CLASS_NAMES  = ["mudflap"]
-
-# Colour palette per model
-COLORS = {
-    "YOLOv11x": (0,  200, 100),   # green
-    "YOLOv5x" : (255, 140,  0),   # orange
-}
+# ── Config ────────────────────────────────────────────────────
+DEVICE      = "cuda" if torch.cuda.is_available() else "cpu"
+MODEL_DIR   = Path("models")
+CLASS_NAMES = ["mudflap"]
+COLORS      = {"YOLOv11x": (0, 200, 100), "YOLOv5x": (255, 140, 0)}
 
-# ── YOLOv11x (Ultralytics) ────────────────────────────────────
+# ── Model loading ─────────────────────────────────────────────
 def load_v11():
     from ultralytics import YOLO
     m = YOLO(str(MODEL_DIR / "yolov11x_mudflap_best.pt"))
@@ -49,168 +27,99 @@ def load_v11():
     print("[✓] YOLOv11x loaded.")
     return m
 
-# ── YOLOv5x (torch.hub) ───────────────────────────────────────
 def load_v5():
-    # trust_repo=True  ← fixes the interactive-prompt EOFError on HF Spaces
     m = torch.hub.load(
         "ultralytics/yolov5", "custom",
         path=str(MODEL_DIR / "yolov5x_mudflap_best.pt"),
-        force_reload=False,
-        trust_repo=True,          # ← KEY FIX
-        device=DEVICE,
+        force_reload=False, trust_repo=True, device=DEVICE,
     )
-    m.conf    = 0.60
-    m.iou     = 0.45
-    m.max_det = 100
+    m.conf = 0.60; m.iou = 0.45; m.max_det = 100
     print("[✓] YOLOv5x loaded.")
     return m
 
-# Load once at startup
 model_v11 = load_v11()
 model_v5  = load_v5()
 
-
-# ─────────────────────────────────────────────────────────────
-# 2. INFERENCE HELPERS
-# ─────────────────────────────────────────────────────────────
+# ── Inference ─────────────────────────────────────────────────
 def run_v11(img_np, conf_thr, iou_thr):
-    """Returns list of (x1,y1,x2,y2,conf,cls_id)."""
-    results = model_v11.predict(
-        source=img_np, conf=conf_thr, iou=iou_thr,
-        imgsz=640, device=DEVICE, verbose=False
-    )
+    results = model_v11.predict(source=img_np, conf=conf_thr, iou=iou_thr,
+                                imgsz=640, device=DEVICE, verbose=False)
     boxes = []
     for r in results:
         for b in r.boxes:
             x1, y1, x2, y2 = b.xyxy[0].cpu().tolist()
-            conf = float(b.conf[0])
-            cls  = int(b.cls[0])
-            boxes.append((x1, y1, x2, y2, conf, cls))
+            boxes.append((x1, y1, x2, y2, float(b.conf[0]), int(b.cls[0])))
     return boxes
 
+def run_v5(img_np, conf_thr, iou_thr):
+    model_v5.conf = conf_thr; model_v5.iou = iou_thr
+    results = model_v5(img_np, size=640)
+    return [(x[0], x[1], x[2], x[3], x[4], int(x[5]))
+            for *x, _ in [r for r in results.xyxy[0].cpu().tolist()]]
 
 def run_v5(img_np, conf_thr, iou_thr):
-    """Returns list of (x1,y1,x2,y2,conf,cls_id)."""
     model_v5.conf = conf_thr
     model_v5.iou  = iou_thr
     results = model_v5(img_np, size=640)
-    boxes   = []
+    boxes = []
     for *xyxy, conf, cls in results.xyxy[0].cpu().tolist():
         boxes.append((xyxy[0], xyxy[1], xyxy[2], xyxy[3], conf, int(cls)))
     return boxes
 
-
-# ─────────────────────────────────────────────────────────────
-# 3. DRAWING UTILITY
-# ─────────────────────────────────────────────────────────────
-def draw_boxes(img_pil: Image.Image, detections: list,
-               color: tuple, label_prefix: str) -> Image.Image:
+# ── Drawing ───────────────────────────────────────────────────
+def draw_boxes(img_pil, detections, color, prefix):
     draw = ImageDraw.Draw(img_pil)
     try:
-        font = ImageFont.truetype(
-            "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 18)
+        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 18)
     except Exception:
         font = ImageFont.load_default()
-
     for (x1, y1, x2, y2, conf, cls) in detections:
         x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
-        lbl  = f"{label_prefix} {CLASS_NAMES[cls]} {conf:.2f}"
+        lbl = f"{prefix} {CLASS_NAMES[cls]} {conf:.2f}"
         r, g, b = color
-
-        # Bounding box (thick border)
         for t in range(3):
             draw.rectangle([x1-t, y1-t, x2+t, y2+t], outline=(r, g, b))
-
-        # Label background
-        tw, th = 0, 0
         try:
-            bbox_txt = draw.textbbox((x1, y1 - 22), lbl, font=font)
-            tw = bbox_txt[2] - bbox_txt[0]
-            th = bbox_txt[3] - bbox_txt[1]
+            bb = draw.textbbox((x1, y1-22), lbl, font=font)
+            tw, th = bb[2]-bb[0], bb[3]-bb[1]
         except Exception:
-            tw, th = len(lbl) * 8, 16
-
-        draw.rectangle([x1, y1 - th - 4, x1 + tw + 4, y1], fill=(r, g, b))
-        draw.text((x1 + 2, y1 - th - 2), lbl, fill=(255, 255, 255), font=font)
-
+            tw, th = len(lbl)*8, 16
+        draw.rectangle([x1, y1-th-4, x1+tw+4, y1], fill=(r, g, b))
+        draw.text((x1+2, y1-th-2), lbl, fill=(255, 255, 255), font=font)
     return img_pil
 
-
-def side_by_side(img_orig: Image.Image,
-                 dets_v11: list,
-                 dets_v5 : list) -> Image.Image:
-    """Create a 3-panel image: original | v11 result | v5 result."""
+def side_by_side(img_orig, dets_v11, dets_v5):
     w, h   = img_orig.size
-    canvas = Image.new("RGB", (w * 3 + 20, h + 40), (30, 30, 30))
-
-    draw = ImageDraw.Draw(canvas)
+    canvas = Image.new("RGB", (w*3+20, h+40), (18, 18, 18))
+    draw   = ImageDraw.Draw(canvas)
     try:
-        font_h = ImageFont.truetype(
-            "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 20)
+        fh = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 20)
     except Exception:
-        font_h = ImageFont.load_default()
-
-    titles = ["Original", "YOLOv11x", "YOLOv5x"]
-    for i, title in enumerate(titles):
-        draw.text((i * (w + 10) + w // 2 - 40, 5), title,
-                  fill=(255, 255, 255), font=font_h)
-
-    panel_orig = img_orig.copy()
-    panel_v11  = draw_boxes(img_orig.copy(), dets_v11, COLORS["YOLOv11x"], "v11")
-    panel_v5   = draw_boxes(img_orig.copy(), dets_v5,  COLORS["YOLOv5x"],  "v5")
-
-    canvas.paste(panel_orig, (0,          40))
-    canvas.paste(panel_v11,  (w + 10,     40))
-    canvas.paste(panel_v5,   (2*w + 20,   40))
+        fh = ImageFont.load_default()
+    for i, title in enumerate(["Original", "YOLOv11x", "YOLOv5x"]):
+        draw.text((i*(w+10)+w//2-40, 8), title, fill=(200, 200, 200), font=fh)
+    canvas.paste(img_orig.copy(),                                    (0,       40))
+    canvas.paste(draw_boxes(img_orig.copy(), dets_v11, COLORS["YOLOv11x"], "v11"), (w+10,   40))
+    canvas.paste(draw_boxes(img_orig.copy(), dets_v5,  COLORS["YOLOv5x"],  "v5"),  (2*w+20, 40))
     return canvas
 
-
-# ─────────────────────────────────────────────────────────────
-# 4. MAIN INFERENCE FUNCTION  (called by Gradio)
-# ─────────────────────────────────────────────────────────────
-def detect(image: Image.Image,
-           model_choice: str,
-           conf_thr: float,
-           iou_thr: float):
+# ── Main detect fn ────────────────────────────────────────────
+def detect(image, model_choice, conf_thr, iou_thr):
     if image is None:
-        return None, "⚠️ Please upload an image."
-
+        return None, "⚠️ Upload an image first."
     img_np = np.array(image.convert("RGB"))
-    t0     = time.time()
-
-    dets_v11, dets_v5 = [], []
-    stats_lines = []
-
-    if model_choice in ("YOLOv11x", "Both (side-by-side)"):
-        t1       = time.time()
-        dets_v11 = run_v11(img_np, conf_thr, iou_thr)
-        dt_v11   = time.time() - t1
-        stats_lines.append(
-            f"**YOLOv11x** — {len(dets_v11)} detection(s) in {dt_v11*1000:.1f} ms"
-        )
-        for d in dets_v11:
-            stats_lines.append(
-                f"  • {CLASS_NAMES[d[5]]}  conf={d[4]:.3f}  "
-                f"box=[{int(d[0])},{int(d[1])},{int(d[2])},{int(d[3])}]"
-            )
+    t0 = time.time()
+    dets_v11, dets_v5, lines = [], [], []
 
-    if model_choice in ("YOLOv5x", "Both (side-by-side)"):
-        t1      = time.time()
-        dets_v5 = run_v5(img_np, conf_thr, iou_thr)
-        dt_v5   = time.time() - t1
-        stats_lines.append(
-            f"**YOLOv5x** — {len(dets_v5)} detection(s) in {dt_v5*1000:.1f} ms"
-        )
-        for d in dets_v5:
-            stats_lines.append(
-                f"  • {CLASS_NAMES[d[5]]}  conf={d[4]:.3f}  "
-                f"box=[{int(d[0])},{int(d[1])},{int(d[2])},{int(d[3])}]"
-            )
+    if model_choice in ("YOLOv11x", "Both"):
+        t1 = time.time(); dets_v11 = run_v11(img_np, conf_thr, iou_thr)
+        lines.append(f"**YOLOv11x** — {len(dets_v11)} detection(s) · {(time.time()-t1)*1000:.0f} ms")
 
-    total_ms = (time.time() - t0) * 1000
-    stats_lines.append(
-        f"\n⏱ Total wall-time: {total_ms:.1f} ms | Device: {DEVICE.upper()}"
-    )
+    if model_choice in ("YOLOv5x", "Both"):
+        t1 = time.time(); dets_v5 = run_v5(img_np, conf_thr, iou_thr)
+        lines.append(f"**YOLOv5x** — {len(dets_v5)} detection(s) · {(time.time()-t1)*1000:.0f} ms")
+
+    lines.append(f"⏱ {(time.time()-t0)*1000:.0f} ms total · {DEVICE.upper()}")
 
     if model_choice == "YOLOv11x":
         out = draw_boxes(image.copy(), dets_v11, COLORS["YOLOv11x"], "v11")
@@ -219,52 +128,120 @@ def detect(image: Image.Image,
     else:
         out = side_by_side(image, dets_v11, dets_v5)
 
-    return out, "\n".join(stats_lines)
+    return out, "\n\n".join(lines)
+
+# ── Custom CSS ────────────────────────────────────────────────
+CSS = """
+@import url('https://fonts.googleapis.com/css2?family=Bebas+Neue&family=DM+Sans:wght@400;500&display=swap');
 
+body, .gradio-container {
+    background: #0f0f0f !important;
+    color: #e8e8e8 !important;
+    font-family: 'DM Sans', sans-serif !important;
+}
+
+#main-title {
+    font-family: 'Bebas Neue', sans-serif !important;
+    font-size: clamp(3rem, 8vw, 6rem) !important;
+    letter-spacing: 0.06em;
+    color: #ffffff;
+    text-align: center;
+    padding: 2rem 0 0.5rem 0;
+    line-height: 1;
+    border-bottom: 2px solid #f0a500;
+    margin-bottom: 2rem;
+}
+
+/* Panels */
+.gr-box, .gr-panel, .svelte-1gfkn6j, .block {
+    background: #1a1a1a !important;
+    border: 1px solid #2a2a2a !important;
+    border-radius: 6px !important;
+}
+
+/* Radio buttons */
+.gr-radio-row label {
+    background: #222 !important;
+    border: 1px solid #333 !important;
+    border-radius: 4px !important;
+    padding: 6px 14px !important;
+    cursor: pointer;
+    transition: all 0.15s;
+}
+.gr-radio-row label:hover { border-color: #f0a500 !important; }
+.gr-radio-row input:checked + label,
+.gr-radio-row label[data-checked] {
+    border-color: #f0a500 !important;
+    color: #f0a500 !important;
+}
 
-# ─────────────────────────────────────────────────────────────
-# 5. GRADIO UI
-# ─────────────────────────────────────────────────────────────
-DESCRIPTION = """
-## 🚛 Trailer Mudflap / Black-sheet Detector
-Detects the **black mud-flap panels** on the sides of Amazon Prime trailers.
+/* Sliders */
+input[type=range] { accent-color: #f0a500; }
+
+/* Detect button */
+#detect-btn {
+    background: #f0a500 !important;
+    color: #0f0f0f !important;
+    font-family: 'Bebas Neue', sans-serif !important;
+    font-size: 1.3rem !important;
+    letter-spacing: 0.1em;
+    border: none !important;
+    border-radius: 4px !important;
+    height: 52px !important;
+    transition: opacity 0.2s;
+}
+#detect-btn:hover { opacity: 0.85; }
 
-### Models
-| Model | Architecture | Notes |
-|-------|-------------|-------|
-| **YOLOv11x** | Ultralytics YOLO v11 X | Latest generation |
-| **YOLOv5x**  | Ultralytics YOLOv5 X   | Battle-tested baseline |
+/* Labels */
+label span, .gr-label { color: #999 !important; font-size: 0.78rem !important; text-transform: uppercase; letter-spacing: 0.08em; }
 
-Both models were fine-tuned with **heavy augmentation** (mosaic, mixup, perspective, weather,
-blur, random shadows, coarse-dropout …) on 300 augmented images from 25 originals.
+/* Example thumbnails */
+.gr-samples-table img { border-radius: 4px; border: 1px solid #333; }
 
-Select *Both* to compare the two models side-by-side.
+/* Stats output */
+#stats-out { font-size: 0.85rem !important; color: #aaa !important; }
 """
 
-with gr.Blocks(title="Mudflap Detector", theme=gr.themes.Soft()) as demo:
-    gr.Markdown(DESCRIPTION)
+# ── UI ────────────────────────────────────────────────────────
+with gr.Blocks(title="Mudflap Detector", css=CSS) as demo:
 
-    with gr.Row():
-        with gr.Column(scale=1):
-            inp_image   = gr.Image(type="pil", label="Input Image")
-            model_sel   = gr.Radio(
-                choices=["YOLOv11x", "YOLOv5x", "Both (side-by-side)"],
-                value="Both (side-by-side)",
-                label="Model",
-            )
-            conf_slider = gr.Slider(
-                0.10, 0.95, value=0.60, step=0.05,
-                label="Confidence Threshold"
+    gr.HTML('<div id="main-title">Mudflap Detector</div>')
+
+    with gr.Row(equal_height=False):
+
+        # ── Left column: inputs ───────────────────────────────
+        with gr.Column(scale=1, min_width=300):
+            inp_image = gr.Image(type="pil", label="Input Image", height=300)
+
+            gr.HTML('<div style="font-size:0.72rem;color:#666;text-transform:uppercase;letter-spacing:.08em;margin:10px 0 6px">Example Images</div>')
+            gr.Examples(
+                examples=[
+                    ["examples/example1.jpg", "Both", 0.60, 0.45],
+                    ["examples/example2.jpg", "YOLOv11x", 0.60, 0.45],
+                ],
+                inputs=[inp_image,
+                        gr.State("Both"),        # placeholder — wired below
+                        gr.State(0.60),
+                        gr.State(0.45)],
+                label=None,
             )
-            iou_slider  = gr.Slider(
-                0.10, 0.90, value=0.45, step=0.05,
-                label="IoU (NMS) Threshold"
+
+            model_sel = gr.Radio(
+                choices=["YOLOv11x", "YOLOv5x", "Both"],
+                value="Both",
+                label="Model",
             )
-            run_btn = gr.Button("🔍 Detect", variant="primary")
 
-        with gr.Column(scale=2):
-            out_image = gr.Image(type="pil", label="Detection Result")
-            out_stats = gr.Markdown(label="Stats")
+            with gr.Row():
+                conf_slider = gr.Slider(0.10, 0.95, value=0.60, step=0.05, label="Confidence")
+                iou_slider  = gr.Slider(0.10, 0.90, value=0.45, step=0.05, label="IoU")
+
+            run_btn = gr.Button("DETECT", elem_id="detect-btn", variant="primary")
+
+        # ── Right column: output ──────────────────────────────
+        with gr.Column(scale=2, min_width=400):
+            out_image = gr.Image(type="pil", label="Result", height=500)
+            out_stats = gr.Markdown(elem_id="stats-out")
 
     run_btn.click(
         fn=detect,
@@ -272,24 +249,5 @@ with gr.Blocks(title="Mudflap Detector", theme=gr.themes.Soft()) as demo:
         outputs=[out_image, out_stats],
     )
 
-    gr.Examples(
-        examples=[
-            ["examples/example1.jpg", "Both (side-by-side)", 0.60, 0.45],
-            ["examples/example2.jpg", "YOLOv11x",           0.60, 0.45],
-        ],
-        inputs=[inp_image, model_sel, conf_slider, iou_slider],
-        outputs=[out_image, out_stats],
-        fn=detect,
-        cache_examples=True,
-        label="Example Images",
-    )
-
-    gr.Markdown("""
----
-**Legend:**  
-🟩 Green boxes → YOLOv11x  |  🟠 Orange boxes → YOLOv5x  
-**Class:** `mudflap` (the black panels on trailer sides)
-""")
-
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)