import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image

# Load model once at startup
model = YOLO("best.pt")

def detect(image):
    # Convert PIL to numpy BGR for YOLO
    img_array = np.array(image)
    img_bgr   = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)

    # Run inference on CPU
    results = model.predict(
        img_bgr,
        conf   = 0.4,
        iou    = 0.45,
        device = "cpu",
    )

    # plot() draws boxes and returns BGR numpy array
    annotated_bgr = results[0].plot()

    # Convert back to RGB for Gradio display
    annotated_rgb = cv2.cvtColor(annotated_bgr, cv2.COLOR_BGR2RGB)
    return Image.fromarray(annotated_rgb)

demo = gr.Interface(
    fn          = detect,
    inputs      = gr.Image(type="pil", label="Upload Image"),
    outputs     = gr.Image(type="pil", label="Detection Result"),
    title       = "Headphone & Calculator Detector",
    description = "Upload an image containing headphones and/or a calculator. The model will detect and label them with bounding boxes.",
    examples    = [],
)

demo.launch()