【目标检测02】非极大值抑制 NMS

1. 前言

在检测图像中的目标时，一个目标可能会被预测出多个矩形框，而实际上我们只需要一个，如何消除冗余的边界框呢？一种方简单的方案是提升置信度的阈值，过滤掉低置信度的边界框。而另一种方案是使用非极大值抑制NMS。NMS的做法是，选出某个类别得分最高的预测框，然后看哪些预测框跟它的IoU大于阈值，就把这些预测框给丢弃掉。这里IoU的阈值是超参数，需要提前设置。

2. 原理

在二分类的场景下：

• 根据置信度得分进行排序
• 选择置信度最高的边界框添加到最终输出列表中，将其从边界框列表中删除
• 计算剩余边界框与输出列表中的边界框的IoU (IoU原理)
• 从边界框列表中删除IoU大于阈值的边界框，把小于阈值的边界框追加到输出列表中
• 重复上述过程，直至边界框列表为空

假设模型输出了11个预测框boxes, 以及对应的置信度scores:

Step0 创建选中列表，keep_list = []

Step1 对置信度进行排序，得到降序排序后边界框的索引位置列表remain_list = [ 3, 5, 10, 2, 9, 0, 1,
6, 4, 7, 8]

Step2 选出boxes[3]，此时keep_list为空，不需要计算IoU，直接将其放入keep_list，keep_list =
[3]， remain_list=[5, 10, 2, 9, 0, 1, 6, 4, 7, 8]

Step3 选出boxes[5]，此时keep_list中已经存在boxes[3]，计算出IoU(boxes[3], boxes[5]) =
0.0，显然小于阈值，则keep_list=[3, 5], remain_list = [10, 2, 9, 0, 1, 6, 4, 7, 8]

Step4 选出boxes[10]，此时keep_list=[3, 5]，计算IoU(boxes[3],
boxes[10])=0.0268，IoU(boxes[5], boxes[10])=0.0268 =
0.24，都小于阈值，则keep_list = [3, 5, 10]，remain_list=[2, 9, 0, 1, 6, 4, 7, 8]

Step5 选出boxes[2]，此时keep_list = [3, 5, 10]，计算IoU(boxes[3], boxes[2]) =
0.88，超过了阈值，直接将boxes[2]丢弃，keep_list=[3, 5, 10]，remain_list=[9, 0, 1, 6, 4, 7, 8]

Step6 选出boxes[9]，此时keep_list = [3, 5, 10]，计算IoU(boxes[3], boxes[9]) =
0.0577，IoU(boxes[5], boxes[9]) = 0.205，IoU(boxes[10], boxes[9]) = 0.88，超过了阈值，将boxes[9]丢弃掉。keep_list=[3, 5, 10]，remain_list=[0, 1, 6, 4, 7, 8]

Step7 重复上述Step6直到remain_list为空

最终得到keep_list=[3, 5, 10]，也就是预测框3、5、10被最终挑选出来了

在多分类的场景下，其实现与二分类的实现原理相同，区别在于需要对每个类别都做非极大值抑制。

3. 代码实现

（1）二分类场景下：

# 非极大值抑制
def nms(bboxes, scores, score_thresh, nms_thresh, pre_nms_topk, i<