greedy-NMS是最传统的(标准)的非极大值抑制算法,soft-NMS是前者的改进版本。下面将分别从两者的原理与代码实现方面进行分析总结。
标准NMS算法的流程:
1.首先对需要进行抑制的同类别bounding-boxes按照预测的置信度conf进行降序排序;
2.取出排好序的ns_boxes中的第一项, 计算其与剩余项间交并比,获得一个交并比list;
3. 判断这些交并比值是否大于 设定的抑制阈值 nms_thresh (关键假设:如果 iou[x] > nms_thresh, 说明 该bounding-box与置信度最大bounding-box同属一个实例的推断,应该被抑制; 反之,该bounding-box属于不同实例),大于阈值就抑制(删掉),小于就保留。执行完一轮后,将剩下的boxes继续执行上述步骤;
4. 停止条件就是 ns_boxes (需要抑制的集合中没有任何选项) 为空。
GPU版代码:
import torch
def nms(ns_boxes, nms_thresh):
# ns_boxes:need_selected_boxes; nms_thresh: nms阈值
s_box = torch.tensor([]).cuda()
write = 0
while ns_boxes.size(0):
ious = bbox_iou_tensor(ns_boxes[0].unsqueeze(0), ns_boxes, False)
iou_mask = ious> nms_thresh
weights = ns_boxes[iou_mask, 4:5]
ns_boxes[0, :4] = (weights * ns_boxes[iou_mask, :4]).sum(0) / weights.sum() # 融合策略
if not write:
s_box = ns_boxes[0]
else:
s_box = torch.cat((s_box, ns_boxes[0]), 0)
ns_boxes = ns_boxes[~iou_mask]
return s_box.view(-1, 7)
soft-nms算法的流程:
总体来说,大部分还是遵循标准NMS算法的思路,不同之处在于:大于nms_thresh阈值的边框没有被直接抑制,而是根据交并比情况对即将抑制边框的置信度进行处理(加权缩小?),然后根据soft_nms_thresh阈值决定是否抑制 (小于软阈值的保留,大于软阈值的真正的被抑制)。 然后保留下的边框置信度采用被修正后的置信度。
对剩余边框按置信度重排序后,重复上述步骤,终止条件与nms算法一致。
GPU版代码
def soft_nms(ns_boxes, nms_thresh, soft_thresh=0.5, gima=0.5):
s_box = torch.tensor([]).cuda()
write = 0
while ns_boxes.size(0):
iou_mask = torch.zeros((ns_boxes.size(0)), dtype=torch.bool) # 筛选掩膜
ious = bbox_iou_tensor(ns_boxes[0].unsqueeze(0), ns_boxes[1:], False)
if not write:
s_box = ns_boxes[0]
write = 1
else:
s_box = torch.cat((s_box, ns_boxes[0]), 0)
for i in range(ious.shape[0]):
if ious[i] > nms_thresh:
ns_boxes[i+1, 4] *= torch.exp(-ious[i]*ious[i] / sigma)
if ns_boxes[i+1, 4] >= soft_thresh: # 调整后的置信度得分小于软阈值则删除该项,反之相反
iou_mask[i+1] = True # 掩码置真
else:
iou_mask[i+1] = True # 小于原来的nms_thresh的项保留
ns_boxes = ns_boxes[iou_mask] # 删除掩码=False的项
resort_idx = torch.sort(ns_boxes[:, 4], descending=True)[1] # 按改变的conf进行排序
ns_boxes = ns_boxes[resort_idx]
return s_box.view(-1, 7)
在上述两个NMS算法中,都有一个公共的函数:计算不同bounding-boxes间的交并比: bbox_iou_tensor(box1, box2, x1y1x2y2=True)
(多次试验结论是:gpu版的iou计算方式笔cpu版(numpy)版本的速度要快)
总结一下这个代码的巧妙之处: 充分利用截断函数torch.clamp() 与最值函数torch.max() .min()
def bbox_iou_tensor(box1, box2, x1y1x2y2=True):
# box1, box2是需要求交并比的两个边框,可以是多维tensor; x1y1x2y2是box的描述形式
# True为对角顶点形式, 如果是[x,y,w,h]中心点+尺度形式要做一下转化
if not x1y1x2y2:
b1_x1, b1_x2 = box1[:,0] - box1[:,2]/2, box1[:,0] + box1[:,2]/2
b1_y1, b1_y2 = box1[:,1] - box1[:,3]/2, box1[:,1] + box1[:,3]/2
b2_x1, b2_x2 = box2[:,0] - box2[:,2]/2, box2[:,0] + box2[:,2]/2
b2_y1, b2_y2 = box2[:,1] - box2[:,3]/2, box2[:,1] + box2[:,3]/2
else:
b1_x1, b1_y1, b1_x2, b1_y2 = box1[:,0], box1[:,1], box1[:,2], box1[:,3]
b2_x1, b2_y1, b2_x2, b2_y2 = box2[:,0], box2[:,1], box2[:,2], box2[:,3]
b1_area = (b1_x2-b1_x1+1) * (b1_y2-b1_y1+1)
b2_area = (b2_x2-b2_x1+1) * (b2_y2-b2_y1+1)
xx1 = torch.max(b1_x1, b2_x1)
yy1 = torch.max(b1_y1, b2_y1)
xx2 = torch.min(b1_x2, b2_x2)
yy2 = torch.min(b1_y2, b2_y2)
intersection = torch.clamp(xx2-xx1+1, min=0)*torch.clamp(yy2-yy1+1, min=0)
ious = intersection / (b1_area + b2_area - intersection + 1e-16)
return ious
辅助理解:
重点是计算交集, 首先计算两边框左顶点的最大值坐标(也即,最大的x和最大的y值);然后计算两边框右底点的最小值坐标(也即,最小的x和最小y)。如果存在交集,那么最大左顶点坐标和最小右底点坐标即为相交边框的对角坐标。如果不存在交集,则利用截断函数将w或h置零。
未完待续…