【目标检测】NMS、SoftNMS

一、NMS

1. 作用

非极大值抑制（Non-Maximum Suppression）会根据IoU筛选出一定区域内属于同一种类得分最大的框。

关于IoU及其变种的介绍请参考IoU、GIoU、DIoU、CIoU、EIoU。一个简单的通过一次NMS直接筛选所有类别的方法，推荐参考视频18:52处讲解的PyTorch中的处理方法。

我们的以下内容均仅针对一张图像的一个类别。

2. 步骤

NMS需要4个参数：预测框坐标（$N\times 4$）、预测分数（$N$）、分数阈值（$1$）、IoU阈值（$1$）
(1) 将预测分数低于分数阈值的预测结果滤除；
(2) 将剩余预测结果按照预测分数降序排列，假设排序为$A(0.9)>B(0.8)>C(0.7)>D(0.6)>E(0.5)$；
(3) 保留最大分数框$A$，计算$A$与剩余框$B,C,D,E$间的IoU，假设分别为$0.5,0.6,0.7,0.8$；
(4) 将IoU大于IoU阈值的预测结果滤除，假设IoU阈值为0.6，则仅剩下$B,C$；
(5) 对剩余框执行(3)、(4)，直至无剩余预测框，返回所有被保留的预测结果。

3. 实现

import torch


def IoU(boxs1, boxs2, x1y1x2y2=False):
    # boxs1.shape = 1 * 4  boxs2.shape = m * 4
    # return iou.shape = m

    if x1y1x2y2:
        b1_x1, b1_y1, b1_x2, b1_y2 = boxs1[:, 0], boxs1[:, 1], boxs1[:, 2], boxs1[:, 3]
        b2_x1, b2_y1, b2_x2, b2_y2 = boxs2[:, 0], boxs2[:, 1], boxs2[:, 2], boxs2[:, 3]
    else:
        b1_x1, b1_x2 = boxs1[:, 0] - boxs1[:, 2] / 2, boxs1[:, 0] + boxs1[:, 2] / 2
        b1_y1, b1_y2 = boxs1[:, 1] - boxs1[:, 3] / 2, boxs1[:, 1] + boxs1[:, 3] / 2
        b2_x1, b2_x2 = boxs2[:, 0] - boxs2[:, 2] / 2, boxs2[:, 0] + boxs2[:, 2] / 2
        b2_y1, b2_y2 = boxs2[:, 1] - boxs2[:, 3] / 2, boxs2[:, 1] + boxs2[:, 3] / 2

    cap_x1, cap_y1 = torch.max(b1_x1, b2_x1), torch.max(b1_y1, b2_y1)
    cap_x2, cap_y2 = torch.min(b1_x2, b2_x2), torch.min(b1_y2, b2_y2)

    in_area = torch.clamp(cap_x2 - cap_x1, min=0) * torch.clamp(cap_y2 - cap_y1, min=0)
    un_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1) + (b2_x2 - b2_x1) * (b2_y2 - b2_y1)
    iou = in_area / torch.clamp(un_area - in_area, min=1e-10)
    return iou


def NMS(boxes, scores, score_th, iou_th):
    ''' 
    :param boxes:      N * 4
    :param scores:     N
    :param score_th:   1
    :param iou_th:     1
    :return:           remaining index
    '''
    results = []
    index = torch.argsort(-scores)
    index = index[scores > score_th]  # 滤除低分数

    while index.size(0):
        i = index[0:1]  # 这样取可以保留单个值的维度
        results.append(i)

        iou = IoU(boxes[i], boxes[index[1:]])

        idx = torch.where(iou < iou_th)[0]  # 小于阈值的被保留
        index = index[idx+1]  # 因为iou的计算不涉及第一个，所以要+1
    return torch.cat(results)


if __name__ == "__main__":
    boxes = torch.rand(255, 4) * 416
    scores = torch.rand(255)
    
    print(NMS(boxes, scores, 0.01, 0.5))

一、Soft-NMS

1. 作用

Soft-NMS认为在进行NMS的时候要同时考虑重合程度和得分。如果存在同一类别的两个目标彼此重叠，即使两个预测框的得分均较高，在NMS下因两个预测框的重合程度较大也只能保留一个，如下图：

两个预测框均是针对马的，两个预测分数（0.95、0.80）也均较高，重叠程度（IoU）很可能大于阈值，如果执行NMS，仅可能保留红框，Soft-NMS为避免该问题，提出根据与最大得分框的IoU降低与之重叠框的得分而非直接去除的方法。

这样，如果两个高得分的框重叠较多，在得分最高的框被保留后，另一个分数同样较高的框虽然得分有所降低，但因其本身得分相比其它框更高，在下一轮比较中有希望脱颖而出。

2. 步骤

Soft-NMS需要4个参数：预测框坐标（$N\times 4$）、预测分数（$N$）、分数阈值（$1$）、IoU阈值（$1$）、加权类型（str）、$\sigma$（高斯加权时会用到的参数）
(1) 将预测结果按照预测分数降序排列，假设排序为$A(0.9)>B(0.8)>C(0.7)>D(0.6)>E(0.2)$；
(2) 保留最大分数框$A$，计算$A$与剩余框$B,C,D,E$间的IoU，假设分别为$0.65,0.65,0.6,0.2$；
(3) 对IoU大于IoU阈值的得分进行加权，假设IoU阈值为0.3，则$E$无需加权，设1-IoU为得分修正权重，则权重分别为$0.35,0.35,0.4,1$，加权后得分变为$0.28,0.245,0.24,0.2$；
(4) 将得分低于得分阈值的预测结果滤除，假设得分阈值为$0.22$，则仅剩下$B,C,D$；
(5) 对剩余框继续执行(2)-(4)，直至无剩余预测框，返回所有被保留的预测结果。
如此，Soft-NMS在第二次执行时$B$作为最高得分框会被保留。但如果是NMS，第二次执行时仅剩下$E$。

Soft-NMS有两种加权策略，其中一种为线形加权，公式如下

其中，$s_i$为加权前后的得分，$1-iou(\mathcal{M},b_i)$为加权权重，仅作用于大于IoU阈值$N_t$的预测框。可以看出，与最大得分框的IoU越大，得分被折损的越多。

另外一种为高斯加权，公式如下：

其中，$\sigma$为传入的超参数。同样地，与最大得分框的IoU越大，得分被折损的越多。

事实上，NMS也可以视为一种加权，公式如下：

其中，与最大得分框IoU大的被加权为0，IoU小的权重为1。

3. 实现

def SoftNMS(boxes, scores, score_th, iou_th, type='gaussian', sigma=0.5):
    '''
    :param boxes:      N * 4
    :param scores:     N
    :param score_th:   1
    :param iou_th:     1
    :param type:       'linear', 'gaussian', 'greedy'
    :param sigma:      parameters used to obtain weights
    :return:           remaining boxes index
    '''
    results = []
    indexes = torch.arange(boxes.size(0))
    while(boxes.size(0)):
        index_sorted = torch.argsort(scores, descending=True)
        boxes = boxes[index_sorted]
        scores = scores[index_sorted]
        indexes = indexes[index_sorted]

        results.append(indexes[0])
        ious = IoU(boxes[0].unsqueeze(0), boxes[1:]).squeeze(0)

        if type == 'linear':
            weights = torch.ones(ious.size())
            weights[ious >= iou_th] -= ious[ious > iou_th]  # 仅iou大的才会被加权，其余的仍为1
        elif type == 'greedy':  # 即NMS
            weights = torch.ones(ious.size())
            weights[ious >= iou_th] = 0
        elif type == 'gaussian':
            weights = torch.exp(-(ious * ious) / sigma)
        else:
            raise ValueError('wrong type!')

        scores = weights * scores[1:]
        boxes = boxes[1:][scores > score_th]
        indexes = indexes[1:][scores > score_th]
        scores = scores[scores > score_th]  # 这个要放最后，不然分数变了无法正确索引
    return torch.stack(results)


if __name__ == "__main__":
    boxes = torch.rand(255, 4) * 416
    scores = torch.rand(255)

    print(SoftNMS(boxes, scores, 0.01, 0.5))

由于Soft-NMS在加权后才执行过滤，与NMS一开始就进行一次过滤的结果会有些许差异。

致谢：

本博客仅做记录使用，无任何商业用途，参考内容如下：
睿智的目标检测31——非极大抑制NMS与Soft-NMS
目标检测中NMS的相关概念与计算
手撕Soft-NMS