0. 为什么选着OpenCV进行推理
按道理，使用OpenVINO推理套件（引擎）进行优化加速推理，或者其他推理套件进行推理，比如TensorRT等等。OpenVINO是英特尔推出的一款全面的工具套件，用于快速部署应用和解决方案，支持计算机视觉的CNN网络结构超过150余种。
但引入推理套件（引擎）势必带来额外的工作量，在对推理速度要求不高的情况下，使用OpenCV自带的DNN模块已经足够满足需求了，且能做到轻量化部署，减少对第三方平台的依赖。
随着版本的更迭，OpenCV的版本当前已经来到了4.5.4，其DNN模块是越来完善了。
1. 前提
# python-opencv版本：4.5.3
# 已经完成分类网络的训练，并将pytorch网络模型转换/保存为onnx文件格式。
# onnx文件 onnx_name = "surface_defect_model.onnx"
至于如何保存为onnx,本人做了记录。
pytorch保存、加载模型, 并将网络模型.pt保存为ONNIX
该网络对于输入的要求为：
dummy = [1, 3, 300, 300] # 按照N C W H 的格式排列
图片要求为最常见的三通道的8UC3格式。
# 网络模型简介：
该模型为图像分类神经网络，它是基于ResNet18进行的迁移训练，训练集为6种不同类型的缺陷图片。
# 图片类别名称
defect_labels = ["In","Sc","Cr","PS","RS","Pa"]由于是在pytorch中进行的迁移学习训练，所以训练数据集都严格按照：
transforms.Compose([transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
transforms.Resize((300, 300))
])进行Normalize预处理的，所以在推理时，也要先将图像严格按照同样的预处理方式处理完毕之后再喂入网络进行推理，否则会产生事与愿违的预测效果。
2. 代码实战
2.1 cv.dnn.blobFromImage不适用
cv.dnn.blobFromImage原型如下：
所以本人自己实现了一个把三通道的CV_8UC3图像转为pytorch中transforms.ToTensor()以及transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])的效果，返回的也是网络需要的N*C*W*H（N=1）形状的数据，可以直接喂入网络进行推理。BlobFromNormalizeRGBImage函数并未更改传入图像的尺寸，其实现如下：
def BlobFromNormalizeRGBImage(img, meanList, stdList):
"""
BlobFromNormalizeRGBImage(img, meanList, stdList)
. @generate a N*C*W*H blob from a RGB(CV_8UC3) picture
. * @param img : a RGB(CV_8UC3) picture ,whose pixel value arrange is 0 ~ 255.
. * @param meanList： the mean value list of RGB channels,it must contains three elements.
. when it is None,ignore it.
. * @param stdList： the standard deviation(sd) value list of RGB channels,it must contains three elements.
. when it is None,ignore it.
. * @returns a N*C*W*H blob .
"""
img = img / 255.0 # 归一化至0 ~ 1区间
R, G, B = cv.split(img)
if meanList is not None:
R = R - meanList[0]
G= G - meanList[1]
B = B - meanList[2]
if stdList is not None:
R = R / stdList[0]
G = G / stdList[1]
B = B / stdList[2]
# 通道合并
merged = cv.merge([R, G, B])
# print("merged.shape:",merged.shape) # merged.shape: (300, 300, 3)
merged = merged.transpose((2, 0, 1))
# print("transpose merged.shape:", merged.shape) # transpose merged.shape: (3, 300, 300)
blob = np.expand_dims(merged, 0)
#生成(1, 3, 300, 300)的格式
# print("blob.shape:", blob.shape) # blob.shape: (1, 3, 300, 300)
return blob2.2 完整代码
# import torch # 脱离pytorch框架，不需要引入
import cv2 as cv
import numpy as np
import os
# 类别
defect_labels = ["In","Sc","Cr","PS","RS","Pa"]
# onnx文件
onnx_name = "surface_defect_model.onnx"
print(cv.__version__)
def BlobFromNormalizeRGBImage(img, meanList, stdList):
"""
BlobFromNormalizeRGBImage(img, meanList, stdList)
. @generate a N*C*W*H blob from a RGB(CV_8UC3) picture
. * @param img : a RGB(CV_8UC3) picture ,whose pixel value arrange is 0 ~ 255.
. * @param meanList： the mean value list of RGB channels,it must contains three elements.
. when it is None,ignore it.
. * @param stdList： the standard deviation(sd) value list of RGB channels,it must contains three elements.
. when it is None,ignore it.
. * @returns a N*C*W*H blob .
"""
img = img / 255.0 # 归一化至0 ~ 1区间
R, G, B = cv.split(img)
if meanList is not None:
R = R - meanList[0]
G= G - meanList[1]
B = B - meanList[2]
if stdList is not None:
R = R / stdList[0]
G = G / stdList[1]
B = B / stdList[2]
# 通道合并
merged = cv.merge([R, G, B])
# print("merged.shape:",merged.shape) # merged.shape: (300, 300, 3)
merged = merged.transpose((2, 0, 1))
# print("transpose merged.shape:", merged.shape) # transpose merged.shape: (3, 300, 300)
blob = np.expand_dims(merged, 0)
#生成(1, 3, 300, 300)的格式
# print("blob.shape:", blob.shape) # blob.shape: (1, 3, 300, 300)
return blob
def opencv_onnx_defect_infer():
#前向推理的输入大小
dummy = [1, 3, 300, 300] # 按照N C W H 的格式排列
width = dummy[2]
height = dummy[3]
# opencv dnn加载
net = cv.dnn.readNetFromONNX(onnx_name)
root_dir = "./test"
#总数
total_pic_cnt = 0
# 预测正确的数量
predict_ok_cnt = 0
fileNames = os.listdir(root_dir)
for f in fileNames:
print("=========================================================")
image = cv.imread(os.path.join(root_dir, f))
# print("原始的尺寸:", image.shape)
image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
image = cv.resize(image, (width, height))
# print("resize之后的尺寸:",image.shape)
# blob = cv.dnn.blobFromImage(image,4.424778,size=None,mean=[0.485, 0.456, 0.406])
blob = BlobFromNormalizeRGBImage(image, meanList=[0.485, 0.456, 0.406],stdList=[0.229, 0.224, 0.225])
# print("BlobFromNormalizeRGBImage:",blob.shape) # (1, 3, 300, 300)
# 设置模型的输入
net.setInput(blob)
out = net.forward()
# print("out.shape:",out.shape)# out.shape: (1, 6)
# Get a class with a highest score.
out = out.flatten()
# print("out. flatten shape:", out.shape)# out. flatten shape: (6,)
classId = np.argmax(out)
# print("classId:", classId)
confidence = out[classId]
# print("confidence:", confidence)
# print("class:", defect_labels[classId])
total_pic_cnt += 1
if f.find(defect_labels[classId]) >= 0:
predict_ok_cnt += 1
print("OK !!!!")
print("real class %s, infer class:%s, confidence:%f, classId:%d" % (f, defect_labels[classId], confidence,classId))
print(out)
else:
print("predict error !!!!")
print("real: %s, predict result:%s, confidence:%f, classId:%d, predict error" % (f, defect_labels[classId],confidence, classId) )
print(out)
#统计结果
print("total_pic_cnt:%d, predict_ok_cnt:%d, rate = %0.3f" % (total_pic_cnt, predict_ok_cnt, predict_ok_cnt / total_pic_cnt))
if __name__ == "__main__":
opencv_onnx_defect_infer()