NMS python实现

转载自：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