我们常常将模型转成onnx通用形式,但有时候也需要测试看一下转成onnx的模型有没有啥问题,我下面写了运行onnx的代码,包括:图像输入、图像resize、BGR转RGB、维度调整、归一化、读取onnx文件并运行、为结果赋予 RGB三彩色、原图与特征图叠加显示、保存输出结果等步骤。
代码里面也有详细注释,很好理解。
路沿检测本质就是语义分割,所以这个代码也适用于语义分割,大家可以尝试改一改,有多少类别,ST_COLORMAP 里面就写多少。我这里就两个类:路和背景,ST_COLORMAP就写了两个。
import os
import onnxruntime as ort
import numpy as np
import cv2
classes = ('road', 'background')
ST_COLORMAP = [[0, 255, 0], [0, 0, 0]]
#[0, 255, 0]代表road的BGR三通道值,[0, 0, 0]代表others的BGR三通道值。
#
def Index2Color(pred):
colormap = np.asarray(ST_COLORMAP, dtype='uint8')
x = np.asarray(pred, dtype='int32')
return colormap[x, :]
def preprocess(img_data):
#归一化
mean_vec = np.array([123.675, 116.28, 103.53])
stddev_vec = np.array([58.395, 57.12, 57.375])
norm_img_data = np.zeros(img_data.shape).astype('float32')
# im = (norm_img_data - mean_vec) / stddev_vec
for i in range(img_data.shape[0]):
norm_img_data[i, :, :] = (img_data[i, :, :] - mean_vec[i]) / stddev_vec[i]
return norm_img_data
output_path = r'D:\Semantic_segmentation\output/'
input_path = r"D:\Semantic_segmentation\images/"
onnx_path = r'D:\Semantic_segmentation\model.onnx'
dir_list = os.listdir(input_path)
for i in dir_list:
#读取图像
img = cv2.imread(input_path+i)
h, w, _ = img.shape
#图像resize
img_resize = cv2.resize(img, (640, 320), interpolation=cv2.INTER_AREA)
#训练前将BGR转RGB
img_rgb = cv2.cvtColor(img_resize, cv2.COLOR_BGR2RGB)
#[320, 640, 3]变成[3, 320, 640]
input_data = np.transpose(img_rgb, (2, 0, 1))
#归一化
input_data = preprocess(input_data)
#调整图像输入
input_data = input_data.reshape([1, 3, 320, 640])
#读取onnx文件,测试运行
sess = ort.InferenceSession(onnx_path)
input_name = sess.get_inputs()[0].name
result = sess.run([], {input_name: input_data})
#为road和others赋予rgb三彩色
result_Color = Index2Color(result[0].squeeze())
#resize到原图大小
result_resize = cv2.resize(result_Color, (w, h), interpolation=cv2.INTER_AREA)
#原图与特征图叠加显示,可自行修改比例
result_out = cv2.addWeighted(img, 1, result_resize, 0.3, 0)
#保存输出结果
cv2.imwrite(output_path+i, result_out)