使用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)

 

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

   显然拿到直接的推理结果还无法在图像上进行准确的目标绘制，接下来还要进行非极大值抑制才能获取准确推理结果。

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

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

   置信度抑制非常简单，即根据置信度阈值，移除阈值以下的候选框，保留阈值以上的候选框进入下一步处理。使用一张图像进行测试，输出结果为(27,85)，即剩下27个检测框。

在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

最终推理结果：