接上一篇yolo。onnx runtime是端测AI部署的重要方法。另一种是tensorflow lite。
本文使用python的onnx runtime，跑一下yolo26n-cls这个模型
pip install onnx
pip install onnxruntime
yolo26n-cls.onnx
#!/bin/python3
import cv2
import numpy as np
import onnxruntime as ort
model_path = "yolo26n-cls.onnx"
image_path = "turtle.webp"
height = 224
width = 128
# read image
img = cv2.imread("turtle.webp")
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (128, 224), interpolation=cv2.INTER_LINEAR)
# convert to tensor
tensor = img.astype("float32") / 255
tensor = np.transpose(tensor, (2, 0, 1))
tensor = np.expand_dims(tensor, axis=0)
# inference
session = ort.InferenceSession("yolo26n-cls.onnx")
outputs = session.run(None, {"images": tensor })
# process result
probabilities = outputs[0][0]
top5_indices = np.argsort(probabilities)[-5:][::-1]
print(top5_indices)代码解释
read image：
opencv是BGR序列，需要专程RGB序列.
yolo26n-cls这个模型接受的张量是(1, 3, 224, 128)，对应batch, channel, height, width.
所以先要缩放成128x224
convert to tensor
img是int8，需要转换成float32
然后需要从HWC 转换成CHW
最后要从3 dim换成4 dim，也就是添加一个batch轴
最后进行推理
推理结果是一个1x1000的数组。去概率最高的5个标签
$ ./onnx-1.py
[ 35 36 37 33 113]
使用onnx加载模型，可以查看模型的标签
import onnx
model = onnx.load(model_path)
print(model)
35号标签是
35: \'mud_turtle\'