# Bottleneck layers
# BN+ReLU+1*1Conv+BN+ReLU+3*3Conv
# 每一层产生k个特征图，原论文中让每个1*1个卷积产生4k个特征图
# growth_rate:网络的增长率，原论文中超参数名称为k
def conv_block(x,growth_rate,name):
    x1=BatchNormalization(name=name+'_0_bn')(x)
    x1=ReLU(name=name+'_0_relu')(x1)
    x1=Conv2D(filters=4*growth_rate,
              kernel_size=(1,1),
              use_bias=False,
              name=name+'_1_conv')(x1)
    x1=BatchNormalization(name=name+'_1_bn')(x1)
    x1=ReLU(name=name+'_1_relu')(x1)
    # 论文原话：对内核大小为3*3的卷积层，输入的每一边都用一个像素补零，以保持特征图大小固定。
    x1=Conv2D(growth_rate,
              kernel_size=(3,3),
              padding='same',
              use_bias=False,
              name=name+'_2_conv')(x1)
    # 将前面所有层的特征堆叠后传到下一层
    out=Concatenate(axis=-1,name=name+'_concat')([x,x1])
    return out

'''
x: input tensor.
    blocks: integer, the number of building blocks.
    name: string, block label.
'''
# DenseNet
def dense_block(x,blocks,name):
    for i in range(blocks):
        x=conv_block(x,32,name=name+'_block'+str(i+1))
    return x

4.2 Transition Layer(过渡层)

  过渡层： $BN+1*1Conv+2*2AvgPool$

# 论文原话：使用1*1卷积和2*2平均池化作为两个连续密集块之间的过渡层
# 过渡层：BN+1*1Conv+2*2AvgPool
# reduction：压缩因子，原论文中设置为0.5
def transition_block(x,reduction,name):
    x=BatchNormalization(name=name+'_bn')(x)
    x=ReLU(name=name+'_relu')(x)
    x=Conv2D(filters=int(x.shape[-1]*reduction),
             kernel_size=(1,1),
             use_bias=False,
             name=name+'_conv')(x)
    x=AveragePooling2D(pool_size=(2,2),strides=2,name=name+'_pool')(x)
    return x

def DenseNet(blocks,include_top=True,weights='imagenet',intput_shape=None,classes=1000):
    img_input=Input(shape=intput_shape)
    x=ZeroPadding2D(padding=((3,3),(3,3)))(img_input)
    # 论文原话：初始卷积层包含2k个大小为7*7的卷积，步长为2(k=32)
    # 其实这里完全可以去掉这两个ZeroPadding2D，直接给这里的卷积和池化设置padding='same'即可，
    # 但是我看谷歌的源码中没用same padding,所以这里我也没用。
    x=Conv2D(filters=64,kernel_size=(7,7),strides=2,use_bias=False,name='conv1/conv')(x)
    x=BatchNormalization(name='conv1/bn')(x)
    x=ReLU(name='conv1/relu')(x)
    x=ZeroPadding2D(padding=((1,1),(1,1)))(x)
    x=MaxPooling2D(pool_size=(3,3),strides=2,name='pool1')(x)

    x=dense_block(x,blocks[0],name='conv2')
    x=transition_block(x,0.5,name='pool2')

    x=dense_block(x,blocks[1],name='conv3')
    x=transition_block(x,0.5,name='pool3')

    x=dense_block(x,blocks[2],name='conv4')
    x=transition_block(x,0.5,name='pool4')

    x=dense_block(x,blocks[3],name='conv5')

    x=BatchNormalization(name='bn')(x)
    x=ReLU(name='relu')(x)
    if include_top:
        x=GlobalAveragePooling2D(name='global_avg_pool')(x)
        x=Dense(classes,activation='softmax')(x)

    model=Model(img_input,x)
    return model

densenet121=DenseNet(blocks=[6,12,24,16],include_top=False,intput_shape=(224,224,3))
densenet121.summary()

# 类别数
num_classes = 17
# 批次大小
batch_size = 32
# 周期数
epochs=100
# 图片大小
image_size = 224

model=tf.keras.Sequential([
    densenet121,
    layers.GlobalAveragePooling2D(),
    layers.Dense(num_classes,activation='softmax')
])
model.summary()

# 训练集数据进行数据增强
train_datagen = ImageDataGenerator(
    rotation_range=20,  # 随机旋转度数
    width_shift_range=0.1,  # 随机水平平移
    height_shift_range=0.1,  # 随机竖直平移
    rescale=1 / 255,  # 数据归一化
    shear_range=10,  # 随机错切变换
    zoom_range=0.1,  # 随机放大
    horizontal_flip=True,  # 水平翻转
    brightness_range=(0.7, 1.3),  # 亮度变化
    fill_mode='nearest',  # 填充方式
)
# 测试集数据只需要归一化就可以
test_datagen = ImageDataGenerator(
    rescale=1 / 255,  # 数据归一化
)

# 训练集数据生成器，可以在训练时自动产生数据进行训练
# 从'data/train'获得训练集数据
# 获得数据后会把图片resize为image_size×image_size的大小
# generator每次会产生batch_size个数据
train_generator = train_datagen.flow_from_directory(
    '../data/花朵分类(17)/train',
    target_size=(image_size, image_size), # 调整图像尺寸
    batch_size=batch_size,
)

# 测试集数据生成器
test_generator = test_datagen.flow_from_directory(
    '../data/花朵分类(17)/test',
    target_size=(image_size, image_size),
    batch_size=batch_size,
)
# 字典的键为17个文件夹的名字，值为对应的分类编号
print(train_generator.class_indices)

# 学习率调节函数，逐渐减小学习率
def adjust_learning_rate(epoch):
    # 前40周期
    if epoch<=40:
        lr = 1e-4
    # 前40到80周期
    elif 40 < epoch <= 80:
        lr = 1e-5
    # 80到100周期
    else:
        lr = 1e-6
    return lr

# 定义优化器
adam = Adam(lr=1e-4)

# 读取模型
checkpoint_save_path = "./checkpoint/DenseNet121(不使用预训练)-花朵分类(17)-.ckpt"
if os.path.exists(checkpoint_save_path + '.index'):
    print('-------------load the model-----------------')
    model.load_weights(checkpoint_save_path)
# 保存模型
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_save_path,
                                                 save_weights_only=True,
                                                 save_best_only=True)

# 定义学习率衰减策略
callbacks = []
callbacks.append(LearningRateScheduler(adjust_learning_rate))
callbacks.append(cp_callback)

# 定义优化器，loss function，训练过程中计算准确率
model.compile(optimizer=adam,loss='categorical_crossentropy',metrics=['accuracy'])

# Tensorflow2.1版本(包括2.1)之后可以直接使用fit训练模型
history = model.fit(x=train_generator,epochs=epochs,validation_data=test_generator,callbacks=callbacks)

# 画出训练集准确率曲线图
plt.plot(np.arange(epochs),history.history['accuracy'],c='b',label='train_accuracy')
# 画出验证集准确率曲线图
plt.plot(np.arange(epochs),history.history['val_accuracy'],c='y',label='val_accuracy')
# 图例
plt.legend()
# x坐标描述
plt.xlabel('epochs')
# y坐标描述
plt.ylabel('accuracy')
# 显示图像
plt.show()

# 画出训练集loss曲线图
plt.plot(np.arange(epochs),history.history['loss'],c='b',label='train_loss')
# 画出验证集loss曲线图
plt.plot(np.arange(epochs),history.history['val_loss'],c='y',label='val_loss')
# 图例
plt.legend()
# x坐标描述
plt.xlabel('epochs')
# y坐标描述
plt.ylabel('loss')
# 显示图像
plt.show()