Spaces:
Runtime error
Runtime error
import gradio as gr | |
import torch | |
import torchvision | |
import numpy as np | |
from PIL import Image | |
# Load model weights | |
model = torch.hub.load('ultralytics/yolov5', 'custom', "model_weights/datasets_1000_41class.pt") | |
# Define a yolo prediction function | |
def yolo(im, size=640): | |
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.imgs[0]) | |
inputs = gr.inputs.Image(type='pil', label="Original Image") | |
outputs = gr.outputs.Image(type="pil", label="Output Image") | |
title = "BandiCount: Detecting Australian native animal species" | |
description = "BandiCount: Detecting Australian native animal species in NSW national parks, using object detection. Upload an image or click an example image to use." | |
article = "" | |
examples = [['data/Bandicoot.jpg'], ['data/BrushtailPossum.jpg'], ['data/Lyrebird.jpg'], ['data/Macropod.jpg'], \ | |
['data/Wombat.jpg'], ['data/baby_wombat.jpg'], ['data/bush_turkey.jpg'], ['data/cat.jpg'], ['data/dingo.jpg'], \ | |
['data/dingo2.jpg'], ['data/echidna.gif'], ['data/fox.jpg'], ['data/goats.jpg'], ['data/kookaburras.jpg'], \ | |
['data/lyrebird2.jpg'], ['data/pademelon.jpg'], ['data/quoll.jpg']] | |
gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, theme="huggingface").launch(enable_queue=True) | |