用Python实现nms

非极大抑制（Non-Maximum Suppression, NMS）是一种常用的算法，用于处理物体检测任务中的重叠边界框问题。下面是一个简单的NMS在Python中的实现步骤[^1]：

1. **导入必要的库**:

import numpy as np

2. **创建一组边界框** (假设每个边界框由左上角坐标、右下角坐标和置信度组成):

boxes = np.array([[100,100,210,210,0.72], [250,250,420,420,0.8], [220,220,320,330,0.92], 
                  [100,100,210,210,0.72], [230,240,325,330,0.81], [220,230,315,340,0.9]])

3. **定义NMS函数**:

def non_max_suppression(boxes, iou_threshold=0.5):
    indices = np.argsort(-boxes[:, -1])  # 按置信度降序排列
    suppressed = set()
    nms_boxes = []
    
    while len(indices) > 0:
        last_index = indices[0]
        nms_boxes.append(boxes[last_index])
        
        # 如果最后一个元素已经抑制过其他元素，则移除
        if last_index in suppressed:
            indices = indices[1:]
            continue
        
        current_box_area = (boxes[last_index, 2] - boxes[last_index, 0]) * (boxes[last_index, 3] - boxes[last_index, 1])
        
        # 计算IoU
        overlap = compute_iou(boxes[last_index], boxes[indices[1:]])  # 假设compute_iou函数已存在

        # 抑制IoU大于阈值的框
        indices = indices[1:].tolist()
        indices = [i for i, j in enumerate(indices) if overlap[i] < iou_threshold]

        # 添加到抑制集合
        suppressed.add(last_index)
        suppressed.update(indices)

    return np.array(nms_boxes)

在这个示例中，`compute_iou`是一个计算两个边界框之间IoU（Intersection over Union）的函数，它不是直接给出的，但通常它是通过计算交集面积除以并集面积来实现的。