候选框非极大值抑制Python

以下是候选框非极大值抑制的Python实现：

def non_max_suppression(boxes, scores, threshold):
    # 按照得分值降序排列所有候选边界框
    sorted_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
    selected_indices = []

    while sorted_indices:
        # 选择得分最高的边界框，并将其添加到最终的结果列表中
        best_index = sorted_indices[0]
        selected_indices.append(best_index)

        # 计算当前选择的边界框与其他候选边界框的重叠区域
        best_box = boxes[best_index]
        other_indices = sorted_indices[1:]
        overlaps = [calculate_overlap(best_box, boxes[i]) for i in other_indices]

        # 从剩余的候选边界框中移除重叠区域高于阈值的边界框
        sorted_indices = [i for i, overlap in zip(other_indices, overlaps) if overlap < threshold]

    # 返回最终的结果边界框列表
    selected_boxes = [boxes[i] for i in selected_indices]
    return selected_boxes

def calculate_overlap(box1, box2):
    # 计算两个边界框的重叠区域（交并比IoU）
    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])

    overlap = intersection / (area1 + area2 - intersection)
    return overlap