import gradio as gr
from PIL import Image
import numpy as np
import torch

model = torch.hub.load('./yolov5', 'custom', path='./best.pt', force_reload=True, source='local')

def yolo(im, size=512):
    g = (size / max(im.size))  # gain
    im = im.resize((int(x * g) for x in im.size), Image.ANTIALIAS)  # resize

    results = model(im)  # inference
    results.render()  # updates results.imgs with boxes and labels
    return Image.fromarray(results.ims[0])

gr.Interface(fn=yolo, 
             inputs=gr.inputs.Image(type = "pil", label = "Original Image"), 
             outputs=gr.outputs.Image(type = "pil", label = "Output Image")).launch()