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)