手撕非极大值抑制算法NMS

前言

非极大值抑制算法（Non-maximum suppression, NMS）是有anchor系列目标检测的标配，如今大部分的 One-Stage 和 Two-Stage 算法在推断(Inference)阶段都使用了NMS作为网络的最后一层，例如YOLOv3、SSD、Faster-RCNN等。

当然NMS在目前最新的 anchor-free 目标检测算法中(CornerNet、CenterNet等)并不是必须的，对这种检测算法提升的精度也有限，但是这并不影响我们学习NMS。

NMS的本质是搜索局部极大值，抑制非极大值元素，在目标检测中，我们经常将其用于消除多余的检测框(从左到右消除了重复的检测框，只保留当前最大confidence的检测框)：

NMS有很多种变体，这里只介绍最为常见的Hard-NMS，我们通常所说的NMS就是指Hard-NMS。

Hard-NMS

Hard-NMS就是我们传统意义上的NMS，也是最常用的NMS算法。

因为NMS主要用于消除多余的检测框，那么消除的标准是什么，我们使用IOU作为标准来进行演示，IOU的原称为Intersection over Union，也就是两个box区域的交集比上并集，下图可以方便理解：

具体介绍可以看这里： 深度学习中IU、IoU(Intersection over Union)的概念理解以及 python 程序实现

因为我们要手撸么，所以废话不多说，直接开始看代码，首先使用Pytorch来看一篇：

def hard_nms(box_scores, iou_threshold, top_k=-1, candidate_size=200):
    """
    Args:
        box_scores (N, 5): box的集合，N为框的数量，5即4(位置信息)+1(可能为物体的概率)
        iou_threshold: 我们用IOU标准去除多余检测框的阈值
        top_k: 保留多少个计算后留下来的候选框，如果为-1则全保留
        candidate_size: 参与计算的boxes数量
    Returns:
         picked: 经过nms计算后保留下来的box
    """
    scores = box_scores[:, -1]                # 首先我们取出box中的最后一个元素也就是当前box检测到物体的概率
    boxes = box_scores[:, :-1]                # 取出box中的四个坐标(左上、右下)
    picked = []  
    _, indexes = scores.sort(descending=True) # 按照降序排列所有的物体的概率，得到 排序 后在原数组中的索引信息 indexes
    indexes = indexes[:candidate_size]        # 只保留前 candidate_size 个 boxes 其余的不考虑了
    while len(indexes) > 0:
        current = indexes[0]                  # 每次取出当前在 indexes 中 检测到物体概率最大的一个 
        picked.append(current.item())         # 将这个最大的存在结果中
        if 0 < top_k == len(picked) or len(indexes) == 1:
            break
        current_box = boxes[current, :]       # 当前第一个也就是最高概率的box
        indexes = indexes[1:]                
        rest_boxes = boxes[indexes, :]        # 剩下其余的box
        iou = iou_of(                         # 将当前的box与剩下其余的boxes用IOU标准进行筛选
            rest_boxes,
            current_box.unsqueeze(0),
        )
        indexes = indexes[iou <= iou_threshold]# 保留与当前box的IOU小于一定阈值的boxes，

    return box_scores[picked, :]

看了上面的代码，我们可以知道大概的流程：

• 选取这类box中scores最大的那一个，记为current_box，并保留它(为什么保留它，因为它预测出当前位置有物体的概率最大啊，对于我们来说当前confidence越大说明当前box中包含物体的可能行就越大)
• 计算current_box与其余的box的IOU
• 如果其IOU大于我们设定的阈值，那么就舍弃这些boxes（由于可能这两个box表示同一目标，因此这两个box的IOU就比较大，会超过我们设定的阈值，所以就保留分数高的那一个）
• 从最后剩余的boxes中，再找出最大scores的那一个(之前那个大的已经保存到输出的数组中，这个是从剩下的里面再挑一个最大的)，如此循环往复

至于上述代码中 iou_of 部分：

def area_of(left_top, right_bottom) -> torch.Tensor:
    """Compute the areas of rectangles given two corners.

    Args:
        left_top (N, 2): left top corner.
        right_bottom (N, 2): right bottom corner.

    Returns:
        area (N): return the area.
    """
    hw = torch.clamp(right_bottom - left_top, min=0.0)
    return hw[..., 0] * hw[..., 1]


def iou_of(boxes0, boxes1, eps=1e-5):
    """Return intersection-over-union (Jaccard index) of boxes.

    Args:
        boxes0 (N, 4): ground truth boxes.
        boxes1 (N or 1, 4): predicted boxes.
        eps: a small number to avoid 0 as denominator.
    Returns:
        iou (N): IoU values.
    """
    overlap_left_top = torch.max(boxes0[..., :2], boxes1[..., :2])
    overlap_right_bottom = torch.min(boxes0[..., 2:], boxes1[..., 2:])

    overlap_area = area_of(overlap_left_top, overlap_right_bottom)
    area0 = area_of(boxes0[..., :2], boxes0[..., 2:])
    area1 = area_of(boxes1[..., :2], boxes1[..., 2:])
    return overlap_area / (area0 + area1 - overlap_area + eps)

手撕NMS

手撕代码用什么撕，当然是用C++撕，这才爽么！

直接看代码，其中使用了OpenCV库中的 Point2f 结构体：

// 这是一个模板函数，接受一个已经排好序的vector，然后降序返回其索引
template <typename T>
vector<int> sort_indexes(const vector<T> &v) {


    vector<int> idx(v.size());
    iota(idx.begin(), idx.end(), 0);
    sort(idx.begin(), idx.end(),
         [&v](int i1, int i2) {return v[i1] > v[i2];});
    return idx;
}
// 这就是我们的NMS函数 输入的坐标已经标准化，所有数值的范围为 0-1
/*
* numBoxes：窗口数目
* points：窗口左上角坐标点
* oppositePoints：窗口右下角坐标点
* score：窗口得分
* overlapThreshold：重叠阈值控制
* numBoxesOut：输出窗口数目
* pointsOut：输出窗口左上角坐标点
* oppositePoints：输出窗口右下角坐标点
* scoreOut：输出窗口得分
*/

int nonMaximumSuppression(int numBoxes, const vector<Point2f>& points,
                          const vector<Point2f>& oppositePoints, const vector<float>& score,
                          float overlapThreshold,
                          int *numBoxesOut, vector<Point2f>& pointsOut,
                          vector<Point2f>& oppositePointsOut, vector<float>& scoreOut)
{

    const float eps = 1e-5;
    int i, j, index;
    float* box_area = (float*)malloc(numBoxes * sizeof(float));  // 定义窗口面积变量并分配空间
    vector<int> indices;
    int* is_suppressed = (int*)malloc(numBoxes * sizeof(int));   // 定义是否抑制表标志并分配空间

    // 初始化indices、is_supperssed、box_area信息
    for (i = 0; i < numBoxes; i++)
    {
        indices.push_back(i);
        is_suppressed[i] = 0;
        box_area[i] = ((oppositePoints[i].x - points[i].x + eps) *
                       (oppositePoints[i].y - points[i].y + eps));
    }

    // 对输入窗口按照分数比值进行排序，排序后的编号放在indices中
    indices = sort_indexes(score);

    for (i = 0; i < numBoxes; i++)                // 循环所有窗口
    {
        if (!is_suppressed[indices[i]])           // 判断窗口是否被抑制
        {
            for (j = i + 1; j < numBoxes; j++)    // 循环当前窗口之后的窗口
            {
                if (!is_suppressed[indices[j]])   // 判断窗口是否被抑制
                {
                    float x1max = max(points[indices[i]].x, points[indices[j]].x);                     // 求两个窗口左上角x坐标最大值
                    float x2min = min(oppositePoints[indices[i]].x, oppositePoints[indices[j]].x);     // 求两个窗口右下角x坐标最小值
                    float y1max = max(points[indices[i]].y, points[indices[j]].y);                     // 求两个窗口左上角y坐标最大值
                    float y2min = min(oppositePoints[indices[i]].y, oppositePoints[indices[j]].y);     // 求两个窗口右下角y坐标最小值
                    float overlapWidth = x2min - x1max + eps;            // 计算两矩形重叠的宽度
                    float overlapHeight = y2min - y1max + eps;           // 计算两矩形重叠的高度
                    if (overlapWidth > 0 && overlapHeight > 0)
                    {
                        float overlapPart = (overlapWidth * overlapHeight) / box_area[indices[j]];    // 计算重叠的比率
                        if (overlapPart > overlapThreshold)          // 判断重叠比率是否超过重叠阈值
                        {
                            is_suppressed[indices[j]] = 1;           // 将窗口j标记为抑制
                        }
                    }
                }
            }
        }
    }

    *numBoxesOut = 0;    // 初始化输出窗口数目0
    for (i = 0; i < numBoxes; i++)
    {
        if (!is_suppressed[i]) (*numBoxesOut)++;    // 统计输出窗口数目
    }

    for (i = 0; i < numBoxes; i++)                  // 遍历所有输入窗口
    {
        if (!is_suppressed[indices[i]])             // 将未发生抑制的窗口信息保存到输出信息中
        {
            Point2f temp_xy(points[indices[i]].x, points[indices[i]].y);
            Point2f temp_Opxy(oppositePoints[indices[i]].x, oppositePoints[indices[i]].y);
            pointsOut.push_back(temp_xy);
            oppositePointsOut.push_back(temp_Opxy);
            scoreOut.push_back(score[indices[i]]);
        }
    }

    indices.clear();
    free(box_area);         // 释放box_area空间
    free(is_suppressed);    // 释放is_suppressed空间

    return 0;
}

好了，代码撕完了，我们只需要输入相应的数据就可以使用，亲测。

但是要注意，此代码仅供学习演示，并没有进行过多的优化，大家有能力者自行优化吧！