YOLOv5改进系列(八) 更换NMS非极大抑制DIoU-NMS、CIoU-NMS、EIoU-NMS、GIoU-NMS 、SIoU-NMS、Soft-NMS

一、NMS（非极大抑制）简介

1.1 什么是NMS？

NMS（non maximum suppression）即非极大抑制，顾名思义就是抑制不是极大值的元素，搜索局部的极大值。在最近几年常见的物体检测算法（包括RCNN、SPPNet、Fast-RCNN、Faster-RCNN等）中，最终都会从一张图片中找出很多个可能是物体的矩形框，然后为每个矩形框为做类别分类概率。

如果用一句话概括NMS的意思就是：筛选出一定区域内属于同一种类别得分最大的框。

如下图，网络模型可以给每个检测框一个score，score越大，说明检测框越接近真实值。

《老友记》yyds！

然后非最大值抑制的目的是删除score小的框，只剩下sore最大作为最终的预测结果。

---

1.2 NMS的计算过程

• 定义置信度阈值和IOU阈值取值
• 按置信度降序排列边界框bounding_box
• 从bbox_list中删除置信度小于阈值的预测框
• 循环遍历剩余框，首先挑选置信度最高的框作为候选框
• 接着计算其他和候选框属于同一类的所有预测框和当前候选框的IOU
• 如果上述任两个框的IOU的值大于IOU阈值，那么从box_list中移除置信度较低的预测框
• 重复此操作，直到遍历完列表中的所有预测框

---

1.3 NMS的局限性

（1）需要手动设置阈值，阈值的设置会直接影响重叠目标的检测，太大造成误检，太小达不到理想情况。

（2）低于阈值的直接设置score为0，做法就比较麻烦了

（3）通过IoU来评估，IoU的做法对目标框尺度和距离的影响不同

---

1.4 NMS的改进思路

（1）根据手动设置阈值的缺陷，通过自适应的方法在目标稀疏时使用小阈值，目标稠密时使用大阈值。例如Adaptive NMS
（2）由于将低于阈值的直接置为0的做法比较困难，所以我们可以通过将其根据IoU大小来进行惩罚衰减，则变得更加平滑。例如Soft NMS，Softer NMS
（3）IoU的做法存在一定缺陷，改进思路是将目标尺度、距离引进IoU的考虑中。如DIoU等

---

二、DIoU-NMS、CIoU-NMS、EIoU-NMS、GIoU-NMS 、 SIoU-NMS

2.1 更换DIoU-NMS

一个成熟的IoU衡量指标应该要考虑预测框与真实框的重叠面积、中心点距离、长宽比三个方面。但是IoU 只考虑到了预测框与真实框重叠区域，并没有考虑到中心点距离、长宽比。

DIoU-NMS在DIoUloss一文中提出，不仅仅考虑IoU，还考虑两个框中心点之间的距离。如果两个框之间IoU比较大，但是两个框的中心距离比较大时，可能会认为这是两个物体的框而不会被过滤掉。由于DIoU的计算考虑到了两框中心点位置的信息，故使用DIoU进行评判的nms效果更符合实际，效果更优。

公式：

---

第①步 修改general.py

在YOLOv5当中，作者是直接调用了Pytorch官方的NMS的API。

也就是在general.py中的non_max_suppression函数中

首先将下面一段代码粘贴到utils/general.py，重新定义NMS模块。这里的计算IoU的函数bbox_iou则是直接引用了YOLOv5中的代码，其简洁的集成了CIoU、SIoU、EIoU、GIoU、DIoU 的计算。

def NMS(boxes, scores, iou_thres, class_nms='CIoU'):
    # 初始化不同的IoU (Intersection over Union) 类型标志
    GIoU = CIoU = DIoU = EIoU = SIoU = False

    # 根据 class_nms 参数选择要使用的 IoU 类型
    if class_nms == 'CIoU':
        CIoU = True
    elif class_nms == 'DIoU':
        DIoU = True
    elif class_nms == 'GIoU':
        GIoU = True
    elif class_nms == 'EIoU':
        EIoU = True
    else:
        SIoU = True

    # 根据分数对边界框索引进行排序（从高到低）
    B = torch.argsort(scores, dim=-1, descending=True)

    # 存储最终选定的边界框索引
    keep = []

    while B.numel() > 0:
        # 选择分数最高的边界框的索引
        index = B[0]
        # 将该索引添加到最终保留的列表中
        keep.append(index)
        if B.numel() == 1:
            break

        # 计算当前边界框与其余边界框的IoU
        iou = bbox_iou(boxes[index, :], boxes[B[1:], :], GIoU=GIoU, DIoU=DIoU, CIoU=CIoU, EIoU=EIoU, SIoU=SIoU)
        # 找到IoU小于等于指定阈值的边界框索引
        inds = torch.nonzero(iou <= iou_thres).reshape(-1)
        # 更新B以排除已经被选择的边界框
        B = B[inds + 1]

    # 返回最终选定的边界框索引作为结果
    return torch.tensor(keep)

---

第②步 更换NMS

然后我们将non_max_suppression 函数中的

i = torchvision.ops.nms(boxes, scores, iou_thres)

替换为

i = NMS(boxes, scores, iou_thres, class_nms='DIoU')

这样就可以开始训练了~

---

2.2 更换其他的NMS

其余几个方法都是一样的，只需要在第②步改个名称即可：

• DloU：

i = NMS(boxes, scores, iou_thres, class_nms='DIoU') 

• SloU：

i = NMS(boxes, scores, iou_thres, class_nms='SIoU') 

• GloU：

i = NMS(boxes, scores, iou_thres, class_nms='GIoU') 

• EloU：

i = NMS(boxes, scores, iou_thres, class_nms='EIoU') 

---

三、Soft-NMS

根据前面对目标检测中NMS的算法描述，易得出传统NMS的不足：如果一个物体在另一个物体重叠区域出现，即当两个目标框接近时，分数更低的框就会因为与之重叠面积过大而被删掉，从而导致对该物体的检测失败并降低了算法的平均检测率。

上图中，有两匹马，都是待检测目标，也有两个检测到的框，得分分别是0.80和0.95
如果用NMS算法，得分最高的框是红色框，得分0.95，而绿色框与红色框通过计算IoU，肯定是大于一般我们设置的0.5的，那么绿色框就会被删除，导致少检测到一匹马的情况。此外，NMS算法设置阈值也比较麻烦，如果设置过小，那么会出先这样的事情，少检测到目标；如果设置过大，又会经常出先误检。

因此，出现升级版Soft-NMS。具体流程就是我们把NMS算法中去除其他边界框改成，修改其他边界框的置信度。

---

第①步 修改general.py

同样，首先将下面一段代码粘贴到utils/general.py，重新定义NMS模块。

def my_soft_nms(bboxes, scores, iou_thresh=0.5, sigma=0.5, score_threshold=0.25):
    # 将边界框张量转换为连续内存布局
    bboxes = bboxes.contiguous()
    x1, y1, x2, y2 = bboxes[:, 0], bboxes[:, 1], bboxes[:, 2], bboxes[:, 3]
    
    # 计算每个边界框的面积
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    
    # 对得分进行降序排序，并获取排序后的索引
    _, order = scores.sort(0, descending=True)
    
    # 存储最终选定的边界框索引
    keep = []

    while order.numel() > 0:
        i = order[0].item()
        keep.append(i)

        # 计算当前边界框与其他边界框的交集坐标
        xx1 = x1[order[1:]].clamp(min=x1[i])
        yy1 = y1[order[1:]].clamp(min=y1[i])
        xx2 = x2[order[1:]].clamp(max=x2[i])
        yy2 = y2[order[1:]].clamp(max=y2[i])

        # 计算交集的面积
        inter = (xx2 - xx1).clamp(min=0) * (yy2 - yy1).clamp(min=0)
        
        # 计算当前边界框与其他边界框的IoU（交并比）
        iou = inter / (areas[i] + areas[order[1:]] - inter)

        # 找到IoU大于指定阈值的边界框索引
        idx = (iou > iou_thresh).nonzero().squeeze()

        if idx.numel() > 0:
            iou = iou[idx]
            
            # 计算IoU对应的新得分，并应用得分衰减
            new_scores = torch.exp(-torch.pow(iou, 2) / sigma)
            
            # 更新具有高IoU的边界框的得分
            scores[order[idx + 1]] *= new_scores

        # 根据得分阈值过滤边界框
        new_order = (scores[order[1:]] > score_threshold).nonzero().squeeze()

        if new_order.numel() == 0:
            break
        else:
            new_scores = scores[order[new_order + 1]]
            max_score_index = torch.argmax(new_scores)

            if max_score_index != 0:
                new_order[[0, max_score_index]] = new_order[[max_score_index, 0]]

            order = order[new_order + 1]

    # 返回最终选定的边界框索引作为结果
    return torch.LongTensor(keep)

---

第②步 更换NMS

将general.py中将NMS改为soft nms。

这步也是和上面一样，将non_max_suppression 函数中的

i = torchvision.ops.nms(boxes, scores, iou_thres)

替换为

     i = my_soft_nms(boxes, scores, iou_thres)  # 

---