您现在的位置是：首页 >技术杂谈 >不同数量的预测框和Ground Truth框计算IoU网站首页技术杂谈

import numpy as np

def calculate_iou(boxes1, boxes2):
    # 转换为 numpy 数组
    boxes1 = np.array(boxes1)
    boxes2 = np.array(boxes2)
    
    # 扩展维度，以便广播计算
    boxes1 = np.expand_dims(boxes1, axis=1)
    boxes2 = np.expand_dims(boxes2, axis=0)
    
    # 计算两组框的交集坐标范围
    x_min = np.maximum(boxes1[:, :, 0], boxes2[:, :, 0])
    y_min = np.maximum(boxes1[:, :, 1], boxes2[:, :, 1])
    x_max = np.minimum(boxes1[:, :, 2], boxes2[:, :, 2])
    y_max = np.minimum(boxes1[:, :, 3], boxes2[:, :, 3])
    
    # 计算交集和并集的面积
    intersection = np.maximum(x_max - x_min, 0) * np.maximum(y_max - y_min, 0)
    union = (boxes1[:, :, 2] - boxes1[:, :, 0]) * (boxes1[:, :, 3] - boxes1[:, :, 1]) + 
            (boxes2[:, :, 2] - boxes2[:, :, 0]) * (boxes2[:, :, 3] - boxes2[:, :, 1]) - 
            intersection
    
    # 计算 IoU
    iou = intersection / np.maximum(union, 1e-8)
    
    return iou

这个代码使用了 NumPy 的广播（broadcasting）机制，它可以使得不同形状的数组在某些维度上进行“自动”扩展，从而使它们可以参与到同样形状的计算中。这种机制可以避免我们手动对数组进行复制或循环操作，从而减少计算的时间和空间消耗。

具体来说，在这个代码中，我们首先将两个数组 boxes1 和 boxes2 转换为 numpy 数组，并对它们的形状进行扩展，以便后续计算。这里使用了 np.expand_dims() 函数，将 boxes1 扩展为 N1 × 1 × 4 的三维数组，将 boxes2 扩展为 1 × N2 × 4 的三维数组。这样，当我们进行计算时，NumPy 就会自动将这两个数组沿着扩展的维度进行广播，从而得到一个 N1 × N2 × 4 的数组，它包含了所有可能的框对之间的坐标范围信息。

接下来，我们对这个三维数组进行切片操作，得到所有可能的框对之间的坐标范围信息。具体来说，我们使用切片 boxes1[:, :, 0] 获取所有预测框的左上角 x 坐标，使用切片 boxes2[:, :, 0] 获取所有 ground truth 框的左上角 x 坐标。然后，使用 NumPy 的 np.maximum() 函数对这两个切片进行按位取最大值，得到两组框的交集左上角 x 坐标。这个操作会得到一个 N1 × N2 的二维数组，它包含了所有可能的框对之间的交集左上角 x 坐标。

类似地，我们对所有可能的框对之间的坐标范围信息进行处理，得到交集和并集的宽度和高度，并使用它们计算交集和并集的面积。最后，我们使用 NumPy 的 np.maximum() 函数计算 IoU，并得到一个 N1 × N2 的二维数组，它包含了所有可能的框对之间的 IoU 值。

通过这种方式，我们可以高效地对不同数量的预测框和 ground truth 框进行 IoU 计算，而不需要手动对它们进行复制或循环操作。这可以大大提高计算的效率，尤其是当预测框和 ground truth 框的数量很大时。