NMS从原理到实现与检测网络的后处理

NMS

class, score, x, y, w, h

• 按class分组
• 按score排序
• 按x,y,w,h计算IoU

处理冗余框 (NMS)

NMS 算法流程
输入：box list (每个box包含自身坐标和对应score), 域值T
输出：keep list
（1）选取当前最高分的框A，将box list所有与A的IoU大于域值T的框删除掉；
（2）将A放入keep list 中；
(3) 选取剩下的框里边分数最高的，重复步骤(1)， 直到box list 为空；
（4）返回 keep list。

NMS代码：

#torch
def nms(boxes, scores, threshold=0.5, top_k=200):
    '''
    Args:
        boxes: 预测出的box, shape[M,4]
        scores: 预测出的置信度，shape[M]
        threshold: 阈值
        top_k: 要考虑的box的最大个数
    Return:
        keep: nms筛选后的box的新的index数组
        count: 保留下来box的个数
    '''
    keep = scores.new(scores.size(0)).zero_().long()
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]

    area = (x2-x1)*(y2-y1)  # 面积,shape[M]
    _, idx = scores.sort(0, descending=True) # 降序排列scores的值大小
    # 取前top_k个进行nms
    idx = idx[:top_k]

    count = 0

    while idx.numel():
        # 记录最大score值的index
        i = idx[0]
        # 保存到keep中
        keep[count] = i
        # keep 的序号
        count += 1

        if idx.size(0) == 1: # 保留框只剩一个
            break

        idx = idx[1:] # 移除已经保存的index

        # 计算boxes[i]和其他boxes之间的iou
        xx1 = x1[idx].clamp(min=x1[i])
        yy1 = y1[idx].clamp(min=y1[i])
        xx2 = x2[idx].clamp(max=x2[i])
        yy2 = y2[idx].clamp(max=y2[i])

        w = (xx2 - xx1).clamp(min=0)
        h = (yy2 - yy1).clamp(min=0)

        # 交集的面积
        inter = w * h  # shape[M-1]
        iou = inter / (area[i] + area[idx] - inter)

        # iou满足条件的idx
        idx = idx[iou.le(threshold)] # Shape[M-1]

    return keep, count

def nms(dets,thresh):
    #boxes location
    x1=dets[:,0]
    y1=dets[:,1]
    x2=dets[:,2]
    y2=dets[:,3]
    #boxes scores
    scores=dets[:,4]
    areas=(x2-x1+1)*(y2-y1+1) #area of each box
    order=scores.argsort()[::-1] #order里存的是box的索引

    keep=[] 
    while order.size>0
        i=order[0] #score最大的box对应的index
        keep.append(i) #将本轮score最大的box的index保留

        #计算剩余的boxes与当前box的重叠程度IoU
        xx1=np.maximum(x1[i],x1[order[1:]])
        yy1=np.maximum(y1[i],y1[order[1:]])
        xx2=np.maximum(x2[i],x2[order[1:]])
        yy2=np.maximum(y2[i],y2[order[1:]])

        w=np.maximum(0.0,xx2-xx1+1)
        h=np.maximum(0.0,yy2-yy1+1)
        inter=w*h
        #IoU
        ovr=inter/(areas[i]+areas[order[1:]]-inter)
        #保留IoU小于设定阈值的boxes
        inds = np.where(ovr <= thresh)[0] 
        order=order[inds+1]
    return keep

对预测图像的缩放