使用OpenCV DNN模块部署YOLOv5中导出的ONNX模型

要使用OpenCV DNN模块部署ONNX模型,你可以按照以下步骤实现代码:

1. 导入必要的库:

import cv2
import numpy as np

2. 加载模型和类别标签:

model_path = 'path_to_model.onnx'  # 替换为ONNX模型文件的路径
class_labels = 'path_to_class_labels.txt'  # 替换为与模型相对应的类别标签文件路径

# 加载类别标签
classes = []
with open(class_labels, 'r') as f:
    classes = [line.strip() for line in f.readlines()]

# 加载模型
net = cv2.dnn.readNetFromONNX(model_path)

3. 加载图像并进行推断 :

image_path = 'path_to_image.jpg'  # 替换为你要测试的图像路径
 
# 加载图像
image = cv2.imread(image_path)

# 将图像转换为blob格式
blob = cv2.dnn.blobFromImage(image, size=(224, 224), mean=(0, 0, 0), swapRB=True, crop=False)

# 设置输入blob
net.setInput(blob)

# 进行推断
detections = net.forward()

 4. 解析推断结果:

for i in range(detections.shape[2]):
    confidence = detections[0, 0, i, 2]
    if confidence > confidence_threshold:  # 设置置信度阈值
        class_id = int(detections[0, 0, i, 1])
        label = classes[class_id]
        box = detections[0, 0, i, 3:7] * np.array([image.shape[1], image.shape[0], image.shape[1], image.shape[0]])
        (x, y, w, h) = box.astype("int")

        cv2.rectangle(image, (x, y), (w, h), (0, 255, 0), 2)
        text = f'{label}: {confidence:.2f}'
        cv2.putText(image, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

# 显示结果图像
cv2.imshow("Output", image)
cv2.waitKey(0)
cv2.destroyAllWindows()

        在上述代码中,我们使用置信度阈值来过滤预测的边界框,并绘制边界框和类别标签,并显示结果图像。 

猜你喜欢

转载自blog.csdn.net/qq_39312146/article/details/132074850