sklearn混淆矩阵的计算和seaborn可视化

扭扭小薯条2024-04-30 5:00

为了计算语义分割的指标miou，需要生成的中间过程就是混淆矩阵。

iou = intersection / union

每个类别的平均iou就是mean iou。

使用sklearn自带的confusion_matrix能很容易生成混淆矩阵，可以进行混淆矩阵的可视化观察哪个类别分割的不好。

from sklearn.metrics import confusion_matrix

# 定义总的混淆矩阵
matrix = np.zeros((self.n, self.n), dtype=np.int64)

# 每个batch的数据都按照一下的方法添加进混淆矩阵:
# pred:  NxC
# label: N

pred_l = pred.max(dim=1)[1]
matrix += confusion_matrix(label.int().cpu().numpy(), pred_l.cpu().numpy(), labels=range(13))

一定要设置 labels=range(13)，否则labels不全可能生成不了完整尺寸的混淆矩阵。

可以使用seaborn进行可视化。

import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt


sns.heatmap(data=matrix / np.sum(matrix, axis=0, keepdims=True), 
            annot=True,
            fmt=".2f", 
            xticklabels=class_name,
            cmap='GnBu')
plt.xticks(rotation=45, ha='right')
plt.ylabel('Pred')
plt.xlabel('True')
plt.title('Confusion Matrix')

plt.tight_layout()

plt.savefig("confusion.png")

matrix按照列进行求和，这样能够按照label值归一化到1。