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


def predict(pilimg: Image.Image, conf_thresh: float, iou_thresh: float) -> Image.Image:
    results = detection_model.predict(pilimg, conf=conf_thresh, iou=iou_thresh)
    img_bgr = results[0].plot()
    return Image.fromarray(img_bgr[..., ::-1])


detection_model = YOLO("best_int8_openvino_model", task="detect")

demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Slider(0.1, 1.0, value=0.5, step=0.05, label="Confidence Threshold"),
        gr.Slider(0.1, 1.0, value=0.6, step=0.05, label="IoU Threshold"),
    ],
    outputs=gr.Image(type="pil", label="Detections"),
    title="Playing Card & License Plate Detector",
    description=(
        "Upload an image containing **playing cards** and/or **license plates**. "
        "The model will draw bounding boxes around detected objects.\n\n"
        "Model: YOLO11s fine-tuned with OpenVINO INT8 export."
    ),
    examples=[],
)

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