【目标检测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=0, c=0):"""nms"""# 置信度排序，结果降序排列，返回对应的索引inds = np.argsort(scores)inds = inds[::-1]# 输出列表keep_inds = []while(len(inds) > 0):# 每次遍历，从边界框列表中拿出一个框，默认从第一个开始处理cur_ind = inds[0]cur_score = scores[cur_ind]if cur_score < score_thresh:breakkeep = Truefor ind in keep_inds:current_box = bboxes[cur_ind]remain_box = bboxes[ind]iou = box_iou_xyxy(current_box, remain_box)# 过滤大于IoU阈值的边界框if iou > nms_thresh:keep = Falsebreakif keep:keep_inds.append(cur_ind)# 从边界框列表中删除已经处理过的框，因为每次都是从第一个框开始处理，因此把0索引的框删掉inds = inds[1:]return np.array(keep_inds)

（2）多分类场景下：

def multiclass_nms(bboxes, scores, score_thresh=0.01, nms_thresh=0.45, pre_nms_topk=1000, pos_nms_topk=100):batch_size = bboxes.shape[0]class_num = scores.shape[1]rets = []for i in range(batch_size):bboxes_i = bboxes[i]scores_i = scores[i]ret = []#对每一个类别单独做非极大值抑制for c in range(class_num):scores_i_c = scores_i[c]keep_inds = nms(bboxes_i, scores_i_c, score_thresh, nms_thresh, pre_nms_topk, i=i, c=c)if len(keep_inds) < 1:continuekeep_bboxes = bboxes_i[keep_inds]keep_scores = scores_i_c[keep_inds]keep_results = np.zeros([keep_scores.shape[0], 6])keep_results[:, 0] = ckeep_results[:, 1] = keep_scores[:]keep_results[:, 2:6] = keep_bboxes[:, :]ret.append(keep_results)if len(ret) < 1:rets.append(ret)continueret_i = np.concatenate(ret, axis=0)scores_i = ret_i[:, 1]if len(scores_i) > pos_nms_topk:inds = np.argsort(scores_i)[::-1]inds = inds[:pos_nms_topk]ret_i = ret_i[inds]rets.append(ret_i)return rets