yolo导出的onnx怎么使用onnxruntime调用

1.上几篇文件,写了怎样标注数据,获取yolo数据,怎样训练数据。这篇文章写下,怎样使用yolo训练好的网络进行推理,怎样使用onnxruntime调用网络推理。

2.yolo模型训练后推理,主要分两种,一种是使用yolo官方提供的api推理,一种使用onnx模式推理。

3.使用yolo官方的api推理:
###把qimage图片转成mat函数
def CoverQImageToMat(self):
ptr_=self.CurrentImage.constBits() ###用pyqt加载的图片
###QImage内部是按照每个像素4个字节的方式组织数据的,即使最后一个alpha通道没有使用,也用0xff来填充
image_mat_=np.array(ptr_).reshape(self.CurrentImage.height(),self.CurrentImage.width(),4)
image_mat_rgb_=cv2.cvtColor(image_mat_,cv2.COLOR_RGBA2RGB)
return image_mat_rgb_

self.Model=YOLO(‘best.pt’)###加载模型
image_mat_ = self.CoverQImageToMat()###图片转换,把图片转成opencv的mat
results_=self.Model(image_mat_)###推理获取结果
for result_ in results_: #####循环结果
boxes_ = result_.boxes ####获取boxes
for box_ in boxes_: ####循环boxes
x_,y_,w_,h_=box_.xywh[0] ####获取识别到对象的左上角坐标,长,宽
c=box_.cls ####获取属于第几类
cls_=self.Model.names[int©] ####查询类名
self.BoxX.append(float(x_))
self.BoxY.append(float(y_))
self.BoxWidth.append(float(w_))
self.BoxHeight.append(float(h_))
self.BoxClass.append(cls_)

4.yolo官方提供的api推理,可以按照官方文档就可以推理,这里还在讲下导出的onnx模型后,怎么加载推理。

4.1self.Model.export(format=‘onnx’)###把yolo的模型导出成onnx格式

4.2 使用onnxruntime调用onnx推理。首先安装onnxruntime.lib,这里使用的版本是
onnxruntime-1.16.3-cp311-cp311-win_amd64.whl
onnx-1.15.0-cp311-cp311-win_amd64.whl
这两个文件可以到https://pypi.org/收索下载,使用pip离线安装

4.3 代码导入onnxruntime
import onnxruntime as ort
import onnx

self.OnnxModel = ort.InferenceSession(‘best.onnx’)###加载onnx模型

###推理函数

def UseOnnxInference(self):      	  
  	  if (self.CurrentImage.width() > 0) & (self.CurrentImage.height() > 0):
  	  	####清空结果缓存
        self.BoxY.clear()	
        self.BoxX.clear()
        self.BoxWidth.clear()
        self.BoxHeight.clear()
        self.BoxClass.clear()
        
        #获取onnx模型输入点
        inputs =self.OnnxModel.get_inputs()
        len(inputs)
        input_onnx_=inputs[0]
        print("Name:",input_onnx_.name)##输入节点名称
        print("Type:", input_onnx_.type)##数据格式
        print("Shape:",input_onnx_.shape)###数据维度

        image_mat_ = self.CoverQImageToMat()  ###图片转换
        print("ImageMatShape:",image_mat_.shape)####查看图片的维度

        img_width_ = image_mat_.shape[1]  ###获取图片输入输出
        img_height_ = image_mat_.shape[0]###获取图片输入输出

        target_image_height_ = 640###输入到onnx推理的图片大小
        target_image_width_ = 640###输入到onnx推理的图片大小


        scale_precentage_=min(target_image_height_/image_mat_.shape[0],target_image_width_/image_mat_.shape[1])
        image_mat_=cv2.resize(image_mat_,None,fx=scale_precentage_,fy=scale_precentage_)#缩放成onnx推理的图片大小

        image_np_=np.array(image_mat_)###图片转成np数组
        print("ImageNpShape:", image_np_.shape)####查看下维度

        image_np_=image_np_.transpose(2,0,1)##转成通道在前面的维度
        print("ImageNpShape:", image_np_.shape)

        image_np_=image_np_.reshape(1,3,640,640)##添加一个新维度
        print(image_np_[0, 0, 0, 0])
        print("ImageNpShape:", image_np_.shape)

        image_np_ = image_np_.astype(np.float32)####把数据转成float32格式

        image_np_=image_np_/255.0##数据归一化
        print(image_np_[0, 0, 0, 0])

        ###获取输出点
        outputs=self.OnnxModel.get_outputs()
        output_onnx_=outputs[0]
        print("Name:",output_onnx_.name)###输出节点名称
        print("Type:",output_onnx_.type)###输出节点格式
        print("Shape:",output_onnx_.shape)###输出节点维度

        ###运行推理
        outputs=self.OnnxModel.run(["output0"],{"images":image_np_})
        len(outputs)
        
        ###outputs处理
        output_onnx_=outputs[0][0]
        print("OutputOnxxShape:",output_onnx_.shape)

        ###转换 行列
        output_onnx_=output_onnx_.transpose()
        print("OutputOnxxShape:", output_onnx_.shape)
        row=output_onnx_[0]
        print(row)###输出下第一个特征

        ###获取多小行特征
        rows_=output_onnx_.shape[0]

        ###获取boxs
        boxs=[]
        for i in range(rows_):
            ###获取最大概率 对应的index
            prob_ = output_onnx_[i][4:].max()

            if prob_>self.ConfidenceThres:
                class_id_ = output_onnx_[i][4:].argmax()
                ###前是个特征代表 x中心 y中心 w宽 h高
                xc_,yc_,w_,h_ = output_onnx_[i][:4]
                ###转换成x1 y1 x2 y2
                x1_=(xc_-w_/2)/target_image_width_*img_width_
                y1_=(yc_-h_/2)/target_image_height_*img_height_
                x2_=(xc_+w_/2)/target_image_width_*img_width_
                y2_=(yc_ + h_ / 2) / target_image_height_ * img_height_

                boxs.append([x1_,y1_,x2_,y2_,class_id_,prob_])

        ###去掉重叠的box
        boxs.sort(key=lambda x:[5],reverse=True)
        result_box_=[]

        lenght_box_=len(boxs)

        while lenght_box_>0:

            box_temp_=boxs[0].copy()

            boxs_temp_ = []

            ##不是重叠的添加进去boxs_temp_
            for box_ in boxs:
                ###计算intersection
                box1_x1_,box1_y1_,box1_x2_,box1_y2_=box_temp_[:4]
                box2_x1_, box2_y1_, box2_x2_, box2_y2_ = box_[:4]

                distance1_=math.sqrt((box1_x1_-box2_x1_)**2+(box1_y1_-box2_y1_)**2)
                distance2_=math.sqrt((box1_x2_-box2_x2_)**2+(box1_y2_-box2_y2_)**2)

                if (distance1_

4.4这里是一个使用onnx检测安全帽子的demo
https://gitee.com/wenyuanmo/py-qt-load-yolo-or-onnx-check-hat/tree/master

你可能感兴趣的:(yolo,YOLO)