import gradio as gr
import torch
from sahi.prediction import ObjectPrediction
from sahi.utils.cv import visualize_object_predictions, read_image
from ultralyticsplus import YOLO, render_result

# your example images
image_path = [
    ['test/web form.jpg', 'foduucom/web-form-ui-field-detection', 640, 0.25, 0.45],
    ['test/web form2.jpg', 'foduucom/web-form-ui-field-detection', 640, 0.25, 0.45]
]

def yolov8_inference(
    image,           # will be a filepath string
    model_path,      # string
    image_size,      # int
    conf_threshold,  # float
    iou_threshold    # float
):
    # load and configure the model
    model = YOLO(model_path)
    model.overrides.update({
        'conf': conf_threshold,
        'iou': iou_threshold,
        'agnostic_nms': False,
        'max_det': 1000
    })

    # read & run
    img = read_image(image)
    results = model.predict(img)
    rendered = render_result(model=model, image=img, result=results[0])
    return rendered

# define components using the new API
inputs = [
    gr.Image(type="filepath", label="Input Image"),
    gr.Dropdown(
        choices=["foduucom/web-form-ui-field-detection"],
        value="foduucom/web-form-ui-field-detection",
        label="Model"
    ),
    gr.Slider(320, 1280, step=32, value=640, label="Image Size"),
    gr.Slider(0.0, 1.0, step=0.05, value=0.25, label="Confidence Threshold"),
    gr.Slider(0.0, 1.0, step=0.05, value=0.45, label="IOU Threshold"),
]

outputs = gr.Image(type="filepath", label="Output Image")

title = "Web-Form UI Field Detection"

# single-tab interface
interface = gr.Interface(
    fn=yolov8_inference,
    inputs=inputs,
    outputs=outputs,
    title=title,
    examples=image_path,
    cache_examples=False,
    theme="huggingface"
)

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