转载自:https://blog.csdn.net/u011436429/article/details/80107042
import cv2
import numpy as np
def nms(bounding_boxes, confidence_score, threshold):
if len(bounding_boxes) == 0:
return [], []
boxes = np.array(bounding_boxes)
score = np.array(confidence_score)
# 取出所有候选区域的左上角和右下角的点坐标
start_x = boxes[:, 0]
start_y = boxes[:, 1]
end_x = boxes[:, 2]
end_y = boxes[:, 3]
picked_boxes = []
picked_scores = []
# 计算每个候选区域的面积
areas = (end_x - start_x + 1) * (end_y - start_y + 1)
# np.argsort(x)为从小到大排序 np.argsort(-x)为从大到小排序
order = np.argsort(score)
# 当order不为空
while order.size > 0:
# 选取最大置信度(score)的index
index = order[-1]
picked_boxes.append(bounding_boxes[index])
picked_score.append(confidence_score[index])
# 计算相交区域左上角和右下角两点坐标
x1 = np.maximum(start_x[index], start_x[order[:-1]])
y1 = np.maximum(start_y[index], start_y[order[:-1]])
x2 = np.minimum(end_x[index], end_x[order[:-1]])
y2 = np.minimum(end_y[index], end_y[order[:-1]])
# 计算相交区域面积
w = np.maximum(0.0, x2 - x1 + 1)
h = np.maximum(0.0, y2 - y1 + 1)
intersection = w * h
# 计算交并比
ratio = intersection / (area[index] + area[order[:-1]] - intersection)
# 保留所有重叠面积小于阈值的框,留作下次处理
left = np.where(ratio < threshold)
order = order[left]
return picked_boxes, picked_score