"""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()