非极大值抑制算法总结(NMS, soft-NMS)

NMS-非极大值抑制算法总结

  • greedy-nms、soft-nms
    • 标准非极大值抑制 - NMS
    • 软阈值NMS - soft-NMS

greedy-nms、soft-nms

greedy-NMS是最传统的(标准)的非极大值抑制算法,soft-NMS是前者的改进版本。下面将分别从两者的原理与代码实现方面进行分析总结。

标准非极大值抑制 - 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)

软阈值NMS - soft-NMS

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置零。

未完待续…

你可能感兴趣的:(pytorch深度学习,pytorch目标检测,NMS,深度学习,pytorch,神经网络)