个人技术笔记

python实现nms

发布时间：2021-08-24  栏目：Python, Tensorflow, 人工智能, 图像处理, 机器视觉, 深度学习  评论：0 Comments

NMS的思路是：对于一张图片中的每一个预测框来说，模型为其每一个类别都预测了一个置信度分数（一般多分类，模型输出后接softmax，每一个类别都得到了一个置信度分数，包括背景类）我们取置信度最高的那一个类别作为预测框中对象所属的类别。1. 首先我们将置信度分数低于置信度阈值a的所有预测框去掉 。2. 然后在同一张图片上，我们按照类别（除开背景类，因为背景类不需要进行NMS），将所有预测框按照置信度从高到低排序，将置信度最高的框作为我们要保留的此类别的第1个预测框，3. 然后按照顺序计算剩下其他预测框与其的IoU，4. 去掉与其IoU大于IoU阈值b的预测框（其实代码实现里是将这些要去掉的预测框其置信度分数置为0），5. 第一次迭代结束，我们已经剔除了与第一个框重合度较高的框。

接着从剩下的预测框中取置信度分数最高的检测框作为我们要保留的第2个预测框，进行第2次迭代。反复下去，我们就过滤掉此类别与同一GT重叠度较高的预测框了，然后对下一个类别处理，直至处理完所有的类别。

 

 

def nms(output, nms_th):
    if len(output) == 0:
        return output

    output = output[np.argsort(-output[:, 0])]
    bboxes = [output[0]]

    for i in np.arange(1, len(output)):
        bbox = output[i]
        flag = 1
        for j in range(len(bboxes)):
            if iou(bbox[1:5], bboxes[j][1:5]) >= nms_th:
                flag = -1
                break
        if flag == 1:
            bboxes.append(bbox)

    bboxes = np.asarray(bboxes, np.float32)
    return bboxes

留下评论

You must be logged in to post a comment.