YOLO系列目标检测后处理-非极大值抑制

        非极大值抑制（NMS）经常应用于各种目标检测模型中，使用NMS可以有效地剔除目标检测结果中多余的检测框，保留最合适的检测框。  

       以YOLOv5为例，模型输入为640*640时，推理输出结果在20*20,40*40，80*80三个尺度上的预测框总和为20*20*3+,40*40*3+80*80*3=25200，每个预测框包含检测框中心xy坐标，预测框宽高wh，预测框置信度，每个分类对应的置信度。要从这上万个预测框中取出正确的检测结果就需要用到非极大值抑制。

         非极大值抑制一般分为置信度抑制和IOU抑制，首先进行的是置信度抑制，即根据设定阈值，从检测结果结果中剔除置信度小于阈值的检测框，保留置信度较高的检测框，这一步非常交单。IOU（交并比）抑制较为复杂，如果目标检测包含多个分类，则需要对每个分类的检测框单独进行IOU抑制，以一个置信度较高的检测框为基准，与另一个同类检测框计算IOU值，如果IOU值大于设定阈值，则认为另一个检测框与基准检测框为同一目标，需要剔除该检测框。

         IOU抑制详解：IOU即为计算两个相同类别检测框的交并比。交并比即两个检测框相交区域与联合区域比值。假设两个检测框完全不相交，那么交集为0，IOU也就为0，那么可以认为两个检测框所预测为不同目标，需要保留。如果两个检测框完全重合，IOU值为1，两个检测框预测为同一个目标，那么就要剔除一个检测框。

         计算交集：两个检测框的相对位置可以分为完全不相交，相交不重合和完全重合几种情况。计算交集就是获取两个检测框重叠部分面积，这个问题可以拆分为分别计算在横向x和纵向y上的重叠长度。假设第一个检测框在横向范围(x1,x2),纵向范围为(y1,y2)第二个检测框为(x3,x4)，(y3,y4)。如果x1>x4或x2y4或y2

#计算交集
def getInter(box1, box2):
    box1_x1, box1_y1, box1_x2, box1_y2 = box1[0] - box1[2] / 2, box1[1] - box1[3] / 2, \
                                         box1[0] + box1[2] / 2, box1[1] + box1[3] / 2
    box2_x1, box2_y1, box2_x2, box2_y2 = box2[0] - box2[2] / 2, box2[1] - box1[3] / 2, \
                                         box2[0] + box2[2] / 2, box2[1] + box2[3] / 2
    if box1_x1 > box2_x2 or box1_x2 < box2_x1:
        return 0
    if box1_y1 > box2_y2 or box1_y2 < box2_y1:
        return 0
    x_list = [box1_x1, box1_x2, box2_x1, box2_x2]
    x_list = np.sort(x_list)
    x_inter = x_list[2] - x_list[1]
    y_list = [box1_y1, box1_y2, box2_y1, box2_y2]
    y_list = np.sort(y_list)
    y_inter = y_list[2] - y_list[1]
    inter = x_inter * y_inter
    return inter

         计算并集：得到相交面积后，计算并集就简单了，取两个检测框面积相加，减去交集即为并集面积。

#计算并集
def getIou(box1, box2, inter_area):
    box1_area = box1[2] * box1[3]
    box2_area = box2[2] * box2[3]
    union = box1_area + box2_area - inter_area
    iou = inter_area / union
    return iou

具体实现：

import numpy as np


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