import gradio as gr
import PIL.Image as Image
from ultralytics import YOLO
import spaces

@spaces.GPU
def yolo_inference(images, model_id, conf_threshold, iou_threshold, max_detection):
    model = YOLO(model_id)
    results = model.predict(
        source=images,
        conf=conf_threshold,
        iou=iou_threshold,
        imgsz=640,
        max_det=max_detection,
        show_labels=True,
        show_conf=True,
    )

    annotated_images = []
    for result in results:
        im_array = result.plot()
        im = Image.fromarray(im_array)
        annotated_images.append(im)
    
    return annotated_images  

interface = gr.Interface(
    fn=yolo_inference,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Dropdown(
            choices=['yolo11n.pt', 'yolo11s.pt', 'yolo11m.pt', 'yolo11l.pt', 'yolo11x.pt', 
                     'yolo11n-seg.pt', 'yolo11s-seg.pt', 'yolo11m-seg.pt', 'yolo11l-seg.pt', 'yolo11x-seg.pt',
                     'yolo11n-pose.pt', 'yolo11s-pose.pt', 'yolo11m-pose.pt', 'yolo11l-pose.pt', 'yolo11x-pose.pt',
                     'yolo11n-obb.pt', 'yolo11s-obb.pt', 'yolo11m-obb.pt', 'yolo11l-obb.pt', 'yolo11x-obb.pt',
                     'yolo11n-cls.pt', 'yolo11s-cls.pt', 'yolo11m-cls.pt', 'yolo11l-cls.pt', 'yolo11x-cls.pt'],
            label="Model Name",
            value="yolo11n.pt",
        ),
        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence Threshold"),
        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU Threshold"),
        gr.Slider(minimum=1, maximum=300, step=1, value=300, label="Max Detection"),
    ],
    outputs=gr.Gallery(label="Annotated Image"),
    title="Yolo11: Object Detection",
    description="Upload image(s) for inference using the latest Ultralytics YOLO11 models.",
    examples=[
        [["zidane.jpg"], "yolo11s.pt", 0.25, 0.45, 300],
        [["bus.jpg"], "yolo11m.pt", 0.25, 0.45, 300],
        [["yolo_vision.jpg"], "yolo11x.pt", 0.25, 0.45, 300],
    ],
)
interface.launch()