PaddleSeg模型的ONNX格式部署

近期在工作时遇到项目需要用ONNX格式实现PaddleSeg模型的本地部署，在此记录供日后自己可以参考。

环境：Python3.8+RTX3060 + CUDA 11.4 + onnxruntime-GPU 1.11.0

在贴码之前，PaddleSeg转ONNX模型的输出结果是一个三维 int64数组，分别对应N, H, W,  代表每一个像素点的分类，比如我的模型只有1个分类——缺陷，那么输出的数组中，0代表这个像素点是背景，1代表是缺陷。另外注意的是，如下图维度-1代表输入ONNX的尺寸可以是任意，但在使用时还是要resize成原模型里的尺寸，否则会报错。

 

读取需要用到包，以及获取输入、输出的名字。

import numpy as np
import onnxruntime
import cv2
import time

model_path = 'output.onnx'
image_path = 'test_dp.jpg'

sess = onnxruntime.InferenceSession(model_path,providers=['CUDAExecutionProvider']) #使用GPU

input_name = sess.get_inputs()[0].name  # 'data'
outputs0 = sess.get_outputs()[0].name

图片的读取，resize与转换为RGB格式；

img = cv2.imread(image_path)
img = cv2.resize(img,(512,512))
img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)

因为opencv读取图片的维度为hwc, 按照模型所需要的先转换为nchw，并对其进行normalize (mean=0.5, std=0.5)，才能喂到ONNX模型中；

img = np.transpose(img,(2,0,1))  # 转为chw
input_blob = ((np.expand_dims(img,axis=0).astype(np.float32)/255.0)-0.5)/0.5 #转为nchw并归一化

运行模型，打印运算时间为0.168s，outputs0, input_name对应为自己模型的输入输出名称；

time0 = time.time()
out = sess.run([outputs0], input_feed={input_name: new_img})
time1 = time.time()
print(time1-time0)

创建空白图像、空白数组用于记录结果并显示；

res_img = np.ones((512,512,3),dtype=np.uint8) #空白图像
tmp =out[0][0]                                #输出数组

#将输出中分类为缺陷的像素点标记出来
for i in range(0,512):
    for j in range(0,512):
        if tmp[i][j] == 1:
            res_img[i][j][0] =88
            res_img[i][j][1] = 197
            res_img[i][j][2] =33

记录下来的输出结果；

将该结果与原图叠加，并显示：

 

res = cv2.addWeighted(img,0.7,res_img,0.3,0)
res_img = cv2.resize(res,(1024,800))
cv2.imshow('?',res_img)
cv2.waitKey(0)

最后是叠加出来的最后结果，检测效果还是比较理想：