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 for either framework (upload .pt/.pth files)
 
-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.
+Revision 2025‑04‑19‑b:
+  • Fixed indentation error in error‑handling branch.
+  • Unavailable weights report as concise messages.
+  • Bounding boxes: 2 px, 60 % opacity.
 """
 
 from __future__ import annotations
@@ -31,7 +31,6 @@ from rfdetr.util.coco_classes import COCO_CLASSES
 ###############################################################################
 
 YOLO_MODEL_MAP = {
-    # Ultralytics hub IDs — downloaded on first use
     "YOLOv12‑n": "yolov12n.pt",
     "YOLOv12‑s": "yolov12s.pt",
     "YOLOv12‑m": "yolov12m.pt",
@@ -57,13 +56,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):
-    """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:
-            mdl = YOLO(YOLO_MODEL_MAP[choice])  # hub download if needed
+            mdl = YOLO(YOLO_MODEL_MAP[choice])
         elif choice in RFDETR_MODEL_MAP:
             mdl = RFDETRBase() if RFDETR_MODEL_MAP[choice] == "base" else RFDETRLarge()
         elif choice.startswith("Custom YOLO"):
@@ -78,24 +75,21 @@ def load_model(choice: str, custom_file: Optional[Path] = None):
             raise ValueError(f"Unsupported model choice: {choice}")
     except Exception as e:
         raise RuntimeError(str(e)) from e
-
     _loaded[choice] = mdl
     return mdl
 
 ###############################################################################
-# Inference helpers — semi‑transparent, thin boxes
+# Inference helpers
 ###############################################################################
 
-box_annotator = sv.BoxAnnotator(thickness=2)  # thinner lines
+box_annotator = sv.BoxAnnotator(thickness=2)
 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
@@ -103,17 +97,17 @@ def run_single_inference(model, image: Image.Image, threshold: float) -> Tuple[I
         result = model.predict(image, verbose=False)[0]
         detections = sv.Detections.from_ultralytics(result)
         label_source = model.names
-
     runtime = time.perf_counter() - start
-
     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
+    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)
+    return Image.fromarray(cv2.cvtColor(blended, cv2.COLOR_BGR2RGB)), runtime
 
 ###############################################################################
 # Gradio callback
@@ -135,11 +129,9 @@ def compare_models(models: List[str], img: Image.Image, threshold: float, custom
             results.append(annotated)
             legends[m] = f"{t*1000:.1f} ms"
         except Exception as e:
-            # show blank slate if model unavailable
             results.append(Image.new("RGB", img.size, (40, 40, 40)))
-                            err_msg = str(e)
-            # Normalize common weight‑missing errors for clarity
-            if "No such file or directory" in err_msg:
+            err_msg = str(e)
+            if "No such file" in err_msg or "not found" in err_msg:
                 legends[m] = "Unavailable (weights not found)"
             else:
                 legends[m] = f"ERROR: {err_msg.splitlines()[0][:120]}"
@@ -147,27 +139,25 @@ def compare_models(models: List[str], img: Image.Image, threshold: float, custom
     return results, legends
 
 ###############################################################################
-# Build & launch Gradio UI
+# UI
 ###############################################################################
 
 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 and 60 % opaque for clarity.""")
-
+        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(0.0, 1.0, 0.5, step=0.05, label="Confidence threshold")
-
-        custom_checkpoint = gr.File(label="Upload custom checkpoint (.pt/.pth)", file_types=[".pt", ".pth"], interactive=True)
+            model_select = gr.CheckboxGroup(ALL_MODELS, value=["YOLOv12‑n"], label="Select models")
+            threshold_slider = gr.Slider(0.0, 1.0, 0.5, 0.05, label="Confidence threshold")
+        custom_file = 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="Latency / status by model")
-
-        run_btn = gr.Button("Run Inference", variant="primary")
-        run_btn.click(compare_models, [model_select, image_in, threshold_slider, custom_checkpoint], [gallery, legends_out])
-
+        gr.Button("Run Inference", variant="primary").click(
+            compare_models, [model_select, image_in, threshold_slider, custom_file], [gallery, legends_out]
+        )
     return demo
 
 if __name__ == "__main__":