Python wrapper for running YOLOv11 object detection using an ONNX model with ONNXRuntime, Including optimized post processing with class targeted NMS.
pip install -r requirements.txtIf you have GPU, uncomment onnxruntime in requirements.txt and uncomment onnxruntime-gpu line
from yolo11 import YOLOv11
import cv2
from pathlib import Path
model = YOLOv11("yolo11s.onnx")
test_img = Path("test_imgs/sample.jpg")
img = cv2.imread(str(test_img))
bbox_list = model.detect(img)
for bbox in bbox_list:
img = bbox.draw(img)
cv2.imwrite(str(test_img.parent / test_img.stem) + "_out.jpg", img)You can take a look at test_vid.py.
To run it, use
python3 test_vid.py <input_video.mp4> <--save>--save will save the output video.
You can pass a function or callable to filter valid classes, making NMS slightly efficient. Example
valid_class_d = {
1: "bicycle",
2: "car",
3: "motorcycle",
4: "airplane",
5: "bus",
6: "train",
7: "truck",
8: "boat",
}
model = YOLOv11(
model_path="yolo11s.onnx",
valid_class_checker=lambda lbl_id, _: 1 <= lbl_id <= 8 # detect only vehicles
# (or)
# valid_class_checker=lambda lbl_id, lbl: lbl_id in valid_class_d and lbl == valid_class_d[lbl_id]
)To get all {class_id, name} pairs
print("\n".join([ f"{k}: {v}" for k,v in model.get_class_id_name_pairs().items() ]))