> 文章列表 > 不同数量的预测框和Ground Truth框计算IoU

不同数量的预测框和Ground Truth框计算IoU

不同数量的预测框和Ground Truth框计算IoU

import numpy as npdef 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# 计算 IoUiou = intersection / np.maximum(union, 1e-8)return iou

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

具体来说,在这个代码中,我们首先将两个数组 boxes1boxes2 转换为 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 框的数量很大时。