nms python代码_NMS python实现

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