import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

plt.rcParams['font.family'] = "SimHei"
gpus = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(gpus[0],True)
#导入mnist的训练集和测试集
mnist = tf.keras.datasets.mnist
(train_x,train_y),(test_x,test_y) = mnist.load_data()
img_testx = test_x
img_testy = test_y

#对属性进行归一化，使取值范围在0-1之间，同时转换为tensor张量，标签值转换为张量，0-9的整数
X_train,X_test = tf.cast(train_x / 255.0,tf.float32),tf.cast(test_x / 255.0,tf.float32)
Y_train,Y_test = tf.cast(train_y , tf.int16),tf.cast(test_y,tf.int16)
X_img = X_test
#建立Sequential模型，说明神经网络的结构，使用add方法添加层
model = tf.keras.Sequential()
model.add(tf.keras.layers.Flatten(input_shape=(28,28))) #Flatten不进行计算，将输入的二维数组转换为一维数组，进行形状转换
model.add(tf.keras.layers.Dense(128,activation="relu")) #添加隐含层，隐含层是全连接层，128个结点，激活函数使用relu函数
model.add(tf.keras.layers.Dense(10,activation="softmax"))#添加输出层，输出层是全连接层，激活函数是softmax函数

#配置训练方法
#优化器使用adam，损失函数使用稀疏交叉熵损失函数，准确率使用稀疏分类准确率函数
model.compile(optimizer='adam',loss='sparse_categorical_crossentropy',metrics=['sparse_categorical_accuracy'])

## 训练模型
## 使用训练集中的数据训练，从中划分20%作为测试数据，用在每轮训练后评估模型的性能，每个小批量使用64条数据，训练5轮
history=model.fit(X_train,Y_train,batch_size=64,epochs=5,validation_split=0.2)
model.save_weights("result.h5")

#使用测试集评估模型
model.evaluate(X_test,Y_test,verbose=2)
#结果可视化
print(history.history)
loss = history.history['loss']          #训练集损失
val_loss = history.history['val_loss']  #测试集损失
acc = history.history['sparse_categorical_accuracy']            #训练集准确率
val_acc = history.history['val_sparse_categorical_accuracy']    #测试集准确率

plt.figure(figsize=(10,3))

plt.subplot(121)
plt.plot(loss,color='b',label='train')
plt.plot(val_loss,color='r',label='test')
plt.ylabel('loss')
plt.legend()

plt.subplot(122)
plt.plot(acc,color='b',label='train')
plt.plot(val_acc,color='r',label='test')
plt.ylabel('Accuracy')
plt.legend()

plt.show()

#使用模型预测随机5个数据
for i in range(5):
    num = np.random.randint(1,10000)
    plt.subplot(1,5,i+1)
    plt.axis("off")
    plt.imshow(test_x[num],cmap="gray")
    # argmax取出值最大的索引，predict中参数的数据范围和维数与训练集一致
    y_pred = np.argmax(model.predict(tf.convert_to_tensor(X_test[num].numpy().reshape(1,28,28))))#使用argmax函数得到预测值
    plt.title("原值="+str(test_y[num])+"\n预测值："+str(y_pred))

plt.show()


#使用模型预测自己的手写数据集
img_arr = []
for i in range(10):
    img = Image.open("F:\\code\\%d.png" % i)
    img_temp = np.array(img)
    img_arr.append(img_temp)

for i in range(10):
    num = np.random.randint(1,10000)
    plt.subplot(2,5,i+1)
    plt.axis("off")
    plt.imshow(img_testx[num],cmap = "gray")
    #argmax取出值最大的索引，predict中参数的数据范围和维数与训练集一致
    #y_pred = np.argmax(model.predict([[X_test[num]]]))
    #tensor = tf.convert_to_tensor(X_test[num].numpy().reshape(1,28*28))
    y_pred = np.argmax(model.predict(tf.convert_to_tensor(X_img[num].numpy().reshape(1,28,28))))
    plt.title("原值=" + str(img_testy[num]) + "\n预测值：" + str(y_pred))

plt.show()