import torch
import ultralytics
from ultralytics import YOLO
import gradio as gr
import cv2

model = YOLO("best.pt")

def detect(image):
    results = model.predict(image, imgsz=640, conf=0.1, verbose=False)
    bgr = results[0].plot()
    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    return rgb

demo = gr.Interface(
    fn=detect,
    inputs=gr.Image(type="numpy", label="Upload image"),
    outputs=gr.Image(type="numpy", label="Detection result"),
)

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