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

weixin_46648511

已于 2025-07-16 09:44:32 修改

阅读量2.5k

点赞数 34

CC 4.0 BY-SA版权

分类专栏： yolo 文章标签： YOLO

于 2024-02-05 23:32:48 首次发布

本文链接：https://blog.youkuaiyun.com/weixin_46648511/article/details/136032406

yolo 专栏收录该内容

24 篇文章

订阅专栏

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(c)]   ####查询类名
                    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.1

self.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_<self.Iou)&(distance2_<self.Iou):
                        pass
                    else:
                        boxs_temp_.append(box_)
                        pass

                boxs.clear()###清空原来的

                for box_ in boxs_temp_:
                    boxs.append(box_)

                lenght_box_ = len(boxs)

                result_box_.append(box_temp_)

            print(len(result_box_))

            for box_ in result_box_:
                box1_x1_, box1_y1_, box1_x2_, box1_y2_ = box_[:4]

                class_id_,prob_=box_[4:]

                width = box1_x2_-box1_x1_
                height = box1_y2_-box1_y1_

                center_x_= (box1_x2_+box1_x1_)/2
                center_y_ = (box1_y2_+box1_y1_)/2

                self.BoxX.append(float(center_x_))
                self.BoxY.append(float(center_y_))
                self.BoxWidth.append(float(width))
                self.BoxHeight.append(float(height))

                if class_id_==0:
                    self.BoxClass.append('head')

                if class_id_==1:
                    self.BoxClass.append('helmet')

                if class_id_==2:
                    self.BoxClass.append('person')