app.py

"""Gradio Space for YOLO image detection."""

from __future__ import annotations

from dataclasses import asdict

import gradio as gr
from PIL import Image

from inference import load_model, predict, resolve_model_path


MODEL_PATH = resolve_model_path()
MODEL = load_model(MODEL_PATH)


def run_detection(
    image: Image.Image | None,
    confidence: float,
    iou: float,
    image_size: int,
    classes: str,
) -> tuple[Image.Image | None, dict[str, object]]:
    """Gradio callback for one uploaded image."""

    if image is None:
        return None, {"error": "Upload an image to run detection."}

    class_filter = classes.strip() or None
    detections, annotated = predict(
        image=image,
        model=MODEL,
        conf=confidence,
        iou=iou,
        imgsz=image_size,
        classes=class_filter,
    )

    return annotated, {
        "model": MODEL_PATH,
        "count": len(detections),
        "detections": [asdict(detection) for detection in detections],
    }


with gr.Blocks(title="YOLO Object Detection") as demo:
    gr.Markdown(
        "# YOLO Object Detection\n"
        "Upload an image, run YOLO detection, and view bounding boxes plus JSON results."
    )

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Image")
            confidence = gr.Slider(0.05, 0.95, value=0.35, step=0.05, label="Confidence")
            iou = gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="IoU")
            image_size = gr.Slider(320, 1536, value=1280, step=32, label="Image size")
            classes = gr.Textbox(
                value="person",
                label="Classes",
                placeholder="person or person,car or 0,2. Leave empty for all classes.",
            )
            detect_button = gr.Button("Detect", variant="primary")
        with gr.Column():
            output_image = gr.Image(type="pil", label="Annotated image")
            output_json = gr.JSON(label="Detections")

    detect_button.click(
        fn=run_detection,
        inputs=[input_image, confidence, iou, image_size, classes],
        outputs=[output_image, output_json],
    )


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