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)

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.
"""

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 = {
    # Ultralytics hub IDs — downloaded on first use
    "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):
    """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
        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 — semi‑transparent, thin boxes
###############################################################################

box_annotator = sv.BoxAnnotator(thickness=2)  # thinner lines
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
    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)  # semi‑transparent boxes
    annotated_pil = Image.fromarray(cv2.cvtColor(blended, cv2.COLOR_BGR2RGB))

    return annotated_pil, 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:
            # 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:
                legends[m] = "Unavailable (weights not found)"
            else:
                legends[m] = f"ERROR: {err_msg.splitlines()[0][:120]}"

    return results, legends

###############################################################################
# Build & launch Gradio 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(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)
        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])

    return demo

if __name__ == "__main__":
    build_demo().launch()