Update app.py

app.py

@@ -5,10 +5,10 @@ Gradio app to compare object‑detection models:
   • Roboflow RF‑DETR (Base, Large)
   • Custom fine‑tuned checkpoints (.pt/.pth upload)
 
-Revision 2025‑04‑19‑d:
-  • Pre‑loads all selected models before running detections, with a visible progress bar.
-  • Progress shows two phases: *Loading weights* and *Running inference*.
-  • Keeps thin, semi‑transparent boxes and concise error labels.
+Revision 2025‑04‑19‑e:
+  • Gallery items now carry captions so you can see which model produced which image (and latency).
+  • Captions display as "Model (xx ms)" or error status.
+  • No other behaviour changed: pre‑loading, progress bar, thin semi‑transparent boxes, concise error labels.
 """
 
 from __future__ import annotations
@@ -31,7 +31,7 @@ from rfdetr.util.coco_classes import COCO_CLASSES
 ###############################################################################
 
 YOLO_MODEL_MAP: Dict[str, str] = {
-    # NOTE: Ultralytics filenames do NOT include the "v" character
+    # Ultralytics filenames omit the "v"
     "YOLOv12‑n": "yolo12n.pt",
     "YOLOv12‑s": "yolo12s.pt",
     "YOLOv12‑m": "yolo12m.pt",
@@ -44,7 +44,6 @@ YOLO_MODEL_MAP: Dict[str, str] = {
     "YOLOv11‑x": "yolo11x.pt",
 }
 
-
 RFDETR_MODEL_MAP = {
     "RF‑DETR‑Base (29M)": "base",
     "RF‑DETR‑Large (128M)": "large",
@@ -58,12 +57,11 @@ ALL_MODELS = list(YOLO_MODEL_MAP.keys()) + list(RFDETR_MODEL_MAP.keys()) + [
 _loaded: Dict[str, object] = {}
 
 def load_model(choice: str, custom_file: Optional[Path] = None):
-    """Fetch and cache a detector instance for *choice*."""
     if choice in _loaded:
         return _loaded[choice]
 
     if choice in YOLO_MODEL_MAP:
-        model = YOLO(YOLO_MODEL_MAP[choice])  # Ultralytics auto‑downloads if missing
+        model = YOLO(YOLO_MODEL_MAP[choice])
     elif choice in RFDETR_MODEL_MAP:
         model = RFDETRBase() if RFDETR_MODEL_MAP[choice] == "base" else RFDETRLarge()
     elif choice.startswith("Custom YOLO"):
@@ -84,8 +82,8 @@ def load_model(choice: str, custom_file: Optional[Path] = None):
 # Inference helpers
 ###############################################################################
 
-BOX_THICKNESS = 2  # thinner boxes
-BOX_ALPHA = 0.6    # 60 % opacity
+BOX_THICKNESS = 2
+BOX_ALPHA = 0.6
 
 box_annotator = sv.BoxAnnotator(thickness=BOX_THICKNESS)
 label_annotator = sv.LabelAnnotator()
@@ -118,7 +116,7 @@ def run_single_inference(model, image: Image.Image, threshold: float) -> Tuple[I
     return Image.fromarray(cv2.cvtColor(blended, cv2.COLOR_BGR2RGB)), runtime
 
 ###############################################################################
-# Gradio generator callback with progress
+# Callback with progress & captions
 ###############################################################################
 
 def compare_models(
@@ -132,50 +130,54 @@ def compare_models(
     if img.mode != "RGB":
         img = img.convert("RGB")
 
-    total_steps = len(models) * 2  # phase 1: load, phase 2: inference
+    total_steps = len(models) * 2
     progress = gr.Progress()
 
-    # ----- Phase 1: preload weights -----
     detectors: Dict[str, object] = {}
     for i, name in enumerate(models, 1):
         try:
             detectors[name] = load_model(name, custom_file)
         except Exception as exc:
-            detectors[name] = exc  # store exception for later reporting
+            detectors[name] = exc
         progress(i, total=total_steps, desc=f"Loading {name}")
 
-    # ----- Phase 2: run inference -----
-    results: List[Image.Image] = []
+    results: List[Tuple[Image.Image, str]] = []
     legends: Dict[str, str] = {}
 
     for j, name in enumerate(models, 1):
-        detector_or_err = detectors[name]
-        step_index = len(models) + j
-        if isinstance(detector_or_err, Exception):
-            # model failed to load
-            results.append(Image.new("RGB", img.size, (40, 40, 40)))
-            emsg = str(detector_or_err)
-            legends[name] = "Unavailable (weights not found)" if "No such file" in emsg or "not found" in emsg else f"ERROR: {emsg.splitlines()[0][:120]}"
-            progress(step_index, total=total_steps, desc=f"Skipped {name}")
+        item = detectors[name]
+        step = len(models) + j
+        if isinstance(item, Exception):
+            placeholder = Image.new("RGB", img.size, (40, 40, 40))
+            emsg = str(item)
+            caption = f"{name} – Unavailable" if "No such file" in emsg or "not found" in emsg else f"{name} – ERROR"
+            results.append((placeholder, caption))
+            legends[name] = caption
+            progress(step, total=total_steps, desc=f"Skipped {name}")
             continue
         try:
-            annotated, latency = run_single_inference(detector_or_err, img, threshold)
-            results.append(annotated)
+            annotated, latency = run_single_inference(item, img, threshold)
+            caption = f"{name} ({latency*1000:.1f} ms)"
+            results.append((annotated, caption))
             legends[name] = f"{latency*1000:.1f} ms"
         except Exception as exc:
-            results.append(Image.new("RGB", img.size, (40, 40, 40)))
+            placeholder = Image.new("RGB", img.size, (40, 40, 40))
+            caption = f"{name} – ERROR"
+            results.append((placeholder, caption))
             legends[name] = f"ERROR: {str(exc).splitlines()[0][:120]}"
-        progress(step_index, total=total_steps, desc=f"Inference {name}")
+        progress(step, total=total_steps, desc=f"Inference {name}")
 
-    yield results, legends  # final output
+    yield results, legends
 
 ###############################################################################
-# Gradio UI
+# UI
 ###############################################################################
 
 def build_demo():
     with gr.Blocks(title="CV Model Comparison") as demo:
-        gr.Markdown("""# 🔍 Compare Object‑Detection Models\nUpload an image, select detectors, then click **Run Inference**.\nThin, semi‑transparent boxes highlight detections.""")
+        gr.Markdown(
+            """# 🔍 Compare Object‑Detection Models\nUpload an image, select detectors, and click **Run Inference**.\nCaptions beneath each result show which model (and latency) generated it."""
+        )
 
         with gr.Row():
             sel_models = gr.CheckboxGroup(ALL_MODELS, value=["YOLOv12‑n"], label="Models")
@@ -188,8 +190,9 @@ def build_demo():
             gallery = gr.Gallery(label="Results", columns=2, height="auto")
         legend_out = gr.JSON(label="Latency / status by model")
 
-        run_btn = gr.Button("Run Inference", variant="primary")
-        run_btn.click(compare_models, [sel_models, img_in, conf_slider, ckpt_file], [gallery, legend_out])
+        gr.Button("Run Inference", variant="primary").click(
+            compare_models, [sel_models, img_in, conf_slider, ckpt_file], [gallery, legend_out]
+        )
 
     return demo