使用来自 Sklearn 的 ConfusionMatrixDisplay 对用于绘制混淆矩阵的颜色图进行归一化

Question

I am trying to create a color map for my 10x10 confusion matrix that is provided by sklearn . I would like to be able to customize the color map to be normalized between [0,1] but I have had no success. I am trying to use ax_ and matplotlib.colors.Normalize but am struggling to get something to work since ConfusionMatrixDisplay is a sklearn object that creates a different than usual matplotlib plot.

My code is the following:

from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix

train_confuse_matrix = confusion_matrix(y_true = ytrain, y_pred = y_train_pred_labels)
print(train_confuse_matrix)
cm_display = ConfusionMatrixDisplay(train_confuse_matrix, display_labels = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'])
print(cm_display)
cm_display.plot(cmap = 'Greens')
plt.show()
plt.clf()


[[3289   56   84   18   55    7   83   61   48  252]
 [   2 3733    0    1    2    1   16    1    3  220]
 [  81   15 3365   64   81   64  273   18    6   17]
 [  17   37   71 3015  127  223  414   44    6   64]
 [   3    1   43   27 3659   24  225   35    0    3]
 [   5   23   38  334  138 3109  224   80    4   25]
 [   3    1   19   10   12    7 3946    1    1    5]
 [   4    7   38   69  154   53   89 3615    2   27]
 [  62   67   12    7   25    3   62    4 3595  153]
 [   2   30    1    2    4    0   15    2    0 3957]]

Answer 1

Let's try imshow and annotate manually:

accuracies = conf_mat/conf_mat.sum(1)
fig, ax = plt.subplots(figsize=(10,8))
cb = ax.imshow(accuracies, cmap='Greens')
plt.xticks(range(len(classes)), classes,rotation=90)
plt.yticks(range(len(classes)), classes)

for i in range(len(classes)):
    for j in range(len(classes)):
        color='green' if accuracies[i,j] < 0.5 else 'white'
        ax.annotate(f'{conf_mat[i,j]}', (i,j), 
                    color=color, va='center', ha='center')

plt.colorbar(cb, ax=ax)
plt.show()

Output: