import os
from PIL import Image
import numpy as np
from keras import regularizers
from keras.utils import np_utils
from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Flatten
from keras.optimizers import SGD, RMSprop, Adam
from keras.layers import Conv2D, MaxPooling2D
from keras.models import load_model
import matplotlib.pyplot as plt

# 提取图像、标签
def x_read_image(filepath):
    img = Image.open('./FA/train/' + filepath).convert('RGB')
    return np.array(img)


# 提取图像、标签
def y_read_image(filepath):
    img = Image.open('./FA/test/' + filepath).convert('RGB')
    return np.array(img)


x_train = []  # 图像
y_train = []  # 标签
for image in os.listdir('./FA/train/'):
    x_train.append(x_read_image(image))
for filepath in os.listdir('./FA/train/'):
    y_train.append(int(filepath.split('_')[0]))
x_train = np.array(x_train)
y_train = np.array(y_train)

x_test = []  # 图像
y_test = []  # 标签
for images2 in os.listdir('./FA/test/'):
    x_test.append(y_read_image(images2))
for filepaths in os.listdir('./FA/test/'):
    y_test.append(int(filepaths.split('_')[0]))
x_test = np.array(x_test)
y_test = np.array(y_test)

# 处理标签
y_train = np_utils.to_categorical(y_train)  # 转为one-hot编码
y_test = np_utils.to_categorical(y_test)
# 归一化
x_train = x_train.astype('float32')
x_train /= 255
x_test = x_test.astype('float32')
x_test /= 255
# 打乱训练集
np.random.seed(200)
np.random.shuffle(x_train)
np.random.seed(200)
np.random.shuffle(y_train)
# CNN
model = Sequential()
# 使用 64 个大小为 3x3 的卷积滤波器。
model.add(Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3), kernel_regularizer=regularizers.l2(0.03)))
model.add(MaxPooling2D(pool_size=(2, 2)))

model.add(Conv2D(16, (3, 3), activation='relu', kernel_regularizer=regularizers.l2(0.03)))
model.add(MaxPooling2D(pool_size=(2, 2)))

model.add(Flatten())
model.add(Dense(16, activation='tanh'))  # Dense(64) 是一个具有 64 个隐藏神经元的全连接层。
model.add(Dropout(0.5))
model.add(Dense(10, activation='softmax'))

# model = load_model('FA4.h5')
adam = Adam(lr=0.0001, beta_1=0.9, beta_2=0.999, epsilon=None, decay=0.0, amsgrad=False)
sgd = SGD(lr=0.002, momentum=0.0, decay=0.0, nesterov=False)
model.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=['accuracy'])

history = model.fit(x_train, y_train, batch_size=32, epochs=512, validation_data=(x_test, y_test))
model.save('FA4.h5')

# # 测试
# score = model.evaluate(x_test, y_test, batch_size=10)
# print(score)
# print(model.summary())
# 绘制训练 & 验证的准确率值
plt.plot(history.history['acc'])
plt.plot(history.history['val_acc'])
plt.title('Model accuracy')
plt.ylabel('Accuracy')
plt.xlabel('Epoch')
plt.legend(['Train', 'Test'], loc='upper left')
plt.show()

# 绘制训练 & 验证的损失值
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('Model loss')
plt.ylabel('Loss')
plt.xlabel('Epoch')
plt.legend(['Train', 'Test'], loc='upper left')
plt.show()