NMS高效实现（Python）

01 什么是非极大值抑制

Github: 非极大值抑制实现

非极大值抑制，简称为NMS算法，虽然在不同应用中实现的具体方式不太一样，但是思想还是一样的。下面我们看一张非极大值抑制的效果图：

左图是人脸检测的候选框结果，每个边界框都有一个置信度得分(Confidence Score)，如果不用非极大值抑制，那么就会有多个候选框出现。而右图则是经过极大值抑制之后的结果，其符合我们对人脸检测的预期效果。效果显然易见，如果不进行非极大值抑制，这结果简直没法看好嘛！下面我们来实现这个算法，包括相应的实现

02 非极大值抑制基本实现

首先，我们需要以下这些参数：

1. 目标边界框的列表
2. 对应的置信度列表
3. 阀值

为了便于性能拓展，我们用类将其封装为抽象数据类型，将相应的属性和方法都封装在一起。现在我们先回顾一下非极大值抑制的基本流程：

while 直到原来置信度框中没有东西:
    1. 根据置信度得分对其进行排序
	2. 选择置信度最高的边界框，将其加入到最终输出的列表，然后我们将其从原来列表中删除
	3. 计算 置信度最高的边界框 与 其它候选框 的 IoU
	4. 删除IoU大于阀值的边界框

但是我们会做一个优化！我们将一些计算进行了矩阵化，事实上让我从零写也不一定写得出，调了好久。其实大家知道极大值抑制最原始思路是怎么做就可以了，要实现的时候来复制粘贴一下吧！

还有一点，numpy的算法效率真的很高，我尝试用numba.jit做了加速，发现相对于numpy来说还是慢了4-5倍。一般来说，最后能优化的就是多进程，速度大概还能提升4-5倍，但是既然这个模块不是瓶颈，其实没有什么必要去优化！都已经用Python了，既然性能不是瓶颈，那么就这样吧。真正要部署实际生产：

• 大多数时候我们在算法开发完成之后会用C++来重写。因为C++也有矩阵运算的包，也有OpenCv，因此大多数情况处理图像的时候，我更偏向于用OpenCv。
• 改到TF-Lite也不是很困难的事情，因为Numpy的基本操作TF都有。

在这里除了可以学习到如何绘制Bounding box，我们还能学习到怎样在原图上绘制Box，这是我们可以学习到的。

import numpy as np
import cv2


def nms(bounding_boxes, confidences, threshold):
    """
    Args:
        bounding_boxes: np.array([(x1, y1, x2, y2), ...])
        confidences: np.array(conf1, conf2, ...),数量需要与bounding box一致,并且一一对应
        threshold: IOU阀值,若两个bounding box的交并比大于该值，则置信度较小的box将会被抑制

    Returns:
        bounding_boxes: 经过NMS后的bounding boxes
        confidences: 经过NMS后的confidences
    """
    len_bound = bounding_boxes.shape[0]
    len_conf = confidences.shape[0]
    if len_bound != len_conf:
        raise ValueError("Bounding box 与 Confidence 的数量不一致")
    if len_bound == 0:
        return np.array([]), np.array([])
    bounding_boxes, confidences = bounding_boxes.astype(np.float), np.array(confidences)

    x1, y1, x2, y2 = bounding_boxes[:, 0], bounding_boxes[:, 1], bounding_boxes[:, 2], bounding_boxes[:, 3]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    idxs = np.argsort(confidences)

    pick = []
    while len(idxs) > 0:
        # 因为idxs是从小到大排列的，last_idx相当于idxs最后一个位置的索引
        last_idx = len(idxs) - 1
        # 取出最大值在数组上的索引
        max_value_idx = idxs[last_idx]
        # 将这个添加到相应索引上
        pick.append(max_value_idx)

        xx1 = np.maximum(x1[max_value_idx], x1[idxs[: last_idx]])
        yy1 = np.maximum(y1[max_value_idx], y1[idxs[: last_idx]])
        xx2 = np.minimum(x2[max_value_idx], x2[idxs[: last_idx]])
        yy2 = np.minimum(y2[max_value_idx], y2[idxs[: last_idx]])

        w, h = np.maximum(0, xx2 - xx1 + 1), np.maximum(0, yy2 - yy1 + 1)

        iou = w * h / areas[idxs[: last_idx]]
        # 删除最大的value,并且删除iou > threshold的bounding boxes
        idxs = np.delete(idxs, np.concatenate(([last_idx], np.where(iou > threshold)[0])))

    # bounding box 返回一定要int类型,否则Opencv无法绘制
    return np.array(bounding_boxes[pick, :]).astype(int), confidences[pick]


def plot_image_boxes(filename, bounding_boxes, confidences, title, color=(0, 0, 255)):
    image = cv2.imread(filename)
    boxes_number = confidences.shape[0]
    for i in range(boxes_number):
        (x1, y1, x2, y2), conf = bounding_boxes[i, :], confidences[i]
        (w, h), baseline = cv2.getTextSize(str(conf),
                                           fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                                           fontScale=1,
                                           thickness=2)
        cv2.rectangle(image, (x1, y1 - (2 * baseline + 5)), (x1 + w, y1), color, -1)
        cv2.rectangle(image, (x1, y1), (x2, y2), color, 1)
        cv2.putText(image, text=str(conf), org=(x1, y1), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                    fontScale=1, color=(0, 0, 0), thickness=2)
    cv2.imshow(title, image)
    cv2.waitKey(0)  # 按空格


if __name__ == '__main__':
    bounding_boxes = np.array([(187, 82, 337, 317), (150, 67, 305, 282), (246, 121, 368, 304)], dtype=int)
    confidences = np.array([0.9, 0.75, 0.8])

    nms_boxes, nms_confs = nms(bounding_boxes, confidences, threshold=0.2)
    plot_image_boxes('nms.jpg', bounding_boxes, confidences, title='Original')
    plot_image_boxes('nms.jpg', nms_boxes, nms_confs, title='NMS')