soft-nms numpy-python实现软非极大值抑制精简代码

soft-nms相对于nms整个改进只需要使用绿色虚线表示的Soft-NMS替换红色虚线表示的NMS。B集合是检测到的所有建议框，S集合是各个建议框得分（分数是指建议框包含物体的可能性大小），Nt是指手动设置的阈值。M为当前得分最高框，bi 为待处理框。

相比其他网站的代码，此soft-nms代码实现非常精短，先给出iou的代码：

def iou(bbox,gt):
    #lt是两个框中间重叠框的最左边和最上边的坐标，rb是两个框中间重叠框的最右边和最下边的坐标
    lt = np.maximum(bbox[:,None,:2],gt[:,:2]) # [N,M,2]
    rb = np.minimum(bbox[:,None,2:4],gt[:,2:4]) # [N,M,2]
    #wh是重叠框的宽和高，+1是因为求边长，边长就等于前后两端的坐标点相减且+1
    wh = np.maximum(rb - lt  + 1 , 0) # [N,M,2]
    #求重叠框的面积
    inter_area = wh[:,:,0] * wh[:,:,1] #[N,M]
    #分别求两个框各自的面积
    bbox_area = (bbox[:,2] - bbox[:,0] + 1) * (bbox[:,3] - bbox[:,1] + 1)  #[N,]
    gt_area = (gt[:,2] - gt[:,0] + 1) * (gt[:,3] - gt[:,1] + 1)  #[M,]
    #iou的公式，重叠框面积 / 两个框面积之和减去重叠框面积
    iou = inter_area / (bbox_area[:,None] + gt_area - inter_area)  #[N,M]
    return iou   

 下面是soft-nms代码：

def soft_nms(bbox,thresh,score_thresh):
    #得分bbox第五列是得分，前四列是x0,y0,x1,y1
    score = bbox[:,4]
    #对得分进行排序
    order = np.argsort(score)
    #记录结果值，每次保存得分最高的那个框的索引，最后再用bbox[keep]取出相应框
    keep = []
    #一直筛选到没有可用的框
    while order.size > 0:
        #取得分最高的框的索引，因为order是升序，所以最后一位是得分最高的
        index = order[-1]
        #保存得分最高的那个框的索引
        keep.append(index)
        #取出这个框
        x = bbox[index]
        #计算iou，x[None,:]是为了保持shape一致，squeeze(0)是去掉第一个维度，不去掉的话结果的shape是[1,5]，再np.where就不对了，必须让其等于[5,]
        sub_bbox_iou = iou(x[None,:],bbox[order[:-1]]).squeeze(0)
        #筛选小于阈值的框，大于阈值的话，就和得分最高的那个框重叠了，所以保留不重叠的框
        index_after = np.where(sub_bbox_iou < thresh)
        #筛选大于阈值的框，然后对其进行分数衰减
        soft_index = np.where(sub_bbox_iou >= thresh)
        #对其进行分数衰减，公式是（1-iou） * score
        bbox[np.concatenate(soft_index),4] = (1 - sub_bbox_iou[np.concatenate(soft_index)]) * bbox[np.concatenate(soft_index),4]
        #判断衰减后的分数，保留大于分数阈值的框
        soft_index = np.where(bbox[np.concatenate(soft_index),4] >= score_thresh)
        #将iou小于阈值和iou大于阈值且分数大于阈值的框合并在一起
        total_index = np.concatenate((np.concatenate(index_after),np.concatenate(soft_index)))
        #重新排序
        order = np.argsort(bbox[total_index,4])
    return keep