预测.onnx的路沿检测（语义分割）模型的结果可视化代码-CSDN博客

本文链接：https://blog.csdn.net/m0_46416211/article/details/135264920

我们常常将模型转成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)