python 利用非极大值抑制处理重叠框

在目标检测任务中，常常会出现多个重叠的框，为了提高检测准确率，需要对这些框进行非极大值抑制（Non-Maximum Suppression，NMS）处理，只保留置信度最高的框，去除其他重叠的框。

下面是一个简单的利用NMS处理重叠框的Python示例代码，假设已经有一个检测结果列表detections，每个元素为一个四元组(x1,y1,x2,y2,score)，表示一个框的左上角坐标（x1,y1）、右下角坐标（x2,y2）和置信度得分score:

def nms(detections, threshold):
    # 按照置信度得分从大到小排序
    detections = sorted(detections, key=lambda x: x[4], reverse=True)
    # 用于保存最终保留的框
    keep = []
    while detections:
        # 取得置信度最高的框
        max_detection = detections.pop(0)
        keep.append(max_detection)
        # 计算当前框与其他框的重叠面积比例
        overlaps = [iou(max_detection, detection) for detection in detections]
        # 保留重叠面积比例小于阈值的框
        detections = [detection for i, detection in enumerate(detections) if overlaps[i] < threshold]
    return keep

def iou(box1, box2):
    # 计算两个框的重叠面积比例
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])
    intersection = max(0, x2 - x1) * max(0, y2 - y1)
    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
    union = area1 + area2 - intersection
    iou = intersection / union
    return iou

在上面的代码中，nms函数接收两个参数，detections是检测结果列表，threshold是重叠面积比例的阈值，如果两个框的重叠面积比例大于等于threshold，则认为它们重叠了。函数先将检测结果按照置信度得分从大到小排序，然后依次取得置信度最高的框，将它保留，计算它与其他框的重叠面积比例，去除重叠面积比例大于等于threshold的框，直到检测结果列表为空。最后返回保留下来的框列表。

iou函数计算两个框的重叠面积比例，它接收两个参数，box1和box2分别表示两个框的左上角坐标和右下角坐标。函数首先计算两个框的交集面积，然后计算它们的并集面积，最后返回交集面积与并集面积的比例。

以上就是一个简单的利用NMS处理重叠框的Python示例代码，可以根据自己的需求进行修改。