在CNN分类中使用tensorflow绘制混淆矩阵

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix

# 加载MNIST数据集
mnist = tf.keras.datasets.mnist
(_, _), (test_images, test_labels) = mnist.load_data()

# 预处理数据
test_images = test_images / 255.0
test_images = np.expand_dims(test_images, axis=-1)

# 加载训练好的模型
model = tf.keras.models.load_model('mnist_model.h5')

# 进行预测
predictions = model.predict(test_images)
predicted_labels = np.argmax(predictions, axis=1)

# 绘制混淆矩阵
cm = confusion_matrix(test_labels, predicted_labels)

# 可视化混淆矩阵
plt.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
plt.title('Confusion Matrix - MNIST Classification')
plt.colorbar()
tick_marks = np.arange(10)
plt.xticks(tick_marks, range(10))
plt.yticks(tick_marks, range(10))
plt.xlabel('Predicted Label')
plt.ylabel('True Label')

# 添加数值标签
thresh = cm.max() / 2.0
for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
    plt.text(j, i, format(cm[i, j], 'd'),
             horizontalalignment="center",
             color="white" if cm[i, j] > thresh else "black")

plt.tight_layout()
plt.show()