Create app.py

app.py

@@ -0,0 +1,149 @@
+"""
+Gradio app to compare object‑detection models:
+  • Ultralytics YOLOv12 (n, s, m, l, x)
+  • Ultralytics YOLOv11 (n, s, m, l, x)
+  • Roboflow RF‑DETR (Base, Large)
+  • Custom fine‑tuned checkpoints for either framework
+Requires Python ≥3.9 plus:
+  pip install gradio ultralytics rfdetr supervision pillow numpy torch torchvision
+If you need ONNX export for RF‑DETR, also: pip install rfdetr[onnxexport]
+"""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+from typing import List, Tuple
+
+import numpy as np
+from PIL import Image
+import gradio as gr
+import supervision as sv
+from ultralytics import YOLO
+from rfdetr import RFDETRBase, RFDETRLarge
+from rfdetr.util.coco_classes import COCO_CLASSES
+
+# -----------------------------------------------------------------------------
+# Model registry & lazy loader
+# -----------------------------------------------------------------------------
+
+YOLO_MODEL_MAP = {
+    # YOLOv12 sizes
+    "YOLOv12‑n": "yolov12n.pt",
+    "YOLOv12‑s": "yolov12s.pt",
+    "YOLOv12‑m": "yolov12m.pt",
+    "YOLOv12‑l": "yolov12l.pt",
+    "YOLOv12‑x": "yolov12x.pt",
+    # YOLOv11 sizes
+    "YOLOv11‑n": "yolov11n.pt",
+    "YOLOv11‑s": "yolov11s.pt",
+    "YOLOv11‑m": "yolov11m.pt",
+    "YOLOv11‑l": "yolov11l.pt",
+    "YOLOv11‑x": "yolov11x.pt",
+}
+
+RFDETR_MODEL_MAP = {
+    "RF‑DETR‑Base (29M)": "base",  # handled explicitly
+    "RF‑DETR‑Large (128M)": "large",
+}
+
+ALL_MODELS = list(YOLO_MODEL_MAP.keys()) + list(RFDETR_MODEL_MAP.keys()) + [
+    "Custom YOLO (.pt/.pth)",
+    "Custom RF‑DETR (.pth)",
+]
+
+_loaded = {}
+
+def load_model(choice: str, custom_path: str | None = None):
+    """Lazy‑load and cache models to avoid re‑download between inferences."""
+    global _loaded
+    if choice in _loaded:
+        return _loaded[choice]
+
+    if choice in YOLO_MODEL_MAP:
+        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"):
+        if not custom_path:
+            raise ValueError("Please provide a path to your YOLO checkpoint.")
+        mdl = YOLO(custom_path)
+    elif choice.startswith("Custom RF‑DETR"):
+        if not custom_path:
+            raise ValueError("Please provide a path to your RF‑DETR checkpoint.")
+        mdl = RFDETRBase(pretrain_weights=custom_path)
+    else:
+        raise ValueError(f"Unsupported model choice: {choice}")
+
+    _loaded[choice] = mdl
+    return mdl
+
+# -----------------------------------------------------------------------------
+# Inference helpers
+# -----------------------------------------------------------------------------
+
+box_annotator = sv.BoxAnnotator()
+label_annotator = sv.LabelAnnotator()
+
+def run_single_inference(model, image: Image.Image, threshold: float) -> Tuple[Image.Image, float]:
+    start = time.perf_counter()
+
+    # RF‑DETR already returns sv.Detections
+    if isinstance(model, (RFDETRBase, RFDETRLarge)):
+        detections = model.predict(image, threshold=threshold)
+        label_source = COCO_CLASSES
+    else:
+        # Ultralytics YOLO inference: returns list of Results
+        result = model.predict(image, verbose=False)[0]
+        detections = sv.Detections.from_ultralytics(result)
+        label_source = model.names  # list of class 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
+
+# -----------------------------------------------------------------------------
+# Gradio UI logic
+# -----------------------------------------------------------------------------
+
+def compare_models(models: List[str], img: Image.Image, threshold: float, custom_path: str | None):
+    if img.mode != "RGB":
+        img = img.convert("RGB")
+    results = []
+    legends = []
+    for m in models:
+        model_obj = load_model(m, custom_path)
+        annotated, t = run_single_inference(model_obj, img, threshold)
+        results.append(annotated)
+        legends.append(f"{m} – {t*1000:.1f} ms")
+    return results, legends
+
+# -----------------------------------------------------------------------------
+# Launch Gradio Interface
+# -----------------------------------------------------------------------------
+
+def build_demo():
+    with gr.Blocks(title="CV Model Comparison") as demo:
+        gr.Markdown("""# 🔍 Compare Object‑Detection Models\nUpload an image and select one or more models to see their predictions side‑by‑side.""")
+
+        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")
+        custom_weight_path = gr.Textbox(label="Path to custom checkpoint (if selected)")
+        image_in = gr.Image(type="pil", label="Upload image")
+
+        with gr.Row():
+            gallery = gr.Gallery(label="Annotated results", columns=2, height="auto")
+
+        legends_out = gr.JSON(label="Runtime (ms)")
+
+        run_btn = gr.Button("Run Inference")
+        run_btn.click(compare_models, inputs=[model_select, image_in, threshold_slider, custom_weight_path], outputs=[gallery, legends_out])
+
+    return demo
+
+# Execute when running directly
+if __name__ == "__main__":
+    build_demo().launch()