一、概述
非极大值抑制(Non-Maximum Suppression
,NMS
),顾名思义就是抑制不是极大值的元素,可以理解为局部最大搜索。这个局部代表的是一个邻域,邻域有两个参数可变,一是邻域的维数,二是邻域的大小。例如在行人检测中,滑动窗口经提取特征,经分类器分类识别后,每个窗口都会得到一个分数。但是滑动窗口会导致很多窗口与其他窗口存在包含或者大部分交叉的情况。这时就需要用到NMS
来选取那些邻域里分数最高(是行人的概率最大),并且抑制那些分数低的窗口。
NMS在计算机视觉领域有着非常重要的应用,如视频目标跟踪、3D重建、目标识别以及纹理分析等。
![](https://i-blog.csdnimg.cn/blog_migrate/2abe0352f82e8ef1a879de53baf2e9be.png#pic_center)
如右上角的子图所示,对于同一个目标,有多个box
重叠,而我们只需要一个最好的即可,那么就需要将其他box
去除。
二、IoU
IoU(intersection over Union
),指的是两个box
的交并比,计算公式如下图所示:
![](https://i-blog.csdnimg.cn/blog_migrate/efc6f65879da28890fa2d0b1f1a58ca8.png#pic_center)
2.1 Python实现
假设box
的坐标以[top, left, bottom, right]
形式给出,即左上和右下两个坐标点,那么怎么计算两个box
的并集和交集?
首先我们应该先计算交集,如果我们考虑两个box
的相对位置(如下图所示),分情况来计算,显然非常麻烦:
![](https://i-blog.csdnimg.cn/blog_migrate/8d9f96fb22306739ff8ab281fb3022bc.jpeg#pic_center)
我们可以先考虑一维的情况:
括号部分即为两条直线的交集:
- 交集的左端点: x m a x = max ( x 1 , x 2 ) x_{max} = \max(x_1, x_2) xmax=max(x1,x2)
- 交集的右端点: y m i n = min ( y 1 , y 2 ) y_{min} = \min (y_1, y_2) ymin=min(y1,y2)
- 如果 y m i n − x m a x < 0 y_{min} - x_{max} < 0 ymin−xmax<0,即右端点在左端点左侧,代表没有交集。
#用python实现一维IoU:
def iou(set_a, set_b):
x1, x2 = set_a # (left, right)
y1, y2 = set_b # (left, right)
low = max(x1, y1)
high = min(x2, y2)
# intersection
if high - low<0:
inter = 0
else:
inter = high - low
# union
union = (x2 - x1) + (y2 - y1) - inter
# iou
iou = inter / union
return iou
二维的IoU
实质上可以通过两个一维IoU
来计算,因为一维的IoU
可以得到某一个方向上交集的上界和下界,两者相减就可以得到交集长度。如果我们将二维IoU
拆分成x
轴方向和y
轴方向,分别计算两个box
的宽和高在这两个方向上的交集长度,就可以得出交集面积(注意,我们不关心交集的具体位置),如下图所示:
而两个box
的并集计算公式为
U
n
i
o
n
=
S
1
+
S
2
−
i
n
t
e
r
s
e
c
t
i
o
n
Union = S_1 + S_2 - intersection
Union=S1+S2−intersection。但是这里要额外说明,box
所在坐标的原点是图片的左上角,因此自己拟定坐标值测试代码的时候要注意,如下图所示:
![](https://i-blog.csdnimg.cn/blog_migrate/3ac27c514950eb59a617cf9a4f6434ba.png#pic_center)
#Python实现二维IoU
def iou(box1, box2):
'''
box:[top, left, bottom, right]
'''
# 得到y轴方向上的交集长度
in_h = min(box1[2], box2[2]) - max(box1[0], box2[0])
# 得到x轴方向上的交集长度
in_w = min(box1[3], box2[3]) - max(box1[1], box2[1])
# 计算交集面积
inter = 0 if in_h < 0 or in_w < 0 else in_h * in_w
# 计算并集面积
union = (box1[2] - box1[0]) * (box1[3] - box1[1]) + \
(box2[2] - box2[0]) * (box2[3] - box2[1]) - inter
iou = inter / union
return iou
2.2 PyTorch实现
在实际的检测任务中,经常有多个Ground Truth
框,算法检测到的Box
需要和所有Ground Truth
计算IoU
。这里intersect
用来计算交集,box_a
的shape
为
(
A
,
4
)
(A, 4)
(A,4),box_b
的shape
为
(
B
,
4
)
(B, 4)
(B,4),
A
,
B
A, B
A,B分别代表有几个box
,也就是说同时对
A
A
A个box
和
B
B
B个box
计算交集。
本质上和上面的代码还是一样的原理,只是为了满足不同shape
张量之间的运算,先要对维度进行扩展。
# PyTorch实现多框对多框的IoU
def intersect(box_a, box_b):
""" We resize both tensors to [A,B,2] without new malloc:
[A,2] -> [A,1,2] -> [A,B,2]
[B,2] -> [1,B,2] -> [A,B,2]
Then we compute the area of intersect between box_a and box_b.
Args:
box_a: (tensor) bounding boxes, Shape: [A,4].
box_b: (tensor) bounding boxes, Shape: [B,4].
Return:
(tensor) intersection area, Shape: [A,B].
"""
A = box_a.size(0)
B = box_b.size(0)
max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), # [A,2]->[A,1,2]->[A,B,2]
box_b[:, 2:].unsqueeze(0).expand(A, B, 2)) # [B,2]->[1,B,2]->[A,B,2]
min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2),
box_b[:, :2].unsqueeze(0).expand(A, B, 2))
# 将交集长度限定在0以上
inter = torch.clamp((max_xy - min_xy), min = 0)
return inter[:, :, 0] * inter[:, :, 1]
def jaccard(box_a, box_b):
"""Compute the jaccard overlap of two sets of boxes. The jaccard overlap
is simply the intersection over union of two boxes. Here we operate on
ground truth boxes and default boxes.
E.g.:
A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
Args:
box_a: (tensor) Ground truth bounding boxes, Shape: [A,4]
box_b: (tensor) Prior boxes from priorbox layers, Shape: [B,4]
Return:
jaccard overlap: (tensor) Shape: [A, B]
"""
inter = intersect(box_a, box_b)
area_a = ((box_a[:, 2]-box_a[:, 0]) *
(box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter) # [A,B]
area_b = ((box_b[:, 2]-box_b[:, 0]) *
(box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter) # [A,B]
union = area_a + area_b - inter
return inter / union # [A,B]
三、NMS算法原理及其实现
NMS
主要是为了除去模型预测后的多余框,所以会设有一个threshold
作为指标。算法步骤如下:
- 选取所有
box
中scores
最大的哪一个,记为box_best
并保留。 - 计算
box_best
与其余box
的IOU
。 - 如果
IOU > threshold
了,那么就舍弃这个box
。因为针对同个目标的box
只需要保留一个,所以score
小的就没有必要保留了。 - 重新进行步骤一,循环往复。
# PyTorch实现NMS算法
# bboxes维度为[N,4],scores维度为[N,], 均为tensor
def nms(self, bboxes, scores, threshold = 0.5):
x1 = bboxes[:,0]
y1 = bboxes[:,1]
x2 = bboxes[:,2]
y2 = bboxes[:,3]
areas = (x2 - x1) * (y2 - y1) # [N,] 每个bbox的面积
_, order = scores.sort(0, descending=True) # 降序排列
keep = []
while order.numel() > 0: # torch.numel()返回张量元素个数
if order.numel() == 1: # 保留框只剩一个
i = order.item()
keep.append(i)
break
else:
i = order[0].item() # 保留scores最大的那个框box[i]
keep.append(i)
# 计算box[i]与其余各框的IOU, 原理同之前的一样, 只是实现方式有所改变。
xx1 = x1[order[1:]].clamp(min = x1[i]) # [N-1,]
yy1 = y1[order[1:]].clamp(min = y1[i])
xx2 = x2[order[1:]].clamp(max = x2[i])
yy2 = y2[order[1:]].clamp(max = y2[i])
inter = (xx2 - xx1).clamp(min = 0) * (yy2 - yy1).clamp(min = 0) # [N-1,]
iou = inter / (areas[i] + areas[order[1:]]-inter) # [N-1,]
# 除去了多余框
idx = (iou <= threshold).nonzero().squeeze()
if idx.numel() == 0:
break
# 修补索引之间的差值,因为idx不包括box_best, 所以索引比order小1
order = order[idx + 1]
return torch.LongTensor(keep) # Pytorch的索引值为LongTensor