非最大值抑制(Non-Maximum Suppression,NMS)是目标检测中一个非常重要的步骤,它在滑动窗口搜索或者锚框生成等方法生成的候选目标中,寻找出最具代表性的目标。
在目标检测中,模型会生成多个候选框或者锚框,每一个候选框都带有一个置信度分数,表示该框内是否包含目标以及预测准确度。之后我们需要在这些候选框中,找出最具代表性的目标框。这时候,NMS就可以派上用场了。
下面是非最大值抑制的基本步骤:
-
对所有候选框按照其内部目标的置信度(confidence score)进行排序,选择置信度最高的候选框,并将其加入输出列表中;
-
对剩余的候选框,计算它们与输出列表中已经选择的框的重叠程度(overlap),如果重叠程度大于一定阈值(比如0.5),则将该候选框从列表中删除;
-
迭代执行步骤1和2,直到所有候选框都被处理过。
这个过程可以用如下的伪代码来表示:
def nms(detections, threshold):
# detections: 输入的所有候选框
# threshold: 重叠程度阈值
boxes = detections[:, :4]
scores = detections[:, 4]
picked = []
_, idxs = scores.sort(0, descending=True)
while len(idxs) > 0:
current_idx = idxs[0]
picked.append(current_idx.item())
if len(idxs) == 1:
break
remain_boxes = boxes[idxs[1:], :]
current_box = boxes[current_idx, :]
iou = calc_iou(current_box, remain_boxes)
idxs = idxs[1:][iou < threshold]
return detections[picked, :]
其中,detections指的是所有候选框(每行包含4个坐标和一个置信度分数),threshold是重叠程度阈值。
该函数首先将所有候选框按照置信度进行降序排列,然后循环进行以下操作:
-
将当前置信度最高的候选框加入输出列表picked中;
-
计算当前候选框与剩余的候选框的重叠程度,并将重叠程度大于阈值的候选框从候选列表idxs中删除;
扫描二维码关注公众号,回复: 15225455 查看本文章 -
循环执行上述操作,直到候选框都被处理完。
最后,函数将被选中的候选框组成新的数组返回。
需要注意的是,NMS算法通常是针对单类别的目标检测,对于多类别的情况需要分别进行处理。此外,NMS的重叠程度阈值是一个关键参数,需要根据具体场景进行调整。较大的阈值可以让算法更加激进,去掉更多的候选框,但可能会漏掉一些真实目标;较小的阈值可以让算法更加保守,保留更多的候选框,但可能会留下一些重复的目标。