DevashishBhake's picture
new file: app.py
5604609
raw
history blame
645 Bytes
import torch
import gradio as gr
from PIL import Image
import numpy as np
model = torch.hub.load('ultralytics/yolov5', 'custom', path=r'./best.pt', force_reload=True)
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()