利用混淆矩阵查看每一类预测结果+miou计算

混淆矩阵的示意图如下：

import numpy as np
cm=np.array([[4,0,0,0],
            [0,1,0,0],
             [0, 0, 2, 0],
             [0, 1, 3, 0]])
plt.figure()
plt.grid(False)
plt.imshow(cm, cmap='jet')
plt.colorbar()
plt.show()

import matplotlib.pyplot as plt
import seaborn as sn
import numpy as np
from sklearn.metrics import confusion_matrix

_pred = np.array([[0, 1, 2],
                  [1, 2, 3],
                  [1, 1, 2]])
_gt = np.array([[0, 1, 1],
                [2, 2, 3],
                [1, 3, 2]])
flat_pred = _pred.flatten().astype('uint8')
print('flat_pred=', flat_pred)
flat_true = _gt.flatten().astype('uint8')
print('flat_true=', flat_true)
label_class = [x for x in range(4)]
# print(label_class)
confusion = confusion_matrix(flat_true, flat_pred, label_class)
print('confusion=', confusion)
list_diag = np.diag(confusion)
print('list_diag=', list_diag)
list_raw_sum = np.sum(confusion, axis=1)
print('list_raw_sum=',list_raw_sum)
each_acc = np.nan_to_num(list_diag.astype('Float32') / list_raw_sum.astype('Float32'))
print('each_acc=', each_acc)
ave_acc=np.mean(each_acc)
print('ave_acc=',ave_acc)


ax = plt.axes()
class_names = ['mountain', 'street', 'glacier', 'person']
sn.heatmap(confusion, annot=True,
           annot_kws={"size": 10},
           xticklabels=class_names,
           yticklabels=class_names, ax=ax)
ax.set_title('Confusion matrix')
plt.savefig('./confusion.jpg')
plt.show()

miou计算:

def gen_matrix(gt_mask, pred_mask, class_num):
    """
    gt_mask(ndarray): shape -> (height, width), 真实的分割图
    pred_mask(ndarray):shape -> (height, width), 预测的分割结果
    class_num: 类别数目，不包含背景
    """
    mask = (gt_mask >= 0) & (gt_mask < class_num)
    print('===mask:', mask)
    print('===pred_mask[mask]:', pred_mask[mask])
    print('==class_num * gt_mask[mask].astype(int):', class_num * gt_mask[mask].astype(int))
    print(class_num * gt_mask[mask].astype(int) + pred_mask[mask])
    # bincount为计数函数，将数组从小到大排序后计数，默认从0到数组最大值计数。
    count = np.bincount(class_num * gt_mask[mask].astype(int) + pred_mask[mask], minlength=class_num ** 2)

    print('==count:', count)
    # 混淆矩阵
    cf_mtx = count.reshape(class_num, class_num)
    print('==cf_mtx:', cf_mtx)
    return cf_mtx


def mean_iou(cf_mtx):
    """
    cf_mtx(ndarray): shape -> (class_num, class_num), 混淆矩阵
    """
    #
    mIou = np.diag(cf_mtx) / (np.sum(cf_mtx, axis=1) + \
                              np.sum(cf_mtx, axis=0) - np.diag(cf_mtx))
    print('===mIou:', mIou)
    # 所有类别iou取平均
    mIou = np.nanmean(mIou)
    return mIou

gt_mask = np.array([[0, 1, 1],
                    [1, 1, 0],
                    [1, 1, 0]])
pred_mask = np.array([[0, 1, 1],
                      [1, 0, 0],
                      [1, 0, 1]])
class_num = 2
cx_matrix = gen_matrix(gt_mask, pred_mask, class_num)
miou = mean_iou(cx_matrix)
print('==miou:', miou)