当前位置：首页 > news >正文

目标检测中的非极大值抑制（NMS）：步骤与Python实现

news 来源：原创 2025/5/14 17:44:06

目标检测中的非极大值抑制（NMS）：步骤与Python实现

非极大值抑制（Non-Maximum Suppression, NMS）是目标检测中用于去除多余边界框的关键步骤。在目标检测任务中，模型通常会为每个目标预测多个边界框，这些边界框可能会有重叠。NMS通过保留最佳的边界框并移除其他重叠的边界框来减少冗余。

NMS的步骤

置信度排序：按照每个边界框的置信度（通常是物体存在的得分）进行降序排序。
选择最高得分的边界框：从排序后的列表中选择置信度最高的边界框，将其添加到最终的检测结果中。
计算交并比（IoU）：对于剩余的每个边界框，计算它与已选择边界框的IoU。
移除重叠边界框：如果某个边界框与已选择的边界框的IoU超过了预设的阈值（如0.5），则认为这两个框表示的是同一个目标，移除置信度较低的边界框。
重复：继续选择剩余边界框中置信度最高的，重复计算IoU和移除过程，直到所有边界框都被检查过。

NMS的Python实现

以下是使用Python和NumPy库实现NMS的示例代码：

import numpy as np

def iou(boxA, boxB):
    # 计算两个边界框的交集
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])
    interArea = max(0, xB - xA) * max(0, yB - yA)
    boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
    boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])
    iou = interArea / float(boxAArea + boxBArea - interArea)
    return iou

def non_max_suppression(boxes, scores, iou_threshold=0.5):
    """
    非极大值抑制
    boxes: 边界框列表，每个框是一个[x1, y1, x2, y2]的数组
    scores: 每个边界框对应的得分列表
    iou_threshold: IoU阈值
    """
    # 按得分排序
    idxs = np.argsort(scores)[::-1]
    selected = []

    while len(idxs) > 0:
        # 选取得分最高的边界框
        current = idxs[0]
        selected.append(current)
        # 计算与其他边界框的IoU
        for i in range(1, len(idxs)):
            if iou(boxes[current], boxes[idxs[i]]) > iou_threshold:
                break
        else:  # 如果没有超过阈值的框，继续
            idxs = idxs[1:]
    return selected

# 示例数据
boxes = np.array([[100, 100, 210, 210], [120, 120, 220, 220], [150, 150, 250, 250]])
scores = np.array([0.9, 0.8, 0.7])
selected_boxes = non_max_suppression(boxes, scores, iou_threshold=0.5)
print("Selected boxes:", selected_boxes)