|
|
import gradio as gr |
|
|
import spaces |
|
|
from PIL import Image |
|
|
from ultralytics import YOLO |
|
|
|
|
|
|
|
|
models = { |
|
|
"yolov10n": YOLO("jameslahm/yolov10n"), |
|
|
"yolov10s": YOLO("jameslahm/yolov10s"), |
|
|
"yolov10m": YOLO("jameslahm/yolov10m"), |
|
|
"yolov10b": YOLO("jameslahm/yolov10b"), |
|
|
"yolov10l": YOLO("jameslahm/yolov10l"), |
|
|
"yolov10x": YOLO("jameslahm/yolov10x"), |
|
|
} |
|
|
|
|
|
@spaces.GPU(duration=30) |
|
|
def yolov10_inference(image, model_id, image_size, conf_threshold, iou_threshold): |
|
|
model = models[model_id] |
|
|
results = model.predict( |
|
|
source=image, |
|
|
imgsz=image_size, |
|
|
conf=conf_threshold, |
|
|
iou=iou_threshold, |
|
|
) |
|
|
annotated_image = results[0].plot() |
|
|
return Image.fromarray(annotated_image[..., ::-1]) |
|
|
|
|
|
def app(): |
|
|
with gr.Blocks() as demo: |
|
|
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.0, |
|
|
maximum=1.0, |
|
|
step=0.05, |
|
|
value=0.25, |
|
|
) |
|
|
iou_threshold = gr.Slider( |
|
|
label="IoU Threshold", |
|
|
minimum=0.0, |
|
|
maximum=1.0, |
|
|
step=0.05, |
|
|
value=0.45, |
|
|
) |
|
|
yolov10_infer = gr.Button(value="Detect Objects") |
|
|
|
|
|
with gr.Column(): |
|
|
output_image = gr.Image(type="pil", label="Annotated Image") |
|
|
|
|
|
yolov10_infer.click( |
|
|
fn=yolov10_inference, |
|
|
inputs=[image, model_id, image_size, conf_threshold, iou_threshold], |
|
|
outputs=[output_image], |
|
|
) |
|
|
|
|
|
gr.Examples( |
|
|
examples=["Rocket.png"], |
|
|
inputs=[image], |
|
|
) |
|
|
return demo |
|
|
|
|
|
if __name__ == "__main__": |
|
|
app().launch() |
