DevashishBhake's picture
Update app.py
9b6527a verified
raw
history blame
No virus
650 Bytes
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()