import keras
from keras.models import Sequential
from keras.layers import Conv2D,MaxPool2D,Dense,Dropout,Flatten
from keras.optimizers import Adam
from keras.preprocessing.image import load_img,ImageDataGenerator,img_to_array

# 定义模型

model = Sequential()

model.add(Conv2D(input_shape=(150,150,3),filters=32,kernel_size=3,padding='same',activation='relu'))
model.add(Conv2D(filters=32,kernel_size=3,padding='same',activation='relu'))
model.add(MaxPool2D(pool_size=2,strides=2)) # 池化窗口2*2 步长2

model.add(Conv2D(filters=64,kernel_size=3,padding='same',activation='relu'))
model.add(Conv2D(filters=64,kernel_size=3,padding='same',activation='relu'))
model.add(MaxPool2D(pool_size=2,strides=2))

model.add(Conv2D(filters=128,kernel_size=3,padding='same',activation='relu'))
model.add(Conv2D(filters=128,kernel_size=3,padding='same',activation='relu'))
model.add(MaxPool2D(pool_size=2,strides=2))

model.add(Flatten())
model.add(Dense(64,activation='relu'))
model.add(Dense(2,activation='softmax'))

adam = Adam(lr=1e-4)

model.compile(optimizer=adam,loss='categorical_crossentropy',metrics=['accuracy'])

# 增加图片种类

train_datagen = ImageDataGenerator(
    rotation_range = 40,     # 随机旋转度数
    width_shift_range = 0.2, # 随机水平平移
    height_shift_range = 0.2,# 随机竖直平移
    rescale = 1/255,         # 数据归一化
    shear_range = 20,       # 随机错切变换
    zoom_range = 0.2,        # 随机放大
    horizontal_flip = True,  # 水平翻转
    fill_mode = 'nearest',   # 填充方式
) 
test_datagen = ImageDataGenerator(
    rescale = 1/255,         # 数据归一化
)

batch_size=32

# 生成训练数据
train_generator = train_datagen.flow_from_directory(
    'image/train',
    target_size=(150,150),
    batch_size=batch_size,
)

# 生成测试数据
test_generator = test_datagen.flow_from_directory(
    'image/test',
    target_size=(150,150),
    batch_size=batch_size,
)

train_generator.class_indices

model.fit_generator(train_generator,
                    steps_per_epoch=len(train_generator),
                    epochs=30,
                    validation_data=test_generator,
                    validation_steps=len(test_generator))

# 画出结构
from keras.utils.vis_utils import plot_model
import matplotlib.pyplot as plt 
# TB：从上往下 LR：左右
plot_model(model,to_file="model.png",show_shapes=True,show_layer_names=True,rankdir='TB')
plt.figure(figsize=(30,30))
img = plt.imread("model.png")
plt.imshow(img)
plt.axis('off')
plt.show()

# pip install h5py
model.save('mask_model_cnn.h5')

# 测试
from keras.models import load_model
import numpy as np

label = np.array(['image_mask','image_nomask'])
# 载入模型
model = load_model('mask_model_cnn.h5')

def process_img(image):
    image = image.resize((150,150))
    image = img_to_array(image)
    image = image/255
    image = np.expand_dims(image,0)
    return image

# 导入图片
image = load_img('image/test/image_mask/0514.jpg')
image

print(label[model.predict_classes(process_img(image))])