【目标检测】非极大值抑制（Non-Maximum Suppression, NMS）步骤与实现

步骤

1. 置信度排序：首先根据预测框的置信度（即预测框包含目标物体的概率）对所有预测框进行降序排序。
2. 选择最佳预测框：选择置信度最高的预测框作为参考框。
3. 计算IoU：计算其他所有预测框与参考框的交并比（Intersection over Union, IoU）。
4. 抑制：删除IoU高于某个阈值的所有预测框，因为这些框与参考框重叠太多，可能是重复预测。
5. 迭代：从剩余的预测框中选择置信度最高的作为新的参考框，重复步骤3和4，直到所有预测框都被处理。

代码实现

导入必要的库

import numpy as np 

定义计算IoU的函数

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])          # 计算并集面积     unionArea = boxAArea + boxBArea - interArea          # 计算IoU     iou = interArea / unionArea if unionArea != 0 else 0          return iou 

定义NMS函数

def nms(boxes, scores, iou_threshold):     picked = []  # 存储被选择的边界框索引     indexes = np.argsort(scores)[::-1]  # 按分数降序排列索引          while len(indexes) > 0:         current = indexes[0]         picked.append(current)  # 选择当前最高分的边界框         indexes = indexes[1:]  # 移除当前最高分的索引                  # 检查剩余边界框与当前选择框的IoU，如果大于阈值则抑制         indexes = [i for i in indexes if iou(boxes[current], boxes[i]) <= iou_threshold]          return picked 

示例使用

# 假设boxes和scores是模型预测的边界框和分数 boxes = np.array([[50, 50, 100, 100], [60, 60, 110, 110], [200, 200, 300, 300]]) scores = np.array([0.9, 0.75, 0.8])  # 设置IoU阈值 iou_threshold = 0.5  # 执行NMS picked_boxes = nms(boxes, scores, iou_threshold)  print("Selected box indices:", picked_boxes) 

注意

• 边界框通常以(x1, y1, x2, y2)的格式表示，其中(x1, y1)是左上角坐标，(x2, y2)是右下角坐标。

Code

完整示例代码已上传至：Machine Learning and Deep Learning Algorithms with NumPy
此项目包含更多AI相关的算法numpy实现，供大家学习参考使用，欢迎star~

备注

个人水平有限，有问题随时交流~