Faster RCNN源码解读(1)-NMS非极大值抑制

参考：
非极大值抑制（Non-Maximum Suppression，NMS）
py_cpu_nms.py

image.png

# ---------------------------
# 非极大值抑制（Non-Maximum Suppression，NMS），顾名思义就是抑制不是极大值的元素
# 在faster rcnn中最后生成的回归框数量特别多且重叠，一般都是用NMS来进行边框过滤
# ---------------------------

from __future__ import absolute_import  # 绝对路径导入

import numpy as np
import torch

# 输入：
#   dets: 边界框集合(???类型),每一个单元为(x1,y1,x2,y2,confidence)
#   thresh: iou过滤阈值
# 输出：nms处理过的边界框

def nms_cpu(dets, thresh):
    dets = dets.numpy()  # ???转换为numpy, N 维数组对象 ndarray
    x1 = dets[:, 0]  # 取出所有的边界框左上角点的x坐标放入x1
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)  # 计算所有边界框的面积
    # numpy的argsort()函数：返回数组值从小到大的索引值，
    # 再加上[::-1]返回数组值从大到小的索引值,
    # 也可以order = np.argsort(-score)
    order = scores.argsort()[::-1]

    # 思想：每次选出scores中最大的那个边界框
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)  # 保留该类剩余box中得分最高的一个
        xx1 = np.maximum(x1[i], x1[order[1:]])  # 获取得分最高的边界框与其他所有框的的相交区域的左上角x坐标
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        # 计算相交的面积,不重叠时面积为0
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h

        # 计算IOU=重叠面积/(面积1+面积2-重叠面积)
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        # 保留iou小于等于阈值的边界框，其它则被过滤了
        # numpy.where() 有两种用法：
        # 1.np.where(condition, x, y):满足条件(condition)，输出x，不满足输出y。
        # 2.np.where(condition):输出满足条件(即非0)元素的坐标(等价于numpy.nonzero)
        inds = np.where(ovr <= thresh)[0]
        # 因为ovr数组的长度比order数组少一个,所以这里要将所有下标后移一位,
        # 获得下一个目标区域的得分最高的一个的索引
        order = order[inds + 1]

    return torch.IntTensor(keep)