|
import spaces |
|
import gradio as gr |
|
import cv2 |
|
import tempfile |
|
from ultralytics import YOLOv10 |
|
|
|
model = YOLOv10.from_pretrained(f'jameslahm/yolov10s') |
|
|
|
@spaces.GPU |
|
def yolov10_inference(image, conf_threshold): |
|
width, _ = image.size |
|
results = model.predict(source=image, imgsz=width, conf=conf_threshold) |
|
annotated_image = results[0].plot() |
|
return annotated_image[:, :, ::-1] |
|
|
|
|
|
def app(): |
|
with gr.Blocks(): |
|
with gr.Row(): |
|
with gr.Column(): |
|
image = gr.Image(type="pil", label="Image", visible=True) |
|
conf_threshold = gr.Slider( |
|
label="Confidence Threshold", |
|
minimum=0.0, |
|
maximum=1.0, |
|
step=0.05, |
|
value=0.25, |
|
) |
|
|
|
with gr.Column(): |
|
output_image = gr.Image(type="numpy", label="Annotated Image", visible=True) |
|
|
|
image.stream( |
|
fn=yolov10_inference, |
|
inputs=[image, conf_threshold], |
|
outputs=[image], |
|
stream_every=0.1, |
|
time_limit=30 |
|
) |
|
|
|
|
|
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'> |
|
<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__ == '__main__': |
|
gradio_app.launch() |
|
|