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

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

news 2024/9/21 6:37:15

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

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

NMS的步骤

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

NMS的Python实现

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

import numpy as npdef 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 ioudef 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)# 计算与其他边界框的IoUfor i in range(1, len(idxs)):if iou(boxes[current], boxes[idxs[i]]) > iou_threshold:breakelse:  # 如果没有超过阈值的框，继续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)