《------往期经典推荐------》
二、机器学习实战专栏【链接】，已更新31期，欢迎关注，持续更新中~~
三、深度学习【Pytorch】专栏【链接】
四、【Stable Diffusion绘画系列】专栏【链接】
五、YOLOv8改进专栏【链接】，持续更新中~~
六、YOLO性能对比专栏【链接】，持续更新中~
《------正文------》
引言
本文将介绍如何使用 onnxruntime 部署YOLOv11模型 。
ONNX 是一个跨平台引擎，用于以 ONNX 格式运行 ML 模型。要将预训练或自定义的 YOLOv11 模型转换为 ONNX 格式，我们将使用 Ultralytics 库，该库简化了此过程，使我们能够仅用几行代码导出模型。
from ultralytics import YOLO
model = YOLO("path/to/your/model.pt") # path to your pretrained or custome YOLO11 model
model.export(format="onnx")
转换成功完成后，我们只需要三个库：ONNXTM，OpenCV 和 Numpy， ONNX 模型和一个 python 脚本来执行推理。
安装库与模型推理
首先我们安装一下需要的库：
pip install opencv-python onnxruntime
接下来是用于推理的脚本，我们首先导入所需的模块，使用 ONNX SDK 加载 YOLOv11 ONNX 模型，并定义模型可以检测到的所有类的列表。因为我使用的是官方的 YOLOv11，它是在包含 80 个类的 COCO 数据集上训练的。如果您使用自己的训练模型，那么类在您的场景中会有所不同。
import onnxruntime as ort
import cv2
import numpy as np
mode_path = "YOLOvn.onnx"
onnx_model = ort.InferenceSession(mode_path)
with open('coco-classes.txt') as file: # loading the classes names from the file.
content = file.read()
classes = content.split('\n')
图片预处理与说明
在运行推理之前，我们将加载图像并对其进行一些必要的图像处理。
image = cv2.imread("images/img1.jpg")
img_w, img_h = image.shape[1], image.shape[0]
img = cv2.resize(image, (640, 640))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = img.transpose(2, 0, 1)
img = img.reshape(1, 3, 640, 640)
在上面的代码片段中，我将图像大小调整为 640x640 像素，因为模型是在这个大小的图像上训练的。另外，我们的 YOLOv11 ONNX 模型期望 RGB 图像的形状为（1，3，640，640），但 OpenCV 默认读取 BGR 格式的图像，通道信息位于第 2 位而不是第 0 位。因此，我们将图像转换为 RGB，并将其转换为所需的输入形状（1，3，640，640）。在这里，
1表示批量大小（我们将一次为模型提供一个图像）。3表示颜色通道（RGB）。640、640是图像的空间维度。
接下来，我们将对像素值进行归一化，并将其范围从[0，255]调整到[0，1]。然后，图像被转换为 float32 以匹配模型的预期输入数据类型。
# Normalize pixel values to the range [0, 1]
img = img / 255.0
# Convert image to float32
img = img.astype(np.float32)
模型推理与后处理
图片预处理完成后，我们使用模型进行图片推理。
outputs = onnx_model.run(None, {"images": img})
在推理结束时，我们得到形状矩阵（1，84，8400）作为输出，代表8400 个检测框，每个检测具有 84 个参数。这是因为官方的 YOLOv11 模型被设计为总是预测图像中的 8400 个对象，而不管实际存在多少对象。我们将删除那些具有低置信度分数的检测。这里，矩阵形状的 84 表示用于每次检测的参数的数量。这包括边界框坐标（x1，y1，x2，y2）和模型训练的 80 个不同类的置信度得分。请注意，对于自定义模型，此结构可能会有所不同。置信度得分的数量始终取决于模型训练的类别数量。例如，如果您训练 YOLOv11 对象检测模型检测 1 个类别，则将有 5 个参数而不是 84 个。前四个将再次是边界框坐标，最后一个将是置信度得分。
现在，为了简单起见，重新整形这个输出矩阵以获得（8400，84）的形状。
results = out[0]
results = results.transpose()
现在，我们的下一步是为每个对象确定最可能的类，并过滤掉低置信度的预测。我们可以通过为每个检测选择具有最高置信度得分的类来做到这一点。此外，我们丢弃所有置信度得分低于选定阈值（在我们的情况下为0.5）的检测。
def filter_Detections(results, thresh = 0.5):
# if model is trained on 1 class only
if len(results[0]) == 5:
# filter out the detections with confidence > thresh
considerable_detections = [detection for detection in results if detection[4] > thresh]
considerable_detections = np.array(considerable_detections)
return considerable_detections
# if model is trained on multiple classes
else:
A = []
for detection in results:
class_id = detection[4:].argmax()
confidence_score = detection[4:].max()
new_detection = np.append(detection[:4],[class_id,confidence_score])
A.append(new_detection)
A = np.array(A)
# filter out the detections with confidence > thresh
considerable_detections = [detection for detection in A if detection[-1] > thresh]
considerable_detections = np.array(considerable_detections)
return considerable_detections
results = filter_Detections(results)
我们打印一下过滤后的结果形状：
print(results.shape)
(22, 6)
上述结果表明，现在剩下 22 个检测，每个检测具有 6 个参数。它们是边界框左上角（x1，y1）和右下角（x2，y2）坐标、类 ID 和置信度值。
此时，仍然可能有一些不必要的检测，这是因为其中一些实际上指向同一个对象。这可以通过使用称为非最大抑制（NMS）的算法来解决，该算法在可能涉及同一对象的检测中选择最佳检测。它通过考虑两个关键指标来实现这一点，这两个指标是置信度值和交集（IOU）。此外，我们需要将剩余的检测重新调整回原始规模。这是因为我们的模型已经输出了对大小为 640x640 的图像的检测，这不是我们原始图像的大小。
def NMS(boxes, conf_scores, iou_thresh = 0.55):
# boxes [[x1,y1, x2,y2], [x1,y1, x2,y2], ...]
x1 = boxes[:,0]
y1 = boxes[:,1]
x2 = boxes[:,2]
y2 = boxes[:,3]
areas = (x2-x1)*(y2-y1)
order = conf_scores.argsort()
keep = []
keep_confidences = []
while len(order) > 0:
idx = order[-1]
A = boxes[idx]
conf = conf_scores[idx]
order = order[:-1]
xx1 = np.take(x1, indices= order)
yy1 = np.take(y1, indices= order)
xx2 = np.take(x2, indices= order)
yy2 = np.take(y2, indices= order)
keep.append(A)
keep_confidences.append(conf)
# iou = inter/union
xx1 = np.maximum(x1[idx], xx1)
yy1 = np.maximum(y1[idx], yy1)
xx2 = np.minimum(x2[idx], xx2)
yy2 = np.minimum(y2[idx], yy2)
w = np.maximum(xx2-xx1, 0)
h = np.maximum(yy2-yy1, 0)
intersection = w*h
# union = areaA + other_areas - intesection
other_areas = np.take(areas, indices= order)
union = areas[idx] + other_areas - intersection
iou = intersection/union
boleans = iou < iou_thresh
order = order[boleans]
# order = [2,0,1] boleans = [True, False, True]
# order = [2,1]
return keep, keep_confidences
# function to rescale bounding boxes
def rescale_back(results,img_w,img_h):
cx, cy, w, h, class_id, confidence = results[:,0], results[:,1], results[:,2], results[:,3], results[:,4], results[:,-1]
cx = cx/640.0 * img_w
cy = cy/640.0 * img_h
w = w/640.0 * img_w
h = h/640.0 * img_h
x1 = cx - w/2
y1 = cy - h/2
x2 = cx + w/2
y2 = cy + h/2
boxes = np.column_stack((x1, y1, x2, y2, class_id))
keep, keep_confidences = NMS(boxes,confidence)
print(np.array(keep).shape)
return keep, keep_confidences
rescaled_results, confidences = rescale_back(results, img_width, img_height)
这里 rescaled_results 包含边界框（x1，y1，x2，y2）和类 id，而 **confidences **存储相应的置信度分数。
最后，我们可以将这些结果可视化在我们的图像上。
for res, conf in zip(rescaled_results, confidences):
x1,y1,x2,y2, cls_id = res
cls_id = int(cls_id)
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
conf = "{:.2f}".format(conf)
# draw the bounding boxes
cv2.rectangle(image,(int(x1),int(y1)),(int(x2),int(y2)),(255,0, 0),1)
cv2.putText(image,classes[cls_id]+' '+conf,(x1,y1-17),
cv2.FONT_HERSHEY_SIMPLEX,0.7,(255,0,0),1)
cv2.imwrite("Output.jpg", image)
好了，这篇文章就介绍到这里，喜欢的小伙伴感谢给点个赞和关注，更多精彩内容持续更新~~
关于本篇文章大家有任何建议或意见，欢迎在评论区留言交流！
扫一扫
2238
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
