NMS非极大值抑制python实现

1、使用python实现

import numpy as np

def iou(box1,box2):

    """
    计算两个边界框的IOU
    
    param box1: [x, y, width, height]
    param box2: [x, y, width, height]
    return: IOU值
    """

    # 计算相交区域的左上和右下坐标
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[0] + box1[2], box2[0] + box2[2])
    y2 = min(box1[1] + box1[3], box2[1] + box2[3])
 
    # 计算相交区域面积
    intersection_area = max(0, x2 - x1) * max(0, y2 - y1)
 
    # 计算两个框的总面积
    box1_area = box1[2] * box1[3]
    box2_area = box2[2] * box2[3]
 
    # 计算IOU
    iou = intersection_area / (box1_area + box2_area - intersection_area)


    return iou

def non_max_suppression(boxes,threshold):
    """
    非极大值抑制
    
    param boxes: 包含边界框和相应分数的列表，其中边界框是[x, y, width, height]，分数为float
    param threshold: 用于判定BOX是否重叠的IOU阈值
    return: 经过NMS处理后的边框列表
    """
    if len(boxes) == 0:
        return []

    boxes = sorted(boxes, key=lambda x:x[4], reverse=True)
    processed_boxes = []

    while len(boxes) > 0:
        # 取得分数最高的框
        box_with_max_score = boxes.pop(0)
        processed_boxes.append(box_with_max_score)

        # 移除与当前分数最高的框IOU大于阈值的框
        for i in range(len(boxes)-1, -1, -1):
            if iou(box_with_max_score, boxes[i]) > threshold:
                boxes.pop(i)

    return processed_boxes

# 示例使用
boxes = [
    [100, 100, 200, 200, 0.7],
    [150, 150, 250, 250, 0.8],
    [120, 120, 220, 220, 0.6]
]
processed_boxes = non_max_suppression(boxes, threshold=0.5)
print(processed_boxes)

2、使用python调用opencv实现

import cv2
 
# 假设已经通过目标检测网络得到了boxes和scores
boxes = [[100, 100, 200, 200], [50, 50, 150, 150]]  # 示例坐标
scores = [0.9, 0.7]  # 示例置信度
 
# 应用非极大值抑制
confidenceThreshold = 0.5  # 置信度阈值
nmsThreshold = 0.3  # NMS阈值
indices = cv2.dnn.NMSBoxes(boxes, scores, confidenceThreshold, nmsThreshold)
 
# 输出被NMS保留下来的边界框索引
print("Surviving boxes:", indices)