import cv2
import gradio as gr
from ultralytics import YOLO

model = YOLO("best.pt")

def predict(img):

    pred = model(img)
    frame_with_bbox = pred[0].plot()
    return frame_with_bbox



gradio_app = gr.Interface(
             fn=predict,
             inputs=gr.Image(type="numpy"),
             outputs=gr.Image(type="numpy"))

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