File size: 4,762 Bytes
2d2959e 1671fee 2d2959e 1671fee 2d2959e 1671fee 2d2959e 1671fee 2d2959e 1671fee 2d2959e 1671fee 2d2959e 23d39e0 2d2959e 1671fee 2d2959e 1671fee 2d2959e 1671fee 2d2959e 1671fee 2d2959e 1671fee 2d2959e 1671fee 2d2959e 1671fee c610138 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | import spaces
import supervision as sv
import PIL.Image as Image
from ultralytics import YOLO
import gradio as gr
import torch
# Model filenames for YOLO models
model_filenames = [
"yolo11n.pt",
"yolo11s.pt",
"yolo11m.pt",
"yolo11l.pt",
"yolo11x.pt"
]
# Box annotator for drawing detection boxes
box_annotator = sv.BoxAnnotator()
# Category dictionary mapping class IDs to labels
category_dict = {
0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus',
6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light', 10: 'fire hydrant',
11: 'stop sign', 12: 'parking meter', 13: 'bench', 14: 'bird', 15: 'cat',
16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 20: 'elephant', 21: 'bear',
22: 'zebra', 23: 'giraffe', 24: 'backpack', 25: 'umbrella', 26: 'handbag',
27: 'tie', 28: 'suitcase', 29: 'frisbee', 30: 'skis', 31: 'snowboard',
32: 'sports ball', 33: 'kite', 34: 'baseball bat', 35: 'baseball glove',
36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle',
40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl',
46: 'banana', 47: 'apple', 48: 'sandwich', 49: 'orange', 50: 'broccoli',
51: 'carrot', 52: 'hot dog', 53: 'pizza', 54: 'donut', 55: 'cake',
56: 'chair', 57: 'couch', 58: 'potted plant', 59: 'bed', 60: 'dining table',
61: 'toilet', 62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote', 66: 'keyboard',
67: 'cell phone', 68: 'microwave', 69: 'oven', 70: 'toaster', 71: 'sink',
72: 'refrigerator', 73: 'book', 74: 'clock', 75: 'vase', 76: 'scissors',
77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'
}
# Function for running YOLO inference
@spaces.GPU
def yolo_inference(image, model_id, conf_threshold, iou_threshold, max_detection):
model = YOLO(model_id)
results = model(source=image, imgsz=640, iou=iou_threshold, conf=conf_threshold, verbose=False, max_det=max_detection)[0]
detections = sv.Detections.from_ultralytics(results)
labels = [
f"{category_dict[class_id]} {confidence:.2f}"
for class_id, confidence in zip(detections.class_id, detections.confidence)
]
annotated_image = box_annotator.annotate(image, detections=detections, labels=labels)
return annotated_image
# Gradio app function
def app():
with gr.Blocks():
with gr.Row():
with gr.Column():
image = gr.Image(type="pil", label="Image", interactive=True)
model_id = gr.Dropdown(
label="Model",
choices=model_filenames,
value=model_filenames[0] if model_filenames else "",
)
conf_threshold = gr.Slider(
label="Confidence Threshold",
minimum=0.1,
maximum=1.0,
step=0.1,
value=0.25,
)
iou_threshold = gr.Slider(
label="IoU Threshold",
minimum=0.1,
maximum=1.0,
step=0.1,
value=0.45,
)
max_detection = gr.Slider(
label="Max Detection",
minimum=1,
maximum=300,
step=1,
value=300,
)
yolov_infer = gr.Button(value="Detect Objects")
with gr.Column():
output_image = gr.Image(type="pil", label="Annotated Image", interactive=False)
yolov_infer.click(
fn=yolo_inference,
inputs=[image, model_id, conf_threshold, iou_threshold, max_detection],
outputs=[output_image],
)
# Move gr.Examples inside the Blocks context
gr.Examples(
examples=[
["zidane.jpg", "yolo11s.pt", 0.25, 0.45, 300],
["bus.jpg", "yolo11m.pt", 0.25, 0.45, 300],
["yolo_vision.jpg", "yolo11x.pt", 0.25, 0.45, 300],
],
fn=yolo_inference,
inputs=[image, model_id, conf_threshold, iou_threshold, max_detection],
outputs=[output_image],
cache_examples=True,
)
# Create Gradio app
gradio_app = gr.Blocks()
with gradio_app:
gr.HTML(
"""
<h1 style='text-align: center'>
Yolo11: Object Detection
</h1>
"""
)
gr.HTML(
"""
<p style='text-align: center'>
Latest ultralytics yolo11 object detection models. Upload an image to run inference.
</p>
"""
)
with gr.Row():
with gr.Column():
app()
# Launch the Gradio app with a specific port
gradio_app.launch(server_port=7861) |