NMS源码
SSD论文链接
NMS介绍
吴恩达对于NMS(非极大值抑制)的介绍:
说白了,NMS的作用就是去掉目标检测任务重复的检测框。 例如,一个目标有多个选择框,现在要去掉多余的选择框。怎么做呢?循环执行步骤1和2, 直到只剩下一个框:
- 1、选出置信度p_c最高的框;
- 2、去掉和这个框IOU>0.7的框。
相关函数
一、torch.clamp( )
torch.clamp(input, min, max, out=None) → Tensor
将输入input张量每个元素夹紧到区间 [min,max],并返回结果到一个新张量。
类似于numpy中的np.clip
操作定义如下:
| min, if x_i < min
y_i = | x_i, if min <= x_i <= max
| max, if x_i > max
参数:
- input (Tensor) – 输入张量
- min (Number) – 限制范围下限
- max (Number) – 限制范围上限
- out (Tensor, optional) – 输出张量
例子:
>>> a = torch.randn(4)
>>> a
1.3869
0.3912
-0.8634
-0.5468
[torch.FloatTensor of size 4]
>>> torch.clamp(a, min=-0.5, max=0.5)
0.5000
0.3912
-0.5000
-0.5000
[torch.FloatTensor of size 4]
二、torch.index_select()
torch.index_select(input, dim, index, out=None) → Tensor
沿着指定维度对输入进行切片。
参数:
- input (Tensor) – 输入张量
- dim (int) – 索引的轴
- index (LongTensor) – 包含索引下标的一维张量
- out (Tensor, optional) – 目标张量
例子:
>>> x = torch.randn(3, 4)
>>> x
1.2045 2.4084 0.4001 1.1372
0.5596 1.5677 0.6219 -0.7954
1.3635 -1.2313 -0.5414 -1.8478
[torch.FloatTensor of size 3x4]
>>> indices = torch.LongTensor([0, 2])
>>> torch.index_select(x, 0, indices)
1.2045 2.4084 0.4001 1.1372
1.3635 -1.2313 -0.5414 -1.8478
[torch.FloatTensor of size 2x4]
>>> torch.index_select(x, 1, indices)
1.2045 0.4001
0.5596 0.6219
1.3635 -0.5414
[torch.FloatTensor of size 3x2]
注意,index_select函数中的参数index表示了有哪些索引值是需要保留的。
三、 torch.numel()
torch.numel(input)->int
返回input 张量中的元素个数。
复现代码
以下为ssd.pytorch中NMS(实际上在任何anchor based的目标检测框架中都适用)。其中:
- 为了减少计算量,作者仅选取置信度前top_k=200个框;
- 代码中包含了IOU的计算。关于IOU计算推荐阅读这篇文章;
def nms(boxes, scores, overlap=0.7, top_k=200):
"""
输入:
boxes: 存储一个图片的所有预测框。[num_positive,4].
scores:置信度。如果为多分类则需要将nms函数套在一个循环内。[num_positive].
overlap: nms抑制时iou的阈值.
top_k: 先选取置信度前top_k个框再进行nms.
返回:
nms后剩余预测框的索引.
"""
keep = scores.new(scores.size(0)).zero_().long()
if boxes.numel() == 0:
return keep
x1 = boxes[:, 0]
y1 = boxes[:, 1]
x2 = boxes[:, 2]
y2 = boxes[:, 3]
area = torch.mul(x2 - x1, y2 - y1)
v, idx = scores.sort(0)
idx = idx[-top_k:]
xx1 = boxes.new()
yy1 = boxes.new()
xx2 = boxes.new()
yy2 = boxes.new()
w = boxes.new()
h = boxes.new()
count = 0
while idx.numel() > 0:
i = idx[-1]
keep[count] = i
count += 1
if idx.size(0) == 1:
break
idx = idx[:-1]
torch.index_select(x1, 0, idx, out=xx1)
torch.index_select(y1, 0, idx, out=yy1)
torch.index_select(x2, 0, idx, out=xx2)
torch.index_select(y2, 0, idx, out=yy2)
xx1 = torch.clamp(xx1, min=x1[i])
yy1 = torch.clamp(yy1, min=y1[i])
xx2 = torch.clamp(xx2, max=x2[i])
yy2 = torch.clamp(yy2, max=y2[i])
w.resize_as_(xx2)
h.resize_as_(yy2)
w = xx2 - xx1
h = yy2 - yy1
w = torch.clamp(w, min=0.0)
h = torch.clamp(h, min=0.0)
inter = w*h
rem_areas = torch.index_select(area, 0, idx)
union = rem_areas + area[i]- inter
IoU = inter/union
idx = idx[IoU.le(overlap)]
return keep, count
本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)