rgp's picture
Update app.py
1535058
raw
history blame
1.38 kB
import torch
import gradio as gr
from huggingface_hub import hf_hub_download
from PIL import Image
REPO_ID = "rgp/yolov5-street-view-detection"
FILENAME = "best.pt"
yolov5_weights = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
model = torch.hub.load('ultralytics/yolov5', 'custom', path=yolov5_weights, force_reload=True) # local repo
def object_detection(im, size=640):
results = model(im) # inference
#results.print() # print results to screen
#results.show() # display results
#results.save() # save as results1.jpg, results2.jpg... etc.
results.render() # updates results.imgs with boxes and labels
return Image.fromarray(results.ims[0])
title = "Identificação de Pedestres e meios de locomoção nas ruas"
description = """Esse modelo é uma pequena demonstração baseada em uma análise de cerca de 680 imagens.
"""
input = gr.inputs.Image(shape=(640, 640), image_mode="RGB", source="upload", label="Imagem", optional=False)
output = gr.outputs.Image(type="pil", label="Output Image")
examples = [["sample_images/IMG_0125.jpeg"], ["sample_images/IMG_0129.jpg"], ["sample_images/IMG_0157.jpg"], ["sample_images/IMG_0158.jpg"], ["sample_images/IMG_012.jpg"]]
gr.Interface(
fn=object_detection,
inputs=input,
outputs=output,
title=title,
description=description,
examples=examples
).launch(debug=True)