python实现onnx模型推理

本文以回归任务模型为例演示推理代码。

import onnxruntime
import cv2
import numpy as np


# 定义一个实现回归预测任务的类
class Regression(object):
    def __init__(self):
        onnx_model_path = 'model.onnx'  # onnx推理模型路径
        self.size = (320, 320)  # 输入图片尺寸
        self.model = onnxruntime.InferenceSession(onnx_model_path, provider=['CPUExecutionProvider'])  # 创建InferenceSession对象并指定CPU提供推理程序
        self.input_name = self.model.get_inputs()[0].name

    #  定义推理方法
    def detect(self, img):
        # 模型训练时图像的预处理操作
        img = self.padding(img)  
        img = cv2.resize(img, self.size) 
        img = np.float32(img)
        input_array = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0)
        # 将图片数据包装成模型输入格式
        ort_inputs = {self.input_name: input_array}
        out = self.model.run(None, ort_inputs)  # 模型推理得到预测结果
        out = np.squeeze(out)  # 删除长度为1的维度
        return out

    def padding(self, img,const=0):
        h,w,c=img.shape
        mm=max(h,w)
        dh1=int((mm-h)/2)
        dw1=int((mm-w)/2)
        dh2=(mm-h)-dh1;dw2=(mm-w)-dw1
        img=np.pad(img,((dh1,dh2),(dw1,dw2),(0,0)),constant_values=const)
        return img