使用onnxruntime进行yolov5s模型推理实现

首先,使用onnxruntime模型推理比使用pytorch快很多,所以模型训练完后,将模型导出为onnx格式并使用onnxruntime进行推理部署是一个不错的选择。接下来就逐步实现yolov5s在onnxruntime上的推理流程。

1、安装onnxruntime

pip install onnxruntime

2、导出yolov5s.pt为onnx,在YOLOv5源码中运行export.py即可将pt文件导出为onnx。

3、使用onnx进行推理,这一步只能获取模型推理产生的候选框,拿到推理结果后还需要进行下一步处理,通过非极大值抑制对候选框进行筛选。

使用onnx加载yolov5s.onnx模型,进行推理计算,打印推理结果:

height, width = 640, 640
img0 = cv2.imread('img.jpg')
img = cv2.resize(img0, (height, width))  # 尺寸变换
img = img / 255.

img = img[:, :, ::-1].transpose((2, 0, 1))  # HWC转CHW
data = np.expand_dims(img, axis=0)  # 扩展维度至[1,3,640,640]
sess = rt.InferenceSession('yolov5s.onnx')
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name
pred_onx = sess.run([label_name], {input_name: data.astype(np.float32)})[0]
pred = np.squeeze(pred_onx)
print(pred)
print(pred.shape)

使用onnxruntime进行yolov5s模型推理实现

此时打印的推理结果矩阵为[25200, 85],意味着有25200个检测框,每个检测框包含85个结果,85个结果中[0~3]为xywh,即检测框中心坐标和检测框宽高,[4]为每个检测框置信度,[5~85]为coco数据集80个分类对应预测结果。

显然,如果得到直接的推理结果,就不能在图像上准确地画出目标,然后需要进行非最大值抑制才能得到准确的推理结果。

[En]

Obviously, if we get the direct reasoning results, we can not accurately draw the target on the image, and then we need to carry out non-maximum suppression to obtain accurate reasoning results.

4、非极大值抑制:非极大值抑制可以采用OpenCV等接口实现,也可以手动实现。这里根据非极大值原理进行手动实现。

非极大值抑制分为置信度和iou抑制。首先进行置信度抑制,此处置信度为检测框置信度而非分类结果的置信度,设置检测框置信度阈值0.25。

置信度抑制很简单,就是根据置信度阈值,去掉阈值以下的候选框,保留阈值以上的候选框进行下一步。一张图片用于测试,输出结果为(270.85),即还剩27个检测框。

[En]

Confidence suppression is very simple, that is, according to the confidence threshold, remove the candidate boxes below the threshold, and retain the candidate boxes above the threshold to proceed to the next step. An image is used for testing, and the output result is (270.85), that is, there are 27 detection boxes left.

使用onnxruntime进行yolov5s模型推理实现

在iou抑制之前,先对所有候选框列表进行置信度排序,然后从候选框列表中找出不同类别候选框形成子列表候选框。

iou抑制整体思路:对每个类别候选框取出最大置信度候选框优先输出,再通过最大置信度候选框与其他候选框计算iou,通过iou阈值判断,大于iou阈值则删除,循环计算直至候选框列表为空。

最后将保留的候选框输出,在图像上绘制检测结果。

非极大值抑制实现:

def nms(pred, conf_thres, iou_thres):
    # 置信度抑制,小于置信度阈值则删除
    conf = pred[..., 4] > conf_thres
    box = pred[conf == True]
    # 类别获取
    cls_conf = box[..., 5:]
    cls = []
    for i in range(len(cls_conf)):
        cls.append(int(np.argmax(cls_conf[i])))
    # 获取类别
    total_cls = list(set(cls))  #删除重复项,获取出现的类别标签列表,example=[0, 17]
    output_box = []   #最终输出的预测框
    # 不同分类候选框置信度
    for i in range(len(total_cls)):
        clss = total_cls[i]   #当前类别标签
        # 从所有候选框中取出当前类别对应的所有候选框
        cls_box = []
        for j in range(len(cls)):
            if cls[j] == clss:
                box[j][5] = clss
                cls_box.append(box[j][:6])
        cls_box = np.array(cls_box)
        box_conf = cls_box[..., 4]   #取出候选框置信度
        box_conf_sort = np.argsort(box_conf)   #获取排序后索引
        max_conf_box = cls_box[box_conf_sort[len(box_conf) - 1]]
        output_box.append(max_conf_box)   #将置信度最高的候选框输出为第一个预测框
        cls_box = np.delete(cls_box, 0, 0)  #删除置信度最高的候选框
        while len(cls_box) > 0:
            max_conf_box = output_box[len(output_box) - 1]     #将输出预测框列表最后一个作为当前最大置信度候选框
            del_index = []
            for j in range(len(cls_box)):
                current_box = cls_box[j]      #当前预测框
                interArea = getInter(max_conf_box, current_box)    #当前预测框与最大预测框交集
                iou = getIou(max_conf_box, current_box, interArea)  # 计算交并比
                if iou > iou_thres:
                    del_index.append(j)   #根据交并比确定需要移出的索引
            cls_box = np.delete(cls_box, del_index, 0)   #删除此轮需要移出的候选框
            if len(cls_box) > 0:
                output_box.append(cls_box[0])
                cls_box = np.delete(cls_box, 0, 0)
    return output_box

最终推理结果:

使用onnxruntime进行yolov5s模型推理实现

Original: https://blog.csdn.net/jameschen9051/article/details/122360989
Author: 追猫人
Title: 使用onnxruntime进行yolov5s模型推理实现

原创文章受到原创版权保护。转载请注明出处:https://www.johngo689.com/30246/

转载文章受原作者版权保护。转载请注明原作者出处!

(0)

大家都在看

发表回复

登录后才能评论
免费咨询
免费咨询
扫码关注
扫码关注
联系站长

站长Johngo!

大数据和算法重度研究者!

持续产出大数据、算法、LeetCode干货,以及业界好资源!

2022012703491714

微信来撩,免费咨询:xiaozhu_tec

分享本页
返回顶部