绘制混淆矩阵热力图

import seaborn as sns
from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
def plot_confusion_matrix(pred_index, targets, labels, save_file, show_number=100):
    sns.set()
    
    y_true = targets
    y_pred = pred_index
    C2= confusion_matrix(y_pred, y_true, labels=labels)
    print("C2.shape", C2.shape) #打印出来看看
    has_targets = np.sum(C2, axis=0)
    has_targets = np.where(has_targets > 0)[0]
    has_preds = np.sum(C2, axis=1)
    has_preds = np.where(has_preds > 0)[0]
    C = C2[:, has_targets]
    C = C[has_preds, :]
    
    x_ticks = [i+0.5 for i in range(has_targets.shape[0])]
    x_labels= has_targets.tolist()
    
    y_ticks = [i+0.5 for i in range(has_preds.shape[0])]
    y_labels=has_preds.tolist()
    
    # import pdb; pdb.set_trace()
    f,ax=plt.subplots(1, 1,  figsize=(8+len(x_labels)//2, 6+len(y_labels)//2))
    print("real shape:", len(x_labels), len(y_labels))
    
    colors = [(0, "white"),(0.1, 'green'), (0.5, "yellow"), (1, "red"), ]
    cmap = mcolors.LinearSegmentedColormap.from_list('custom_colormap', colors)
    sns.heatmap(C, annot=True, annot_kws={"size": 10, 'weight':'bold'}, ax=ax, cmap=cmap, fmt=".0f", linewidths=0, cbar=True) #画热力图
    # sns.heatmap(C, annot=True, annot_kws={"size": 10, 'weight':'bold'}, ax=ax, cmap=plt.cm.Blues, fmt=".0f", linewidths=0, cbar=True) #画热力图
    # sns.heatmap(C3, linewidths = 0.01, annot=True, ax=ax) #画热力图
    ax.set_title('confusion matrix') #标题
    ax.set_xlabel('true') #x轴
    ax.set_xticks(ticks=x_ticks, labels=x_labels)
    ax.set_ylabel('predict') #y轴
    ax.set_yticks(ticks=y_ticks, labels=y_labels)
    
    print("generate finished") #打印出来看看
    plt.savefig(save_file, format='png')
    print("save png finished") #打印出来看看
    # import pdb; pdb.set_trace()


save_file = out_file_path.replace(".pkl", ".png")
plot_confusion_matrix(pred_indexs, targets, [i for i in range(class_number)], save_file)