【源码阅读-CV】IoU计算

import numpy as np


def matrix_iou(a, b):
    # a: N1x4 b:N2x4
    #            ——————————
    #           ｜         ｜
    #   --------｜---------｜-----
    #   |       ｜ P1      ｜     |
    #   |       ｜         ｜     |
    #   |        ——————————- P2   |
    #   --------------------------
    # 计算出两个框相交位于最内侧的左上角点P1
    lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])  # N1xN2x2
    # 计算出两个框相交位于最内侧的右下角点P2
    rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])  # N1xN2x2

    # 当lt>rb时，也就是两个框不相交的时候，面积等于0
    # 等价 area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)
    wh = (rb-lt).clip(min=0)
    area_i = wh[:, :, 0] * wh[:, :, 1]  # N1xN2
    area_a = np.prod(a[:, 2:] - a[:, :2], axis=1) # N1
    area_b = np.prod(b[:, 2:] - b[:, :2], axis=1) # N2
    # area_a[:, np.newaxis] + area_b把结果变成N1xN2
    return area_i / (area_a[:, np.newaxis] + area_b - area_i) # N1xN2


a = np.array([[0, 0, 4, 4],
              [0, 0, 5, 5]])
b = np.array([[0, 0, 4, 4],
              [0, 0, 5, 5]])
print(matrix_iou(a, b))

#######
[[1.   0.64]
 [0.64 1.  ]]