Spaces:
Running
Running
#!/usr/bin/env python | |
# coding: utf-8 | |
# In[ ]: | |
# importing required libraries | |
from transformers import pipeline | |
import gradio as gr | |
from PIL import Image, ImageDraw | |
# main function for object detection | |
def detector(raw): | |
# Resize the image | |
WIDTH = 800 | |
width, height = raw.size | |
ratio = float(WIDTH) / float(width) | |
new_h = height * ratio | |
ip_img = raw.resize((int(WIDTH), int(new_h)), Image.Resampling.LANCZOS) | |
# load the model pipeline and predict | |
outs = pipeline(model="hustvl/yolos-tiny")(ip_img) | |
# draw the image on the canvas | |
draw = ImageDraw.Draw(ip_img) | |
# draw the boxes with labels | |
for object in outs: | |
score = f"{object['score']*100:.2f}%" | |
label = object['label'] | |
xmin, ymin, xmax, ymax = object['box'].values() | |
draw.rectangle((xmin, ymin, xmax, ymax), outline='red', width=1) | |
draw.text((xmin, ymin), f"{label}: {score}", fill="black") | |
return ip_img | |
demo = gr.Interface(fn=detector, | |
inputs=gr.Image(type='pil'), | |
outputs=gr.Image(type='pil'), allow_flagging=False) | |
demo.queue(True) | |
demo.launch(debug=True, inline=False, show_api=False, share=False) | |