Update app.py

app.py

@@ -5,9 +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‑c:
-  • Re‑indented entire file with 4‑space consistency to remove `IndentationError`.
-  • Thin, semi‑transparent 60 % boxes; concise error labels.
+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.
 """
 
 from __future__ import annotations
@@ -26,7 +27,7 @@ from rfdetr import RFDETRBase, RFDETRLarge
 from rfdetr.util.coco_classes import COCO_CLASSES
 
 ###############################################################################
-# Model registry & lazy loader
+# Model registry & cache
 ###############################################################################
 
 YOLO_MODEL_MAP: Dict[str, str] = {
@@ -55,27 +56,24 @@ 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):
-    """Return and cache a detector matching *choice*."""
+    """Fetch and cache a detector instance for *choice*."""
     if choice in _loaded:
         return _loaded[choice]
 
-    try:
-        if choice in YOLO_MODEL_MAP:
-            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"):
-            if custom_file is None:
-                raise ValueError("Upload a YOLO .pt/.pth checkpoint first.")
-            model = YOLO(str(custom_file))
-        elif choice.startswith("Custom RF‑DETR"):
-            if custom_file is None:
-                raise ValueError("Upload an RF‑DETR .pth checkpoint first.")
-            model = RFDETRBase(pretrain_weights=str(custom_file))
-        else:
-            raise ValueError(f"Unsupported model choice: {choice}")
-    except Exception as exc:
-        raise RuntimeError(str(exc)) from exc
+    if choice in YOLO_MODEL_MAP:
+        model = YOLO(YOLO_MODEL_MAP[choice])  # Ultralytics auto‑downloads if missing
+    elif choice in RFDETR_MODEL_MAP:
+        model = RFDETRBase() if RFDETR_MODEL_MAP[choice] == "base" else RFDETRLarge()
+    elif choice.startswith("Custom YOLO"):
+        if custom_file is None:
+            raise RuntimeError("Upload a YOLO .pt/.pth checkpoint first.")
+        model = YOLO(str(custom_file))
+    elif choice.startswith("Custom RF‑DETR"):
+        if custom_file is None:
+            raise RuntimeError("Upload an RF‑DETR .pth checkpoint first.")
+        model = RFDETRBase(pretrain_weights=str(custom_file))
+    else:
+        raise RuntimeError(f"Unsupported model choice: {choice}")
 
     _loaded[choice] = model
     return model
@@ -84,8 +82,8 @@ def load_model(choice: str, custom_file: Optional[Path] = None):
 # Inference helpers
 ###############################################################################
 
-BOX_THICKNESS = 2
-BOX_ALPHA = 0.6
+BOX_THICKNESS = 2  # thinner boxes
+BOX_ALPHA = 0.6    # 60 % opacity
 
 box_annotator = sv.BoxAnnotator(thickness=BOX_THICKNESS)
 label_annotator = sv.LabelAnnotator()
@@ -100,28 +98,25 @@ def run_single_inference(model, image: Image.Image, threshold: float) -> Tuple[I
         detections = model.predict(image, threshold=threshold)
         label_src = COCO_CLASSES
     else:
-        ul_result = model.predict(image, verbose=False)[0]
-        detections = sv.Detections.from_ultralytics(ul_result)
+        ul_res = model.predict(image, verbose=False)[0]
+        detections = sv.Detections.from_ultralytics(ul_res)
         label_src = model.names  # type: ignore
 
     runtime = time.perf_counter() - start
 
-    base_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-    overlay = base_bgr.copy()
-
+    img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    overlay = img_bgr.copy()
     overlay = box_annotator.annotate(overlay, detections)
     overlay = label_annotator.annotate(
         overlay,
         detections,
-        [f"{label_src[cid]} {conf:.2f}" for cid, conf in zip(detections.class_id, detections.confidence)],
+        [f"{label_src[c]} {p:.2f}" for c, p in zip(detections.class_id, detections.confidence)],
     )
-
-    blended = _blend(base_bgr, overlay)
-    out_pil = Image.fromarray(cv2.cvtColor(blended, cv2.COLOR_BGR2RGB))
-    return out_pil, runtime
+    blended = _blend(img_bgr, overlay)
+    return Image.fromarray(cv2.cvtColor(blended, cv2.COLOR_BGR2RGB)), runtime
 
 ###############################################################################
-# Gradio callback
+# Gradio generator callback with progress
 ###############################################################################
 
 def compare_models(
@@ -135,24 +130,42 @@ def compare_models(
     if img.mode != "RGB":
         img = img.convert("RGB")
 
+    total_steps = len(models) * 2  # phase 1: load, phase 2: inference
+    progress = gr.Progress(total=total_steps)
+
+    # ----- 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
+        progress.update(i, desc=f"Loading {name}")
+
+    # ----- Phase 2: run inference -----
     results: List[Image.Image] = []
     legends: Dict[str, str] = {}
 
-    for model_name in models:
+    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.update(step_index, desc=f"Skipped {name}")
+            continue
         try:
-            detector = load_model(model_name, custom_file)
-            annotated, latency = run_single_inference(detector, img, threshold)
+            annotated, latency = run_single_inference(detector_or_err, img, threshold)
             results.append(annotated)
-            legends[model_name] = f"{latency*1000:.1f} ms"
+            legends[name] = f"{latency*1000:.1f} ms"
         except Exception as exc:
             results.append(Image.new("RGB", img.size, (40, 40, 40)))
-            emsg = str(exc)
-            if "No such file" in emsg or "not found" in emsg:
-                legends[model_name] = "Unavailable (weights not found)"
-            else:
-                legends[model_name] = f"ERROR: {emsg.splitlines()[0][:120]}"
+            legends[name] = f"ERROR: {str(exc).splitlines()[0][:120]}"
+        progress.update(step_index, desc=f"Inference {name}")
 
-    return results, legends
+    yield results, legends  # final output
 
 ###############################################################################
 # Gradio UI
@@ -160,9 +173,7 @@ def compare_models(
 
 def build_demo():
     with gr.Blocks(title="CV Model Comparison") as demo:
-        gr.Markdown(
-            """# 🔍 Compare Object‑Detection Models\nUpload an image, choose detectors, and optionally add a custom checkpoint.\nBounding boxes are thin (2 px) and 60 % transparent for clarity."""
-        )
+        gr.Markdown("""# 🔍 Compare Object‑Detection Models\nUpload an image, select detectors, then click **Run Inference**.\nThin, semi‑transparent boxes highlight detections.""")
 
         with gr.Row():
             sel_models = gr.CheckboxGroup(ALL_MODELS, value=["YOLOv12‑n"], label="Models")