import gradio as gr
from PIL import Image
import torch

model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')

# Confidence threshold
model.conf = 0.15

# NMS IoU threshold
model.iou = 0.50

imgs = '2.jpg'


def object_detection(img):

  result = model(img, size=416)
  result.show()

  return Image.fromarray(result.imgs[0])

iface = gr.Interface(
    object_detection,
    gr.inputs.Image(),
    "image"
)

iface.launch(
    # debug=True
)


object_detection(imgs)