Gradio example

[robmarkcole]

I created an app to allow using supervision with gradio to viz ultralytics predictions. It allows filtering on confidence and area of detections. I implemented a basic cache to make it more performant than this example. I will extend to enable metric calcs soon. Any other suggestions welcome!

# app.py
# Run: python app.py
#
# Single-image object detection web app using Ultralytics YOLOv11 and Supervision.
# Configure MODEL_PATH, CONF_THRESHOLD, IOU_THRESHOLD, and MAX_EXAMPLES via environment variables if desired.
import os
from dotenv import load_dotenv

import cv2
import numpy as np
import pandas as pd
from PIL import Image

import gradio as gr
import supervision as sv
from ultralytics import YOLO

# Load environment variables
load_dotenv()
MODEL_PATH = os.getenv("MODEL_PATH", "yolo11n.pt")
DEFAULT_CONF_THRESHOLD = float(os.getenv("CONF_THRESHOLD", 0.5))
DEFAULT_IOU_THRESHOLD = float(os.getenv("IOU_THRESHOLD", 0.55))
FOLDER_PATH = os.getenv("FOLDER_PATH", "./images")
# Maximum number of example images to display
MAX_EXAMPLES = int(os.getenv("MAX_EXAMPLES", 20))
# Default minimum detection area in pixels
DEFAULT_AREA_THRESHOLD = int(os.getenv("AREA_THRESHOLD", 100))

# Initialize YOLO model
model = YOLO(MODEL_PATH)

# Simple cache for raw detection outputs keyed by image hash and thresholds
_DETECTION_CACHE: dict[tuple[int, float, float], sv.Detections] = {}

def predict_image(
    img: Image,
    conf_threshold: float,
    iou_threshold: float,
    show_labels: bool,
    area_threshold: int,
    show_table: bool,
) -> tuple[np.ndarray, pd.DataFrame]:
    """
    Runs YOLO inference on a PIL image, caches detections, converts to Supervision detections,
    filters by confidence and area thresholds, annotates segmentation masks (and optional labels),
    and returns the annotated RGB image and a detection table.
    """
    # Convert PIL to BGR numpy array
    image_bgr = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
    # Compute cache key
    key = (hash(img.tobytes()), conf_threshold, iou_threshold)
    # Get or run inference
    if key in _DETECTION_CACHE:
        detections = _DETECTION_CACHE[key]
    else:
        results = model.predict(
            source=image_bgr,
            conf=0.0,
            iou=iou_threshold,
            show_labels=False,
            show_conf=False,
        )
        detections = sv.Detections.from_ultralytics(results[0])
        _DETECTION_CACHE[key] = detections
    # Filter by confidence and area
    filtered = detections[detections.confidence > conf_threshold]
    filtered = filtered[filtered.area > area_threshold]
    # Annotate mask
    mask_annotator = sv.MaskAnnotator()
    annotated_bgr = mask_annotator.annotate(scene=image_bgr, detections=filtered)
    # Optionally annotate labels (confidence + area)
    if show_labels:
        labels = [f"{conf:.2f} ({int(area)})" for conf, area in zip(filtered.confidence, filtered.area)]
        label_annotator = sv.LabelAnnotator()
        annotated_bgr = label_annotator.annotate(
            scene=annotated_bgr,
            detections=filtered,
            labels=labels,
        )
    # Convert to RGB
    annotated_rgb = annotated_bgr[..., ::-1]
    # Build table if requested
    if show_table:
        table = pd.DataFrame({
            "class_id": filtered.class_id,
            "confidence": filtered.confidence,
            "area": filtered.area,
        })
    else:
        table = pd.DataFrame()
    return annotated_rgb, table

# Prepare example images
image_files = [
    os.path.join(FOLDER_PATH, f)
    for f in os.listdir(FOLDER_PATH)
    if f.lower().endswith((".jpg", ".jpeg", ".png"))
][:MAX_EXAMPLES]

# Build Gradio interface
iface = gr.Interface(
    fn=predict_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(0, 1, value=DEFAULT_CONF_THRESHOLD, label="Confidence threshold"),
        gr.Slider(0, 1, value=DEFAULT_IOU_THRESHOLD, label="IoU threshold"),
        gr.Checkbox(label="Show labels", value=True),
        gr.Slider(1, 500, value=DEFAULT_AREA_THRESHOLD, label="Min detection area (px)"),
        gr.Checkbox(label="Show detection table", value=False),
    ],
    outputs=[
        gr.Image(type="numpy", label="Result"),
        gr.Dataframe(headers=["class_id", "confidence", "area"], label="Detections Table"),
    ],
    title="YOLO11 + Supervision Annotator",
    description="Upload an image, filter and annotate detections, and optionally view a detection table.",
    examples=[[img, DEFAULT_CONF_THRESHOLD, DEFAULT_IOU_THRESHOLD, True, DEFAULT_AREA_THRESHOLD, False] for img in image_files],
)

if __name__ == "__main__":
    iface.launch(debug=True)