Spaces:
Running
on
Zero
Running
on
Zero
import PIL.Image as Image | |
import gradio as gr | |
from ultralytics import ASSETS, YOLOv10 | |
def predict_image(img, model_id, image_size, conf_threshold): | |
model = YOLOv10.from_pretrained(f'jameslahm/{model_id}') | |
results = model.predict( | |
source=img, | |
conf=conf_threshold, | |
show_labels=True, | |
show_conf=True, | |
imgsz=image_size, | |
) | |
for r in results: | |
im_array = r.plot() | |
im = Image.fromarray(im_array[..., ::-1]) | |
return im | |
def app(): | |
with gr.Blocks(): | |
with gr.Row(): | |
with gr.Column(): | |
image = gr.Image(type="pil", label="Image") | |
model_id = gr.Dropdown( | |
label="Model", | |
choices=[ | |
"yolov10n", | |
"yolov10s", | |
"yolov10m", | |
"yolov10b", | |
"yolov10l", | |
"yolov10x", | |
], | |
value="yolov10m", | |
) | |
image_size = gr.Slider( | |
label="Image Size", | |
minimum=320, | |
maximum=1280, | |
step=32, | |
value=640, | |
) | |
conf_threshold = gr.Slider( | |
label="Confidence Threshold", | |
minimum=0.1, | |
maximum=1.0, | |
step=0.1, | |
value=0.25, | |
) | |
yolov10_infer = gr.Button(value="Detect Objects") | |
with gr.Column(): | |
output_image = gr.Image(type="pil", label="Annotated Image") | |
yolov10_infer.click( | |
fn=predict_image, | |
inputs=[ | |
image, | |
model_id, | |
image_size, | |
conf_threshold, | |
], | |
outputs=[output_image], | |
) | |
gr.Examples( | |
examples=[ | |
[ | |
"bug.jpg", | |
"yolov10x", | |
640, | |
0.25, | |
], | |
[ | |
"zidane.jpg", | |
"yolov10m", | |
640, | |
0.25, | |
], | |
], | |
fn=predict_image, | |
inputs=[ | |
image, | |
model_id, | |
image_size, | |
conf_threshold, | |
iou_threshold, | |
], | |
outputs=[output_image], | |
cache_examples=True, | |
) | |
gradio_app = gr.Blocks() | |
with gradio_app: | |
gr.HTML( | |
""" | |
<h1 style='text-align: center'> | |
YOLOv10: Real-Time End-to-End Object Detection | |
</h1> | |
""") | |
gr.HTML( | |
""" | |
<h3 style='text-align: center'> | |
YOLOv10: Real-Time End-to-End Object Detection | |
<a href='https://arxiv.org/abs/2405.14458' target='_blank'>arXiv</a> | <a href='https://github.com/THU-MIG/yolov10' target='_blank'>github</a> | |
</h3> | |
""") | |
with gr.Row(): | |
with gr.Column(): | |
app() | |
if __name__ == '' | |
gradio_app.launch(debug=True) |