import os
import gradio as gr

from ultralytics import YOLO

model = YOLO("yolo26m-obb.onnx")

example_list = [
    ["examples/example_1.png", 0.25, 0.45],
    ["examples/example_2.jpg", 0.25, 0.45]
]
title = "Ultralytics Gradio YOLO26"
description = "Upload images for YOLO26 obb detection."

def predict_image(img, conf_threshold, iou_threshold):
    if img is None:
        return None
    
    conf = conf_threshold if conf_threshold is not None else 0.25
    iou = iou_threshold if iou_threshold is not None else 0.45
    
    results = model.predict(
        source=img,
        conf=conf,
        iou=iou,
        show_labels=True,
        show_conf=True,
    )
    return results[0].plot(boxes=True, probs=False, line_width=1,) if results else None

iface = gr.Interface(
    fn=predict_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
    ],
    examples=example_list,
    outputs=gr.Image(type="pil", label="Result"),
    title=title,
    description=description,
)
iface.launch()