Update app.py

app.py

@@ -5,18 +5,19 @@ Gradio app to compare object‑detection models:
   • Roboflow RF‑DETR (Base, Large)
   • Custom fine‑tuned checkpoints for either framework (upload .pt/.pth files)
 
-Python ≥3.9
-Install:
-  pip install -r requirements.txt
-Optionally, add GPU‑specific PyTorch wheels or `rfdetr[onnxexport]` for ONNX export.
+Changes in this revision (2025‑04‑19):
+  • Thinner, semi‑transparent bounding boxes for better visibility in crowded scenes.
+  • Legend now shows a clean dict of runtimes (or concise errors) instead of auto‑indexed JSON.
+  • File uploader is fully integrated for custom checkpoints.
 """
 
 from __future__ import annotations
 
 import time
 from pathlib import Path
-from typing import List, Tuple, Optional
+from typing import List, Tuple, Dict, Optional
 
+import cv2
 import numpy as np
 from PIL import Image
 import gradio as gr
@@ -30,7 +31,7 @@ from rfdetr.util.coco_classes import COCO_CLASSES
 ###############################################################################
 
 YOLO_MODEL_MAP = {
-    # Names follow Ultralytics hub convention; they will be auto‑downloaded
+    # Ultralytics hub IDs — downloaded on first use
     "YOLOv12‑n": "yolov12n.pt",
     "YOLOv12‑s": "yolov12s.pt",
     "YOLOv12‑m": "yolov12m.pt",
@@ -53,21 +54,16 @@ ALL_MODELS = list(YOLO_MODEL_MAP.keys()) + list(RFDETR_MODEL_MAP.keys()) + [
     "Custom RF‑DETR (.pth)",
 ]
 
-_loaded = {}  # cache of already‑instantiated models
+_loaded: Dict[str, object] = {}
 
 def load_model(choice: str, custom_file: Optional[Path] = None):
-    """Return (and cache) a model instance for *choice*.
-    custom_file is a Path object (uploaded file) used when choice is custom.
-    Raises RuntimeError with helpful message if loading fails.
-    """
-    global _loaded
+    """Lazy‑load and cache a detector. Returns a model instance or raises RuntimeError."""
     if choice in _loaded:
         return _loaded[choice]
 
     try:
         if choice in YOLO_MODEL_MAP:
-            weight_id = YOLO_MODEL_MAP[choice]
-            mdl = YOLO(weight_id)  # Ultralytics downloads if not local
+            mdl = YOLO(YOLO_MODEL_MAP[choice])  # hub download if needed
         elif choice in RFDETR_MODEL_MAP:
             mdl = RFDETRBase() if RFDETR_MODEL_MAP[choice] == "base" else RFDETRLarge()
         elif choice.startswith("Custom YOLO"):
@@ -80,43 +76,47 @@ def load_model(choice: str, custom_file: Optional[Path] = None):
             mdl = RFDETRBase(pretrain_weights=str(custom_file))
         else:
             raise ValueError(f"Unsupported model choice: {choice}")
-    except FileNotFoundError as e:
-        raise RuntimeError(
-            f"Weights for '{choice}' not found locally and could not be downloaded. "
-            "Place the .pt file in the working directory, supply a custom checkpoint, "
-            "or ensure the model is released on the Ultralytics hub.\n" + str(e)
-        ) from e
+    except Exception as e:
+        raise RuntimeError(str(e)) from e
 
     _loaded[choice] = mdl
     return mdl
 
 ###############################################################################
-# Inference helpers
+# Inference helpers — semi‑transparent, thin boxes
 ###############################################################################
 
-box_annotator = sv.BoxAnnotator()
+box_annotator = sv.BoxAnnotator(thickness=2)  # thinner lines
 label_annotator = sv.LabelAnnotator()
 
+def blend_overlay(base_np: np.ndarray, overlay_np: np.ndarray, alpha: float = 0.6) -> np.ndarray:
+    """Blend two BGR images with given alpha for overlay."""
+    return cv2.addWeighted(overlay_np, alpha, base_np, 1 - alpha, 0)
+
 def run_single_inference(model, image: Image.Image, threshold: float) -> Tuple[Image.Image, float]:
     start = time.perf_counter()
 
     if isinstance(model, (RFDETRBase, RFDETRLarge)):
         detections = model.predict(image, threshold=threshold)
         label_source = COCO_CLASSES
-    else:  # Ultralytics YOLO
+    else:
         result = model.predict(image, verbose=False)[0]
         detections = sv.Detections.from_ultralytics(result)
         label_source = model.names
 
     runtime = time.perf_counter() - start
 
-    labels = [f"{label_source[cid]} {conf:.2f}" for cid, conf in zip(detections.class_id, detections.confidence)]
-    annotated = box_annotator.annotate(image.copy(), detections)
-    annotated = label_annotator.annotate(annotated, detections, labels)
-    return annotated, runtime
+    img_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    overlay = img_np.copy()
+    overlay = box_annotator.annotate(overlay, detections)
+    overlay = label_annotator.annotate(overlay, detections, [f"{label_source[c]} {p:.2f}" for c, p in zip(detections.class_id, detections.confidence)])
+    blended = blend_overlay(img_np, overlay, alpha=0.6)  # semi‑transparent boxes
+    annotated_pil = Image.fromarray(cv2.cvtColor(blended, cv2.COLOR_BGR2RGB))
+
+    return annotated_pil, runtime
 
 ###############################################################################
-# Gradio UI logic
+# Gradio callback
 ###############################################################################
 
 def compare_models(models: List[str], img: Image.Image, threshold: float, custom_file: Optional[Path]):
@@ -125,45 +125,44 @@ def compare_models(models: List[str], img: Image.Image, threshold: float, custom
     if img.mode != "RGB":
         img = img.convert("RGB")
 
-    results, legends = [], []
+    results: List[Image.Image] = []
+    legends: Dict[str, str] = {}
+
     for m in models:
         try:
             model_obj = load_model(m, custom_file)
             annotated, t = run_single_inference(model_obj, img, threshold)
             results.append(annotated)
-            legends.append(f"{m} – {t*1000:.1f} ms")
+            legends[m] = f"{t*1000:.1f} ms"
         except Exception as e:
-            # Append a blank image with the error message overlayed
-            error_img = Image.new("RGB", img.size, color=(30, 30, 30))
-            legends.append(f"{m} – ERROR: {e}")
-            results.append(error_img)
+            # show blank slate if model unavailable
+            results.append(Image.new("RGB", img.size, (40, 40, 40)))
+            err = str(e).split("\n")[0][:120]  # shorten
+            legends[m] = f"ERROR: {err}"
+
     return results, legends
 
 ###############################################################################
-# Build & launch demo
+# Build & launch Gradio UI
 ###############################################################################
 
 def build_demo():
     with gr.Blocks(title="CV Model Comparison") as demo:
-        gr.Markdown("""# 🔍 Compare Object‑Detection Models\nUpload an image, select detectors, and optionally upload a custom checkpoint.\nThe app annotates predictions and reports per‑model latency.""")
+        gr.Markdown("""# 🔍 Compare Object‑Detection Models\nUpload an image, choose detectors, and optionally add a custom checkpoint.\nBounding boxes are thin and 60 % opaque for clarity.""")
 
         with gr.Row():
             model_select = gr.CheckboxGroup(choices=ALL_MODELS, value=["YOLOv12‑n"], label="Select models")
-            threshold_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, step=0.05, label="Confidence threshold")
+            threshold_slider = gr.Slider(0.0, 1.0, 0.5, step=0.05, label="Confidence threshold")
 
-        custom_checkpoint = gr.File(label="Upload custom YOLO / RF‑DETR checkpoint", file_types=[".pt", ".pth"], interactive=True)
-        image_in = gr.Image(type="pil", label="Upload image", sources=["upload", "webcam"], show_label=True)
+        custom_checkpoint = gr.File(label="Upload custom checkpoint (.pt/.pth)", file_types=[".pt", ".pth"], interactive=True)
+        image_in = gr.Image(type="pil", label="Image", sources=["upload", "webcam"])
 
         with gr.Row():
             gallery = gr.Gallery(label="Annotated results", columns=2, height="auto")
-        legends_out = gr.JSON(label="Runtime (ms) or error messages")
+        legends_out = gr.JSON(label="Latency / status by model")
 
         run_btn = gr.Button("Run Inference", variant="primary")
-        run_btn.click(
-            fn=compare_models,
-            inputs=[model_select, image_in, threshold_slider, custom_checkpoint],
-            outputs=[gallery, legends_out],
-        )
+        run_btn.click(compare_models, [model_select, image_in, threshold_slider, custom_checkpoint], [gallery, legends_out])
 
     return demo