轉(zhuǎn)載自: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)
# 取出所有候選區(qū)域的左上角和右下角的點坐標
start_x = boxes[:, 0]
start_y = boxes[:, 1]
end_x = boxes[:, 2]
end_y = boxes[:, 3]
picked_boxes = []
picked_scores = []
# 計算每個候選區(qū)域的面積
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])
# 計算相交區(qū)域左上角和右下角兩點坐標
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]])
# 計算相交區(qū)域面積
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