onnxruntime推理

今天转换了一个pytorch模型到onnx模型。然后想要测试一下onnx模型是不是可以正常使用。
然后就发现自己几乎没有接触过onnx推理的python代码。
所以就整理了一下，还是参考了一下自己之前的博客onnx模型转换及推理

使用的是一个图像分割的模型转的onnx模型。
然后发现onnx真的支持的算子很少，连一个sigmoid函数都没有找到，看到后处理中还要用到sigmoid函数，这个虽然很简单，居然还要自己实现。

import os
import onnx
from onnx import optimizer
from onnx import helper
import onnxruntime
import cv2
import numpy as np


def  sigmoid(x):
    return 1.0 / (1.0 + np.exp(-x))

# Load the ONNX model
# model = onnx.load("BiSeNet.onnx")
#
# # 检查模型格式是否完整及正确
# onnx.checker.check_model(model)
# # 获取输出层，包含层名称、维度信息
# input = model.graph.input
# print(input)
# output = model.graph.output
# print(output)
# # step1: 数据预处理
# in_size = 512
# norm_mean = (0.5, 0.5, 0.5)  # 比imagenet的mean效果好
# norm_std = (0.5, 0.5, 0.5)

# 创建一个InferenceSession的实例，并将模型的地址传递给该实例
sess = onnxruntime.InferenceSession('BiSeNet.onnx')
# 调用实例sess的润方法进行推理
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name

# step1: 数据预处理
in_size = 512
norm_mean = (0.5, 0.5, 0.5)  # 比imagenet的mean效果好
norm_std = (0.5, 0.5, 0.5)
# transform = A.Compose([
#     A.Resize(width=in_size, height=in_size),
#     A.Normalize(norm_mean, norm_std),
# ])
path_img = "test.jpg"
img_bgr = cv2.imdecode(np.fromfile(path_img), cv2.IMREAD_UNCHANGED)
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)  # 注意 opencv默认bgr，需要转换为rgb
img_rgb = cv2.resize(img_rgb, (in_size, in_size))
img_rgb = img_rgb.astype(np.float32)/255.0
img_rgb = (img_rgb - norm_mean) / norm_std
# transformed = transform(image=img_rgb, mask=img_rgb)  # 预处理
# img_rgb = transformed['image']  # 取出图片
img_in = np.transpose(img_rgb, (2, 0, 1)).astype(np.float32)
img_in = np.expand_dims(img_in, axis=0)


outputs = sess.run([output_name], {input_name: img_in})

print("process over! ")

pre_label = np.squeeze(outputs[0])
pre_label = sigmoid(pre_label)
# step4：后处理显示图像
background = np.zeros_like(img_bgr, dtype=np.uint8)
background[:] = 255
alpha_bgr = pre_label
alpha_bgr = cv2.cvtColor(alpha_bgr, cv2.COLOR_GRAY2BGR)
h, w, c = img_bgr.shape
alpha_bgr = cv2.resize(alpha_bgr, (w, h))
# fusion
result = np.uint8(img_bgr * alpha_bgr + background * (1 - alpha_bgr))
out_img = np.concatenate([img_bgr, result], axis=1)
result_root_dir = "./test"
path_result = os.path.join(result_root_dir, path_img)
cv2.imwrite(path_result, out_img)