|
import spaces |
|
import gradio as gr |
|
import cv2 |
|
import tempfile |
|
from ultralytics import YOLOv10 |
|
|
|
image_processor = RTDetrImageProcessor.from_pretrained("PekingU/rtdetr_r50vd") |
|
model = RTDetrForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd") |
|
|
|
@spaces.GPU |
|
def yolov10_inference(image, conf_threshold): |
|
|
|
inputs = image_processor(images=image, return_tensors="pt") |
|
|
|
with torch.no_grad(): |
|
outputs = model(**inputs) |
|
|
|
results = image_processor.post_process_object_detection(outputs, target_sizes=torch.tensor([image.size[::-1]]), threshold=0.3) |
|
|
|
|
|
|
|
def app(): |
|
with gr.Blocks(): |
|
with gr.Row(): |
|
with gr.Column(): |
|
image = gr.Image(type="pil", label="Image", visible=True, sources="webcam", height=500, width=500) |
|
conf_threshold = gr.Slider( |
|
label="Confidence Threshold", |
|
minimum=0.0, |
|
maximum=1.0, |
|
step=0.05, |
|
value=0.25, |
|
) |
|
image.stream( |
|
fn=yolov10_inference, |
|
inputs=[image, conf_threshold], |
|
outputs=[image], |
|
stream_every=0.2, |
|
time_limit=30 |
|
) |
|
|
|
|
|
gradio_app = gr.Blocks() |
|
with gradio_app: |
|
gr.HTML( |
|
""" |
|
<h1 style='text-align: center'> |
|
YOLOv10 Webcam Stream |
|
</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() |
|
|