2021-07-31

NMS算法步骤以及python实现

非极大值抑制的方法是：先假设有6个矩形框，根据分类器的类别分类概率做排序，假设从小到大属于车辆的概率 分别为A、B、C、D、E、F。

(1)从最大概率矩形框F开始，分别判断A~E与F的重叠度IOU是否大于某个设定的阈值;

(2)假设B、D与F的重叠度超过阈值，那么就扔掉B、D；并标记第一个矩形框F保留。

(3)从剩下的矩形框A、C、E中，选择概率最大的E，然后判断E与A、C的重叠度，重叠度大于一定的阈值，那么就扔掉；并标记E是我们保留下来的第二个矩形框。

(4)重复，找到所有被保留下来的矩形框。

def NMS(dects, threshhold):
    """
    detcs:二维数组(n_samples,5)
    5列：x1,y1,x2,y2,score
    threshhold: IOU阈值
    """
    x1 = dects[:, 0]
    y1 = dects[:, 1]
    x2 = dects[:, 2]
    y2 = dects[:, 3]
    score = dects[:, 4]
    ndects = dects.shape[0]  # box的数量
    area = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = score.argsort()[::-1]  # score从大到小排列的indexs,一维数组
    keep = []  # 保存符合条件的index
    suppressed = np.array([0] * ndects)  # 初始化为0，若大于threshhold,变为1，表示被抑制

    for _i in range(ndects):
        i = order[_i]  # 从得分最高的开始遍历
        if suppressed[i] == 1:
            continue
        keep.append(i)
        for _j in range(i + 1, ndects):
            j = order[_j]
            if suppressed[j] == 1:  # 若已经被抑制，跳过
                continue
            xx1 = np.max(x1[i], x1[j])  # 求两个box的交集面积interface
            yy1 = np.max(y1[i], y1[j])
            xx2 = np.min(x2[i], x2[j])
            yy2 = np.min(y2[i], y2[j])
            w = np.max(0, xx2 - xx1 + 1)
            h = np.max(0, yy2 - yy1 + 1)
            interface = w * h
            overlap = interface / (area[i] + area[j] - interface)  # 计算IOU（交/并）

            if overlap >= threshhold:  # IOU若大于阈值，则抑制
                suppressed[j] = 1
    return keep

soft-NMS

soft-nms解决当两个ground truth的目标的确重叠度很高时错误过滤的问题

def box_soft_nms(bboxes, scores, labels, nms_threshold=0.3, soft_threshold=0.3, sigma=0.5, mode='union'):
    """
    soft-nms implentation according the soft-nms paper
    :param bboxes: all pred bbox
    :param scores: all pred cls
    :param labels: all detect class label，注：scores只是单纯的得分，需配合label才知道具体对应的是哪个类
    :param nms_threshold: origin nms thres, for judging to reduce the cls score of high IoU pred bbox
    :param soft_threshold: after cls score of high IoU pred bbox been reduced, soft_thres further filtering low score pred bbox
    :return:
    """
    unique_labels = labels.cpu().unique().cuda() # 获取pascal voc 20类标签
    box_keep = []
    labels_keep = []
    scores_keep = []
    for c in unique_labels: # 相当于NMS中对每一类的操作，对应step-1
        c_boxes = bboxes[labels == c] # bboxes、scores、labels一一对应，按照label == c就可以取出对应类别 c 的c_boxes、c_scores
        c_scores = scores[labels == c]
        weights = c_scores.clone()
        x1 = c_boxes[:, 0]
        y1 = c_boxes[:, 1]
        x2 = c_boxes[:, 2]
        y2 = c_boxes[:, 3]
        areas = (x2 - x1 + 1) * (y2 - y1 + 1) # bbox面积
        _, order = weights.sort(0, descending=True) # bbox根据score降序排序，对应NMS中step-2
        while order.numel() > 0: # 对应NMS中step-5
            i = order[0] # 当前order中的top-1，保存之
            box_keep.append(c_boxes[i]) # 保存bbox
            labels_keep.append(c) # 保存cls_id
            scores_keep.append(c_scores[i]) # 保存cls_score
            if order.numel() == 1: # 当前order就这么一个bbox了，那不玩了，下一个类的bbox操作吧
                break
            xx1 = x1[order[1:]].clamp(min=x1[i]) # 别忘了x1[i]对应x1[order[0]]，也即top-1，寻找Inp区域的坐标
            yy1 = y1[order[1:]].clamp(min=y1[i])
            xx2 = x2[order[1:]].clamp(max=x2[i])
            yy2 = y2[order[1:]].clamp(max=y2[i])
            w = (xx2 - xx1 + 1).clamp(min=0) # Inp区域的宽、高、面积
            h = (yy2 - yy1 + 1).clamp(min=0)
            inter = w * h
            ovr = inter / (areas[i] + areas[order[1:]] - inter)
            # 经过origin NMS thres，得到高IoU的bboxes index，
            # origin NMS操作就直接剔除掉这些bbox了，soft-NMS就是对这些bbox对应的score做权重降低
            ids_t= (ovr>=nms_threshold).nonzero().squeeze() # 高IoU的bbox，与inds = np.where(ovr >= nms_threshold)[0]功能类似
            weights[[order[ids_t+1]]] *= torch.exp(-(ovr[ids_t] * ovr[ids_t]) / sigma)
            # soft-nms对高IoU pred bbox的score调整了一次，soft_threshold仅用于对score抑制，score太小就不考虑了
            ids = (weights[order[1:]] >= soft_threshold).nonzero().squeeze() # 这一轮未被抑制的bbox
            if ids.numel() == 0: # 竟然全被干掉了，下一个类的bbox操作吧
                break
            c_boxes = c_boxes[order[1:]][ids] # 先取得c_boxes[order[1:]]，再在其基础之上操作[ids]，获得这一轮未被抑制的bbox
            c_scores = weights[order[1:]][ids]
            _, order = c_scores.sort(0, descending=True)
            if c_boxes.dim()==1:
                c_boxes=c_boxes.unsqueeze(0)
                c_scores=c_scores.unsqueeze(0)
            x1 = c_boxes[:, 0] # 因为bbox已经做了筛选了，areas需要重新计算一遍，抑制的bbox剔除掉
            y1 = c_boxes[:, 1]
            x2 = c_boxes[:, 2]
            y2 = c_boxes[:, 3]
            areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    return box_keep, labels_keep, scores_keep # scores_keep保存的是未做权重降低的score，降低权重的score仅用于soft-nms操作