使用netron工具查看onnx网络结构
如下图,可以看出此次要使用的网络输入为:
网络的输出为:
准备输入数据
将要进行推理的图像数据的数据类型和形状转换为需要的数据类型和需要的输入形状
# 1.准备输入数据
img = cv2.imread("../../0.samples/0.png") # 读取图像
img_data = cv2.resize(img, (256, 256)) # 对图像进行大小变换
img_data = img_data.swapaxes(1, 2).swapaxes(0, 1) # 将维度为[256, 256, 3]转换为[3, 256, 256]
img_data = img_data[np.newaxis, :, :, :] # 对图像进行维度变换,使其满足网络的输入要求
img_data = np.array(img_data) # 将图像转换为nparray
img_data = img_data/255. # 对图像数据进行归一化
img_data = img_data.astype(np.float32) # 转换图像数据类型,使其满足网络的输入要求
创建一个inference实例,把模型传递给inference实例
# 2.创建一个InferenceSession的实例,并将模型的地址传递给该实例
sess = onnxruntime.InferenceSession('../../0.onnx/uNet_32.onnx', providers=['CUDAExecutionProvider'])
调用inference实例,获取输入输出名称
# 3.调用InferenceSession的实例,获取输入输出名称
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name
调用InferenceSession实例的run方法进行推理,根据输出名称获取推理结果
# 4.调用InferenceSession实例的run方法进行推理,根据输出名称获取推理结果
outputs = sess.run([output_name], {input_name: img_data})
对inference实例输出的结果进行处理,获得所需要的结果
# 5.推理结果处理
predict_imgs = np.squeeze(outputs[0]) # 推理结果转换为nparray
predict_imgs = predict_imgs * 255 # 结果数值范围转换
predict_imgs = np.clip(predict_imgs, 0, 255) # 结果数值范围转换
predict_imgs = cv2.resize(predict_imgs, (1920, 1080)) # 结果图像大小转换
predict_imgs = predict_imgs.astype('uint8') # 结果数据类型转换
python版API推理流程
import cv2
import onnxruntime
import numpy as np
# 1.准备输入数据
img = cv2.imread("../../0.samples/0.png") # 读取图像
img_data = cv2.resize(img, (256, 256)) # 对图像进行大小变换
img_data = img_data[np.newaxis, :, :, :] # 对图像进行维度变换,使其满足网络的输入要求
img_data = np.array(img_data) # 将图像转换为nparray
img_data = img_data/255. # 对图像数据进行归一化
img_data = img_data.astype(np.float32) # 转换图像数据类型,使其满足网络的输入要求
# 2.创建一个InferenceSession的实例,并将模型的地址传递给该实例
sess = onnxruntime.InferenceSession('../../0.onnx/uNet_32.onnx')
# 3.调用InferenceSession的实例,获取输入输出名称
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name
# 4.调用InferenceSession实例的run方法进行推理,根据输出名称获取推理结果
outputs = sess.run([output_name], {input_name: img_data})
# 5.推理结果处理
predict_imgs = np.squeeze(outputs[0]) # 推理结果转换为nparray
predict_imgs = predict_imgs * 255 # 结果数值范围转换
predict_imgs = np.clip(predict_imgs, 0, 255) # 结果数值范围转换
predict_imgs = cv2.resize(predict_imgs, (1920, 1080)) # 结果图像大小转换
predict_imgs = predict_imgs.astype('uint8') # 结果数据类型转换
# 查看输出结果
cv2.imshow("1", predict_imgs)
cv2.waitKey(0)