【目标检测】NMS、SoftNMS

最新推荐文章于 2024-06-17 19:52:02 发布

原创

最新推荐文章于 2024-06-17 19:52:02 发布 · 1.6k 阅读

40 ·

CC 4.0 BY-SA版权

文章标签：

#目标检测 #人工智能 #计算机视觉

一、NMS

1. 作用

非极大值抑制（Non-Maximum Suppression）会根据IoU筛选出一定区域内属于同一种类得分最大的框。

关于IoU及其变种的介绍请参考IoU、GIoU、DIoU、CIoU、EIoU。一个简单的通过一次NMS直接筛选所有类别的方法，推荐参考视频18:52处讲解的PyTorch中的处理方法。

我们的以下内容均仅针对一张图像的一个类别。

2. 步骤

NMS需要4个参数：预测框坐标（ $N\times 4$ ）、预测分数（ $N$ ）、分数阈值（ $1$ ）、IoU阈值（ $1$ ）
(1) 将预测分数低于分数阈值的预测结果滤除；
(2) 将剩余预测结果按照预测分数降序排列，假设排序为 $A (0.9) > B (0.8) > C (0.7) > D (0.6) > E (0.5)$ ；
(3) 保留最大分数框 $A$ ，计算 $A$ 与剩余框 $B, C, D, E$ 间的IoU，假设分别为 $0.5, 0.6, 0.7, 0.8$ ；
(4) 将IoU大于IoU阈值的预测结果滤除，假设IoU阈值为0.6，则仅剩下 $B, C$ ；
(5) 对剩余框执行(3)、(4)，直至无剩余预测框，返回所有被保留的预测结果。

3. 实现

import torch


def IoU(boxs1, boxs2, x1y1x2y2=False):
    # boxs1.shape = 1 * 4  boxs2.shape = m * 4
    # return iou.shape = m

    if x1y1x2y2:
        b1_x1, b1_y1, b1_x2, b1_y2 = boxs1[:, 0], boxs1[:, 1], boxs1[:, 2], boxs1[:, 3]
        b2_x1, b2_y1, b2_x2, b2_y2 = boxs2[:, 0], boxs2[<