import gradio as gr
from PIL import Image, ImageDraw, ImageFont
from ultralytics import YOLO
import spaces

@spaces.GPU
def yolo_inference(images, model_id, conf_threshold, iou_threshold, max_detection):
    if images is None:
        # Create a blank image
        width, height = 640, 480
        blank_image = Image.new("RGB", (width, height), color="white")
        draw = ImageDraw.Draw(blank_image)
        message = "No image provided"
        font = ImageFont.load_default(size=40)
        bbox = draw.textbbox((0, 0), message, font=font)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]
        text_x = (width - text_width) / 2
        text_y = (height - text_height) / 2
        draw.text((text_x, text_y), message, fill="black", font=font)
        return blank_image
    
    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,
    )

    for r in results:
        image_array = r.plot()
        image = Image.fromarray(image_array[..., ::-1])
    return image

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.Image(type="pil", label="Annotated Image"),
    cache_examples=True,
    title="Yolo11: Object Detection, Instance Segmentation, Pose/Keypoints, Oriented Detection, Classification",
    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],
        ["Tricycle.jpg", "yolo11x-cls.pt", 0.25, 0.45, 300],
        ["tcganadolu.jpg", "yolo11m-obb.pt", 0.25, 0.45, 300],
        ["San Diego Airport.jpg", "yolo11x-seg.pt", 0.25, 0.45, 300],
        ["Theodore_Roosevelt.png", "yolo11l-pose.pt", 0.25, 0.45, 300],
    ],
)
interface.launch()