黑色的是被检测的人,蓝色是检测框,蓝色数字是每个框的概率分数。
当我们检测一个物体的时候。
会有多个检测框,此时,我们需要用NMS 来删除多余的框。
我们首先选出得分最高的框–在这里就是得分0.8的框。
此时我们需要判断得分最高的框和另一个框的iou值,
假设我们设定iou的标准阈值是0.5,
如果2个框的iou算出来大于我们设定iou的标准阈值是0.5,那么我们就删掉另一个框,保留得分最高的框。
如果小于,就不删。
注意,这个操作必须是针对同一个物体的。如果是这样的情况:
下面我们来实现NMS:
import torch
from iou import intersection_over_union
def non_max_suppression(
bboxes,
iou_threshold,
prob_threshold,
box_format="corners"
):
# bboxes的结构假设是这样的【类别,框的概率,x1,y1,x2,y2】
assert type(bboxes) == list
bboxes = [box for box in bboxes if box[1] > prob_threshold]
bboxes = sorted(bboxes,key= lambda x:x[1],reverse=True)
bboxes_after_nms = []
while bboxes:
chosen_box = bboxes.pop(0)
bboxes = [
box for box in bboxes
if box[0] != chosen_box[0]
or intersection_over_union(
torch.tensor(chosen_box[2:]),
torch.tensor(box[2:])
)
< iou_threshold
]
bboxes_after_nms.append(chosen_box)
return bboxes_after_nms
iou就是我目标检测专栏的iou