python numpy极简版代码实现IOU和NMS

搞视觉的出去求职，这道题也是高频出现要求手写的题目。当然搞CV不会IOU/NMS，堪比翟天临博士没听过知网。
开始面试前，对这道题一定要做到必知必会！
讲解见我另外一篇博客：numpy实现目标检测中的IOU和NMS

代码实现版本见：

import numpy as np

def compute_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
    iou = intersection / union
    return iou
def nms(boxes, scores, threshold):
    # 对得分进行降序排序
    sorted_indices = np.argsort(-scores)

    # 初始化结果列表
    result = []

    # 遍历排序后的得分列表
    for i in sorted_indices:
        # 如果当前得分对应的框没有被添加到结果列表中
        if len(result) == 0 or compute_iou(result[-1], boxes[i]) < threshold:
            result.append(boxes[i])

    # 返回结果列表
    return result

def main():
    # 定义一些测试数据
    boxes = np.array([[10, 20, 30, 40], [20, 30, 40, 50], [30, 40, 50, 60], [40, 50, 60, 70]])
    scores = np.array([0.9, 0.8, 0.7, 0.6])
    threshold = 0.5

    # 调用nms函数
    result = nms(boxes, scores, threshold)

    # 打印结果
    print("NMS结果：", result)

if __name__ == "__main__":
    main()