Python实现非极大值抑制算法（NMS）

import cv2
import numpy as np


def nms(bounding_boxes, confidence_score, threshold):
    '''
    :param bounding_boxes: 候选框列表，[左上角坐标, 右下角坐标], [min_x, min_y, max_x, max_y], 原点在图像左上角
    :param confidence_score: 候选框置信度
    :param threshold: IOU阈值
    :return: 抑制后的bbox和置信度
    '''
    # 如果没有bbox，则返回空列表
    if len(bounding_boxes) == 0:
        return [], []

    # bbox转为numpy格式方便计算
    boxes = np.array(bounding_boxes)

    # 分别取出bbox的坐标
    start_x = boxes[:, 0]
    start_y = boxes[:, 1]
    end_x = boxes[:, 2]
    end_y = boxes[:, 3]

    # 置信度转为numpy格式方便计算
    score = np.array(confidence_score)  # [0.9  0.75 0.8  0.85]

    # 筛选后的bbox和置信度
    picked_boxes = []
    picked_score = []

    # 计算每一个框的面积
    areas = (end_x - start_x + 1) * (end_y - start_y + 1)

    # 将score中的元素从小到大排列，提取其对应的index(索引)，然后输出到order
    order = np.argsort(score)   # [1 2 3 0]

    # Iterate bounding boxes
    while order.size > 0:

        # The index of largest confidence score
        # 取出最大置信度的索引
        index = order[-1]

        # Pick the bounding box with largest confidence score
        # 将最大置信度和最大置信度对应的框添加进筛选列表里
        picked_boxes.append(bounding_boxes[index])
        picked_score.append(confidence_score[index])

        # 求置信度最大的框与其他所有框相交的长宽，为下面计算相交面积做准备
        # 令左上角为原点，
        # 两个框的左上角坐标x取大值，右下角坐标x取小值，小值-大值+1==相交区域的长度
        # 两个框的左上角坐标y取大值，右下角坐标y取小值，小值-大值+1==相交区域的高度
        # 这里可以在草稿纸上画个图，清晰明了
        x1 = np.maximum(start_x[index], start_x[order[:-1]])
        x2 = np.minimum(end_x[index], end_x[order[:-1]])
        y1 = np.maximum(start_y[index], start_y[order[:-1]])
        y2 = np.minimum(end_y[index], end_y[order[:-1]])

        # 计算相交面积，当两个框不相交时，w和h必有一个为0，面积也为0
        w = np.maximum(0.0, x2 - x1 + 1)
        h = np.maximum(0.0, y2 - y1 + 1)
        intersection = w * h

        # 计算IOU
        ratio = intersection / (areas[index] + areas[order[:-1]] - intersection)

        # 保留小于阈值的框的索引
        left = np.where(ratio < threshold)
        # 根据该索引修正order中的索引（order里放的是按置信度从小到大排列的索引）
        order = order[left]

    return picked_boxes, picked_score


# 图像路径
image_name = 'lena.jpg'

# 自己设定候选框
bounding_boxes = [(210, 180, 337, 380), (180, 120, 330, 340), (270, 160, 350, 360), (220, 210, 345, 410)]
confidence_score = [0.9, 0.75, 0.8, 0.85]

# 读取图像
image = cv2.imread(image_name)

# 复制一份原图矩阵
org = image.copy()

# 画框参数
font = cv2.FONT_HERSHEY_SIMPLEX
font_scale = 1
thickness = 2

# IOU阈值设定
threshold = 0.4

# 画框（未运行NMS）
for (start_x, start_y, end_x, end_y), confidence in zip(bounding_boxes, confidence_score):
    (w, h), baseline = cv2.getTextSize(str(confidence), font, font_scale, thickness)
    cv2.rectangle(org, (start_x, start_y - (2 * baseline + 5)), (start_x + w, start_y), (0, 255, 255), -1)
    cv2.rectangle(org, (start_x, start_y), (end_x, end_y), (0, 255, 255), 2)
    cv2.putText(org, str(confidence), (start_x, start_y), font, font_scale, (0, 0, 0), thickness)

# 运行NMS算法
picked_boxes, picked_score = nms(bounding_boxes, confidence_score, threshold)

# 画框（运行了NMS后）
for (start_x, start_y, end_x, end_y), confidence in zip(picked_boxes, picked_score):
    (w, h), baseline = cv2.getTextSize(str(confidence), font, font_scale, thickness)
    cv2.rectangle(image, (start_x, start_y - (2 * baseline + 5)), (start_x + w, start_y), (0, 255, 255), -1)
    cv2.rectangle(image, (start_x, start_y), (end_x, end_y), (0, 255, 255), 2)
    cv2.putText(image, str(confidence), (start_x, start_y), font, font_scale, (0, 0, 0), thickness)

# 展示图像
cv2.imshow('Original', org)
cv2.imshow('NMS', image)
cv2.waitKey(0)