NMS的numpy写法
简单易懂的NMS的numpy写法
目标检测中的NMS，输入：boxes形如N*5,N个(x1,y1,x2,y2,score), thresh阈值为float型。
计算时，首先获取基于分数的降序排序order，计算全部box的面积area，对每个得分最高的boxes[i]，计算与其余低分boxes[order[1:]]的交并比ious，去除ious高于阈值的box。1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23def nms_plain(boxes, threshold):
x1 = boxes[:,0]
y1 = boxes[:,1]
x2 = boxes[:,2]
y2 = boxes[:,3]
scores = boxes[:,4]
order = scores.argsort()[::-1]
areas = (x2-x1+1)*(y2-y1+1)
keep = []
while order.size>0:
i = order[0]
keep.append(i)
xx1 = np.maximum(x1[i],x1[order[1:]] )
yy1 = np.maximum(y1[i],y1[order[1:]] )
xx2 = np.minimum(x2[i],x2[order[1:]] )
yy2 = np.minimum(y2[i],y2[order[1:]] )
w = np.maximum(0, xx2-xx1)
h = np.maximum(0,yy2-yy1)
overlaps = w*h
ious = overlaps/(areas[i] + areas[order[1:]] - overlaps)
inds = np.where(ious<threshold)[0]
order = order[inds+1]
return keep
简洁的NMS的numpy写法
下面是一个简洁的NMS写法，主要是prod操作计算面积免去了许多多余行。1
2
3
4
5
6
7
8
9
10
11
12
13
14def nms_neat(boxes, threshold):
areas = np.prod(boxes[:,2:4]-boxes[:,:2],axis=1)
order = boxes[:,4].argsort()[::-1]
keep = []
while order.size>0:
i = order[0]
keep.append(i)
tl = np.maximum(boxes[i][:2],boxes[order[1:]][:,:2])
br = np.minimum(boxes[i][2:4],boxes[order[1:]][:,2:4])
overlaps = np.prod(br-tl,axis=1)*(br>tl).all(axis=1)
ious = overlaps/(areas[i]+areas[order[1:]]-overlaps)
inds = np.where(ious<threshold)[0]
order=order[inds+1]
return keep
Soft-NMS的简洁的numpy写法
和NMS相比，Soft-NMS把高于thresh的box置信度置为score*weight,而非直接置为0。weight有多种计算方式，linear和Gaussian都有。给出一个Soft-NMS的简洁写法：1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30def soft_nms(boxes, method = 'linear', thresh = 0.001, Nt = 0.3, sigma = 0.5):
'''
the soft nms implement using python
:param dets: the pred_bboxes
:param method: the policy of decay pred_bbox score in soft nms
:param thresh: the threshold
:param Nt: Nt
:return: the index of pred_bbox after soft nms
'''
areas = np.prod(boxes[:,2:4]-boxes[:,:2], axis=1)
order = boxes[4].argsort()[::-1]
keep = []
while order.size>0:
i = order[0]
keep.append(i)
tl = np.maximum(boxes[i][:2], boxes[order[1:]][:,:2])
br = np.minimum(boxes[i][2:4], boxes[order[1:]][:,2:4])
overlaps = np.prod(br-tl, axis = 1)*(br>tl).all(axis = 1)
ious = overlaps/(areas[i]+areas[order[1:]]-overlaps)
weight = np.ones_like(overlaps)
if method == 'linear':
weight[np.where(ious>Nt)] = 1-ious[np.where(ious>Nt)]
elif method == 'gaussian':
weight = np.exp(-(ious*ious)/sigma)
else:
weight[np.where(ious>Nt)] = 0
boxes[order[1:]][:,4] *= weight
inds = np.where(boxes[order[1:]][:,4]>thresh)[0]
order = order[inds+1]
return keep
最后是画图和主函数验证：1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33def plot_bbox(boxes, c = 'r'):
x1 = boxes[:,0]
y1 = boxes[:,1]
x2 = boxes[:,2]
y2 = boxes[:,3]
plt.plot([x1,x2], [y1,y1], c)
plt.plot([x1,x1], [y1,y2], c)
plt.plot([x1,x2], [y2,y2], c)
plt.plot([x2,x2], [y1,y2], c)
plt.title(" nms")
if __name__ == "__main__":
boxes=np.array([[100,100,210,210,0.72],
[250,250,420,420,0.8],
[220,220,320,330,0.92],
[100,100,210,210,0.72],
[230,240,325,330,0.81],
[220,230,315,340,0.9]])
keep = nms_plain(boxes,0.5)
plt.figure(1)
ax1 = plt.subplot(1,2,1)
ax2 = plt.subplot(1,2,2)
plt.sca(ax1)
plot_bbox(boxes,'k') # before nms
keep = nms_neat(boxes, 0.5)
# keep = soft_nms(boxes,'dd')
plt.sca(ax2)
plot_bbox(boxes[keep], 'r')# after nms
plt.show()
NMS的pytorch写法
NMS的pytorch写法和Numpy类似，稍有不同，代码如下1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31def nms_pytorch(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() # 注意此时idx为[N-1,] 而order为[N,]
if idx.numel() == 0:
break
order = order[idx+1] # 修补索引之间的差值
return torch.LongTensor(keep) # Pytorch的索引值为LongTensor
NMS的cpp写法
本文的cpp写法是基于libtorch（pytorch c++ api），仿照pytorch，输入boxes和scores以及thresh变量。1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43//输入boxes：Nx4; socres: N; thresh:介于(0,1)的float变量
//输出vector<int>向量，存储保留的box的index值
vector<int> nms_libtorch(torch::Tensor boxes, torch::Tensor scores, float thresh);\\函数声明
vector<int> nms_libtorch(torch::Tensor bboxes, torch::Tensor scores, float thresh) {\\函数定义
auto x1 = bboxes.select(-1, 0);
auto y1 = bboxes.select(-1, 1);
auto x2 = bboxes.select(-1, 2);
auto y2 = bboxes.select(-1, 3);
auto areas = (x2 - x1)*(y2 - y1); //[N, ] 每个bbox的面积
auto tuple_sorted = scores.sort(0, true); //降序排列
auto order = std::get<1>(tuple_sorted);
vector<int> keep;
while (order.numel() > 0) {// torch.numel()返回张量元素个数
if (order.numel() == 1) {// 保留框只剩一个
auto i = order.item();
keep.push_back(i.toInt());
break;
}
else {
auto i = order[0].item();// 保留scores最大的那个框box[i]
keep.push_back(i.toInt());
}
//计算box[i]与其余各框的IOU(思路很好)
auto order_mask = order.narrow(0, 1, order.size(-1) - 1);
x1.index({ order_mask });
x1.index({ order_mask }).clamp(x1[keep.back()].item().toFloat(), 1e10);
auto xx1 = x1.index({ order_mask }).clamp(x1[keep.back()].item().toFloat(),1e10);// [N - 1, ]
auto yy1 = y1.index({ order_mask }).clamp(y1[keep.back()].item().toFloat(), 1e10);
auto xx2 = x2.index({ order_mask }).clamp(0, x2[keep.back()].item().toFloat());
auto yy2 = y2.index({ order_mask }).clamp(0, y2[keep.back()].item().toFloat());
auto inter = (xx2 - xx1).clamp(0, 1e10) * (yy2 - yy1).clamp(0, 1e10);// [N - 1, ]
auto iou = inter / (areas[keep.back()] + areas.index({ order.narrow(0,1,order.size(-1) - 1) }) - inter);//[N - 1, ]
auto idx = (iou <= thresh).nonzero().squeeze();//注意此时idx为[N - 1, ] 而order为[N, ]
if (idx.numel() == 0) {
break;
}
order = order.index({ idx + 1 }); //修补索引之间的差值
}
return keep;
}
事实上，除了使用libtorch的写法，将libtorch张量存储到c++的其他数据结构中，比如vector或者数组中同样可以实现NMS功能，本文不再赘述。
