非极大值抑制算法(NMS)的python实现

前言

  本篇博客主要是介绍非极大值抑制NMS算法的python实现，并根据实例检测实现效果。

1. NMS概述

  非极大值抑制(Non-Maximum Supression, NMS)，顾名思义，就是抑制非极大值，在目标检测领域中经常使用到，主要是用来对候选框进行去重处理。
  NMS算法的大致流程如下：
  (1) 根据概率分数score对候选框进行排序
  (2) 选择概率分数最大的bbox，记录这个bbox到输出列表中，并删除和这个选框IoU大于一定阈值的bbox；
  (3) 继续选择概率分数最大的边框，并添加到输出列表中，重复上处过程直至没有候选框为止。

2. 绘制候选框

  以上述迪迦奥特曼为例，我们首先提供三个候选框，并在候选框的基础之上生成其他的候选框，实现代码如下：

import cv2
import numpy as np
import copy


seed = 10001
np.random.seed(seed)

bounding_boxes = [
        [545, 125, 765, 440],
        [890, 100, 1115, 430],
        [1275, 170, 1490, 490]
    ]

confidence_score = [0.95, 0.98, 0.96]

num_anchor = 10
anchors = copy.deepcopy(bounding_boxes)
scores = copy.deepcopy(confidence_score)


if __name__ == '__main__':
    img = cv2.imread('dijia.png')

    for i in range(num_anchor):
        index = np.random.randint(0, 3)
        offset = np.random.randint(-50, 50, size=4)
        score = np.random.uniform(0.5, 0.9)
        anchors.append(list(bounding_boxes[index] - offset))
        scores.append(round(score, 2))

    for i in range(len(scores)):
        cv2.rectangle(img, pt1=tuple(anchors[i][:2]), pt2=tuple(anchors[i][2:]), color=(0, 255, 0), thickness=2)
        cv2.putText(img, text=str(scores[i]), org=tuple(anchors[i][:2]), fontFace=cv2.FONT_HERSHEY_SCRIPT_SIMPLEX,
                    fontScale=1, color=(255, 0, 255), thickness=2)
    cv2.imshow('dijia', img)
    cv2.waitKey()
    cv2.imwrite('dijia1.png', img)

# cv2.rectangle: 	pt1: 左上角坐标
# 					pt2: 右下角坐标
# 					color: 边框颜色(B, G, R)
# 					thickness: 边框粗细

# cv2.putText: 	text: 文字信息
# 				org: 起始点坐标(左下角)
# 				fontFace: 字体类型
# 				fontScale: 字体大小
# 				color: 字体颜色(B, G, R)
# 				thickness: 字体粗细

  代码运行结果如下：

3. NMS代码实现

  在实现NMS算法之前先看一下IoU是如何计算的：
  IoU就是我们常说的交并比(Intersection over Union, IoU)，顾名思义，就是交集与并集的比值，反映的是两个物体间的重合程度。计算公式如下：$$IoU=\frac{A\cap B}{A\cup B}$$

  根据上图所示，IoU 就等于左边灰色面积与右边灰色面积的比值。
  下面来看一下NMS算法的具体实现：

def nms(bboxes, scores, threshold=0.5):
    x1 = bboxes[:, 0]
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 2]
    y2 = bboxes[:, 3]

    areas = (x2 - x1) * (y2 - y1)

    # 从大到小对应的的索引
    order = scores.argsort()[::-1]

    # 记录输出的bbox
    keep = []
    while order.size > 0:
        i = order[0]
        # 记录本轮最大的score对应的index
        keep.append(i)

        if order.size == 1:
            break

        # 计算当前bbox与剩余的bbox之间的IoU
        # 计算IoU需要两个bbox中最大左上角的坐标点和最小右下角的坐标点
        # 即重合区域的左上角坐标点和右下角坐标点
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        # 如果两个bbox之间没有重合, 那么有可能出现负值
        w = np.maximum(0.0, (xx2 - xx1))
        h = np.maximum(0.0, (yy2 - yy1))
        inter = w * h

        iou = inter / (areas[i] + areas[order[1:]] - inter)

        # 删除IoU大于指定阈值的bbox(重合度高), 保留小于指定阈值的bbox
        ids = np.where(iou <= threshold)[0]
        # 因为ids表示剩余的bbox的索引长度
        # +1恢复到order的长度
        order = order[ids + 1]

    return keep

4. 完整代码

import cv2
import numpy as np
import copy


seed = 10001
np.random.seed(seed)

bounding_boxes = [
        [545, 125, 765, 440],
        [890, 100, 1115, 430],
        [1275, 170, 1490, 490]
    ]

confidence_score = [0.95, 0.98, 0.96]

num_anchor = 10
anchors = copy.deepcopy(bounding_boxes)
scores = copy.deepcopy(confidence_score)


def nms(bboxes, scores, threshold=0.5):
    x1 = bboxes[:, 0]
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 2]
    y2 = bboxes[:, 3]

    areas = (x2 - x1) * (y2 - y1)

    # 从大到小对应的的索引
    order = scores.argsort()[::-1]

    # 记录输出的bbox
    keep = []
    while order.size > 0:
        i = order[0]
        # 记录本轮最大的score对应的index
        keep.append(i)

        if order.size == 1:
            break

        # 计算当前bbox与剩余的bbox之间的IoU
        # 计算IoU需要两个bbox中最大左上角的坐标点和最小右下角的坐标点
        # 即重合区域的左上角坐标点和右下角坐标点
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        # 如果两个bbox之间没有重合, 那么有可能出现负值
        w = np.maximum(0.0, (xx2 - xx1))
        h = np.maximum(0.0, (yy2 - yy1))
        inter = w * h

        iou = inter / (areas[i] + areas[order[1:]] - inter)

        # 删除IoU大于指定阈值的bbox(重合度高), 保留小于指定阈值的bbox
        ids = np.where(iou <= threshold)[0]
        # 因为ids表示剩余的bbox的索引长度
        # +1恢复到order的长度
        order = order[ids + 1]

    return keep


if __name__ == '__main__':
    img = cv2.imread('dijia.png')

    for i in range(num_anchor):
        index = np.random.randint(0, 3)
        offset = np.random.randint(-50, 50, size=4)
        score = np.random.uniform(0.5, 0.9)
        anchors.append(list(bounding_boxes[index] - offset))
        scores.append(round(score, 2))

    anchors = np.asarray(anchors)
    scores = np.asarray(scores)

    keep = nms(anchors, scores, threshold=0.5)
    proposals = anchors[keep]
    proposals_score = scores[keep]

    for i in range(len(proposals)):
        cv2.rectangle(img, pt1=tuple(proposals[i][:2]), pt2=tuple(proposals[i][2:]), color=(0, 255, 0), thickness=2)
        cv2.putText(img, text=str(proposals_score[i]), org=tuple(proposals[i][:2]), fontFace=cv2.FONT_HERSHEY_SCRIPT_SIMPLEX,
                    fontScale=1, color=(255, 0, 255), thickness=2)
    cv2.imshow('dijia', img)
    cv2.waitKey()
    cv2.imwrite('dijia2.png', img)

  代码运行如下：

结束语

  从NMS的实现代码中可以看到主要是对数组的操作，而这部分可以通过GPU进行加速处理，比如PyToch，PaddlePaddle等深度学习框架，其目标检测模块中的NMS算法已经内置实现，可以直接在GPU上使用。