import numpy as np

import matplotlib.pyplot as plt

class DrawConfusionMatrix:

     def __init__(self, labels_name, normalize=True):

         """

normalize ：是否设元素为百分比形式

         """

         self.normalize = normalize

         self.labels_name = labels_name

         self.num_classes = len(labels_name)

         self.matrix = np.zeros((self.num_classes, self.num_classes), dtype="float32")

     def update(self, predicts, labels):

         """

         :param predicts: 一维预测向量， eg ： array([0,5,1,6,3,...],dtype=int64)

         :param labels:    一维标签向量： eg ： array([0,5,0,6,2,...],dtype=int64)

         :return:

         """

         for predict, label in zip(predicts, labels):

             self.matrix[predict, label] += 1

     def getMatrix(self,normalize=True):

         """

         根据传入的 normalize 判断要进行 percent 的转换，

         如果 normalize 为 True ，外汇跟单gendan5.com则矩阵元素转换为百分比形式，

         如果 normalize 为 False ，则矩阵元素就为数量

         Returns: 返回一个以百分比或者数量为元素的矩阵

         """

         if normalize:

             per_sum = self.matrix.sum(axis=1)  # 计算每行的和，用于百分比计算

             for i in range(self.num_classes):

                 self.matrix[i] =(self.matrix[i] / per_sum[i])   # 百分比转换

             self.matrix=np.around(self.matrix, 2)   # 保留 2 位小数点

             self.matrix[np.isnan(self.matrix)] = 0  # 可能存在 NaN ，将其设为 0

         return self.matrix

     def drawMatrix(self):

         self.matrix = self.getMatrix(self.normalize)

         plt.imshow(self.matrix, cmap=plt.cm.Blues)  # 仅画出颜色格子，没有值

         plt.title("Normalized confusion matrix")  # title

         plt.xlabel("Predict label")

         plt.ylabel("Truth label")

         plt.yticks(range(self.num_classes), self.labels_name)  # y 轴标签

         plt.xticks(range(self.num_classes), self.labels_name, rotation=45)  # x 轴标签

         for x in range(self.num_classes):

             for y in range(self.num_classes):

                 value = float(format('%.2f' % self.matrix[y, x]))  # 数值处理

                 plt.text(x, y, value, verticalalignment='center', horizontalalignment='center')  # 写值

         plt.tight_layout()  # 自动调整子图参数，使之填充整个图像区域

         plt.colorbar()  # 色条

         plt.savefig('./ConfusionMatrix.png', bbox_inches='tight')  # bbox_inches='tight' 可确保标签信息显示全

         plt.show()