import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap, Normalize

conf_arr = [[33, 2, 0, 0, 0, 0, 0, 0, 0, 1, 3],
            [3, 31, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 4, 41, 0, 0, 0, 0, 0, 0, 0, 1],
            [0, 1, 0, 30, 0, 6, 0, 0, 0, 0, 1],
            [0, 0, 0, 0, 38, 10, 0, 0, 0, 0, 0],
            [0, 0, 0, 3, 1, 39, 0, 0, 0, 0, 4],
            [0, 2, 2, 0, 4, 1, 31, 0, 0, 0, 2],
            [0, 1, 0, 0, 0, 0, 0, 36, 0, 2, 0],
            [0, 0, 0, 0, 0, 0, 1, 5, 37, 5, 1],
            [3, 0, 0, 0, 0, 0, 0, 0, 0, 39, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 38]]

# 自定义渐变颜色
colors = ['#ffffff', '#ffcccc', '#ff6666', '#ff0000', '#990000']
custom_cmap = LinearSegmentedColormap.from_list('custom_cmap', colors)

# 归一化对象，定义最小值和最大值
norm = Normalize(vmin=0, vmax=np.max(conf_arr))

norm_conf = []
for i in conf_arr:
    a = 0
    tmp_arr = []
    a = sum(i, 0)
    for j in i:
        tmp_arr.append(float(j)/float(a))
    norm_conf.append(tmp_arr)

fig = plt.figure()
plt.clf()
ax = fig.add_subplot(111)
ax.set_aspect(1)
res = ax.imshow(np.array(conf_arr), cmap=custom_cmap, norm=norm, interpolation='nearest')

# 获取矩阵的宽度和高度
width, height = np.array(conf_arr).shape

# 使用 range 替换 xrange
for x in range(width):
    for y in range(height):
        ax.annotate(str(conf_arr[x][y]), xy=(y, x),
                    horizontalalignment='center',
                    verticalalignment='center',
                    color='black')

cb = fig.colorbar(res)

# 调整 alphabet 使其不会超出索引范围
alphabet = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'[:width]
plt.xticks(range(width), alphabet)
plt.yticks(range(height), alphabet)

plt.savefig('confusion_matrix.png', format='png')
plt.show()

效果图如下：