【语义分割】——计算IOU

参考自：https://github.com/dilligencer-zrj/code_zoo/blob/master/compute_mIOU

python版本
直接就是对每个类别进行求解交集和并集

numpy版本
采用numpy的bitcount分布。这里的分布做了一种变换target × nclass + pred，这样预测正确的像素点都在hist矩阵的对角线上。

import numpy as np
from numpy.lib.twodim_base import diag

size = (20, 30)
classes = 19
pred = np.random.randint(0, classes, size)
target = np.random.randint(0, classes, size)
pred.flatten()

# 计算iou
def IOU(pred, target, nclass):
    ious = []
    for i in range(classes):
        pred_ins = pred == i 
        target_ins = target == i
        inser = pred_ins[target_ins].sum()
        union = pred_ins.sum() + target_ins.sum() - inser
        iou = inser / union
        ious.append(iou)
    
    return ious

# 采用numpy的函数bincount
def IOU2(pred, target, classes):

    hist = np.zeros((classes, classes))
    hist = np.bincount(pred.flatten() * classes + target.flatten())
    hist = np.reshape(hist, (classes, classes))
    ious = np.diag(hist) / (hist.sum(0) + hist.sum(1) - np.diag(hist))

    return ious


if __name__ == "__main__":
    ious = IOU(pred, target, classes)
    ious2 = IOU2(pred, target, classes)
    print(ious)
    print(ious2)

reference

1. https://blog.csdn.net/xlinsist/article/details/51346523