NMS方法的总结可以参考我之前的文章:
https://blog.csdn.net/qq_34919792/article/details/108186234
非极大值抑制(Non-Maximum Suppression,NMS),顾名思义就是抑制不是极大值的元素。在检测中,我们通过将IOU大于一定阈值的框做一个筛选,只保留置信度最高的框。
网上比较经典的实现思路
def py_cpu_nms(dets, thresh):
"""Pure Python NMS baseline."""
#x1、y1、x2、y2、以及score赋值
x1 = dets[:, 0]
y1 = dets[:, 1]
x2 = dets[:, 2]
y2 = dets[:, 3]
scores = dets[:, 4]
#每一个检测框的面积
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
#按照score置信度降序排序
order = scores.argsort()[::-1]
keep = [] #保留的结果框集合
while order.size > 0:
i = order[0]
keep.append(i) #保留该类剩余box中得分最高的一个
#得到相交区域,左上及右下
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:]])
#计算相交的面积,不重叠时面积为0
w = np.maximum(0.0, xx2 - xx1 + 1)
h = np.maximum(0.0, yy2 - yy1 + 1)
inter = w * h
#计算IoU:重叠面积 /(面积1+面积2-重叠面积)
ovr = inter / (areas[i] + areas[order[1:]] - inter)
#保留IoU小于阈值的box
inds = np.where(ovr <= thresh)[0]
order = order[inds + 1] #因为ovr数组的长度比order数组少一个,所以这里要将所有下标后移一位
return keep
其实在NMS可以做pytorch版本,速度和cpython加速后一样,即和pytorch自带的库一样,但是更适合改。
def IOU(box_a, box_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]
def torch_nms(boxes, scores, iou_threshold):
_, idx = scores.sort(0, descending=True) # descending表示降序
boxes_idx = boxes[idx]
iou = IOU(boxes_idx, boxes_idx).triu_(diagonal=1) #取上三角矩阵,不包含对角线
B = iou
while 1:
A = B
maxA, _ = torch.max(A, dim=0)
E = (maxA <= iou_threshold).float().unsqueeze(1).expand_as(A)
B = iou.mul(E)
if A.equal(B) == True:
break
keep= idx[maxA <= iou_threshold]
return keep
这样的代码其实很好改进,比如可以把IOU去改成GIOU或者DIOU或者自己魔改魔改,也可以去改进NMS,提供个DIOU的改进示例。
def DIOU(box_a, box_b, delta = 0.9):
inter = intersect(box_a, box_b)
area_a = ((box_a[:, :, 2]-box_a[:, :, 0]) *
(box_a[:, :, 3]-box_a[:, :, 1])).unsqueeze(2).expand_as(inter) # [A,B]
area_b = ((box_b[:, :, 2]-box_b[:, :, 0]) *
(box_b[:, :, 3]-box_b[:, :, 1])).unsqueeze(1).expand_as(inter) # [A,B]
union = area_a + area_b - inter
x1 = ((box_a[:, :, 2]+box_a[:, :, 0]) / 2).unsqueeze(2).expand_as(inter)
y1 = ((box_a[:, :, 3]+box_a[:, :, 1]) / 2).unsqueeze(2).expand_as(inter)
x2 = ((box_b[:, :, 2]+box_b[:, :, 0]) / 2).unsqueeze(1).expand_as(inter)
y2 = ((box_b[:, :, 3]+box_b[:, :, 1]) / 2).unsqueeze(1).expand_as(inter)
t1 = box_a[:, :, 1].unsqueeze(2).expand_as(inter)
b1 = box_a[:, :, 3].unsqueeze(2).expand_as(inter)
l1 = box_a[:, :, 0].unsqueeze(2).expand_as(inter)
r1 = box_a[:, :, 2].unsqueeze(2).expand_as(inter)
t2 = box_b[:, :, 1].unsqueeze(1).expand_as(inter)
b2 = box_b[:, :, 3].unsqueeze(1).expand_as(inter)
l2 = box_b[:, :, 0].unsqueeze(1).expand_as(inter)
r2 = box_b[:, :, 2].unsqueeze(1).expand_as(inter)
cr = torch.max(r1, r2)
cl = torch.min(l1, l2)
ct = torch.min(t1, t2)
cb = torch.max(b1, b2)
D = (((x2 - x1)**2 + (y2 - y1)**2) / ((cr-cl)**2 + (cb-ct)**2 + 1e-7))
out = inter / union - D ** delta
return out if use_batch else out.squeeze(0)
def torch_nms(boxes, scores, iou_threshold):
_, idx = scores.sort(0, descending=True) # descending表示降序
boxes_idx = boxes[idx]
iou = DIOU(boxes_idx, boxes_idx).triu_(diagonal=1) #取上三角矩阵,不包含对角线
B = iou
while 1:
A = B
maxA, _ = torch.max(A, dim=0)
E = (maxA <= iou_threshold).float().unsqueeze(1).expand_as(A)
B = iou.mul(E)
if A.equal(B) == True:
break
keep= idx[maxA <= iou_threshold]
return keep