python绘制混淆矩阵

test_generator = test_datagen.flow_from_directory(
        'dataset/test',
        target_size=(48, 48),
        shuffle = False ,#随机打乱默认为true
        # batch_size=16,
        color_mode="grayscale",
        class_mode = 'categorical')
predictions = model.predict_generator(test_generator)
predicted_classes = np.argmax(predictions, axis=1)
true_classes = test_generator.classes
labels = ['Anger', 'Disgust', 'Fear', 'Happy', 'Sad', 'Surprise', 'Neutral']
confusion_matrix = confusion_matrix(true_classes, predicted_classes)

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator


def plotCM(classes, matrix, savname):
    """classes: a list of class names"""
    # Normalize by row
    matrix = matrix.astype(np.float)
    # linesum = matrix.sum(1)
    # linesum = np.dot(linesum.reshape(-1, 1), np.ones((1, matrix.shape[1])))
    # matrix /= linesum
    # plot
    plt.switch_backend('agg')
    fig = plt.figure()
    ax = fig.add_subplot(111)
    cax = ax.matshow(matrix)
    fig.colorbar(cax)
    ax.xaxis.set_major_locator(MultipleLocator(1))
    ax.yaxis.set_major_locator(MultipleLocator(1))
    for i in range(matrix.shape[0]):
        for j in range(matrix.shape[1]):
            ax.text(j, i, str('%.2f' % (matrix[i, j])), va='center', ha='center')
    ax.set_xticklabels([''] + classes, rotation=90)
    ax.set_yticklabels([''] + classes)
    # save
    plt.savefig(savname)
plotCM(labels, confusion_matrix, 'matrix.jpg')