import gradio as gr
from ultralytics import YOLO

# Load the model
model = YOLO('best.pt')

def predict(image):
    results = model.predict(image, conf=0.3)
    return results[0].plot()

# Create the Gradio interface
demo = gr.Interface(
    fn=predict, 
    inputs=gr.Image(type="pil"), 
    outputs=gr.Image(type="numpy"),
    title="Instrument Detector"
)
demo.launch()