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

# Load model directly from file in the Space
detection_model = YOLO("best.pt")

# Prediction function
def predict(pilimg):
    results = detection_model.predict(pilimg, conf=0.5, iou=0.6)
    img_bgr = results[0].plot()
    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # Convert BGR to RGB
    return out_pilimg

# Gradio interface
gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Mask Detection Demo"
).launch()