app.py

"""
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 (upload .pt/.pth files)

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

import time
from pathlib import Path
from typing import List, Tuple, Dict, Optional

import cv2
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‑n": "yolov12n.pt",
    "YOLOv12‑s": "yolov12s.pt",
    "YOLOv12‑m": "yolov12m.pt",
    "YOLOv12‑l": "yolov12l.pt",
    "YOLOv12‑x": "yolov12x.pt",
    "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",
    "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: Dict[str, object] = {}

def load_model(choice: str, custom_file: Optional[Path] = None):
    if choice in _loaded:
        return _loaded[choice]
    try:
        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_file:
                raise ValueError("Upload a YOLO .pt/.pth checkpoint first.")
            mdl = YOLO(str(custom_file))
        elif choice.startswith("Custom RF‑DETR"):
            if not custom_file:
                raise ValueError("Upload an RF‑DETR .pth checkpoint first.")
            mdl = RFDETRBase(pretrain_weights=str(custom_file))
        else:
            raise ValueError(f"Unsupported model choice: {choice}")
    except Exception as e:
        raise RuntimeError(str(e)) from e
    _loaded[choice] = mdl
    return mdl

###############################################################################
# Inference helpers
###############################################################################

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:
    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:
        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)
    return Image.fromarray(cv2.cvtColor(blended, cv2.COLOR_BGR2RGB)), runtime

###############################################################################
# Gradio callback
###############################################################################

def compare_models(models: List[str], img: Image.Image, threshold: float, custom_file: Optional[Path]):
    if img is None:
        raise gr.Error("Please upload an image first.")
    if img.mode != "RGB":
        img = img.convert("RGB")

    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[m] = f"{t*1000:.1f} ms"
        except Exception as e:
            results.append(Image.new("RGB", img.size, (40, 40, 40)))
            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]}"

    return results, legends

###############################################################################
# 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."""
        )
        with gr.Row():
            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")
        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__":
    build_demo().launch()