import  tensorflow as tf
from    tensorflow import keras

base_model = keras.applications.DenseNet121(weights='imagenet')
base_model.summary()

import  tensorflow as tf
from    tensorflow import keras
from    tensorflow.keras import layers, models, Sequential, backend
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Dense, Flatten, Dropout, BatchNormalization, Activation, GlobalAveragePooling2D
from tensorflow.keras.layers import Concatenate, Lambda, Input, ZeroPadding2D, AveragePooling2D

def conv_block(x, nb_filter, dropout_rate=None, name=None):
   
    inter_channel = nb_filter*4

    # 1x1 convolution
    x = BatchNormalization(epsilon=1.1e-5, axis=3, name=name+'_bn1')(x)
    x = Activation('relu', name=name+'_relu1')(x)
    x = Conv2D(inter_channel, 1, 1, name=name+'_conv1', use_bias=False)(x)

    if dropout_rate:
        x = Dropout(dropout_rate)(x)

    # 3x3 convolution
    x = BatchNormalization(epsilon=1.1e-5, axis=3, name=name+'_bn2')(x)
    x = Activation('relu', name=name+'_relu2')(x)
    x = ZeroPadding2D((1, 1), name=name+'_zeropadding2')(x)
    x = Conv2D(nb_filter, 3, 1, name=name+'_conv2', use_bias=False)(x)

    if dropout_rate:
        x = Dropout(dropout_rate)(x)

    return x

def dense_block(x, stage, nb_layers, nb_filter, growth_rate, dropout_rate=None,
                grow_nb_filters=True, name =None):

    concat_feat = x # store the last layer output

    for i in range(nb_layers):
       
        branch = i+1
        x =conv_block(concat_feat, growth_rate, dropout_rate, name=name+str(stage)+'_block'+str(branch)) # 在参考的基础，修改的地方这里应该是相同的growth_rate=32
        concat_feat = Concatenate(axis=3, name=name+str(stage)+'_block'+str(branch))([concat_feat, x])

        if grow_nb_filters:
            nb_filter += growth_rate

    return concat_feat, nb_filter

ef transition_block (x,stage, nb_filter, compression=1.0, dropout_rate=None, name=None):

    x = BatchNormalization(epsilon=1.1e-5, axis=3, name=name+str(stage)+'_bn')(x)
    x = Activation('relu', name=name+str(stage)+'_relu')(x)
   
    x = Conv2D(int(nb_filter*compression), 1, 1, name=name+str(stage)+'_conv', use_bias=False)(x)

    if dropout_rate:
        x = Dropout(dropout_rate)(x)
   
    x = AveragePooling2D((2,2), strides=(2,2), name=name+str(stage)+'_pooling2d')(x)

    return x

def DenseNet(nb_dense_block=4, growth_rate=32, nb_filter=64, reduction=0.0, dropout_rate=0.0, weight_decay=1e-4, classes=1000, weights_path=None):

    compression = 1.0-reduction
    nb_filter=64
    nb_layers = [6, 12, 24, 16]  #For DenseNet-121
   
    img_input = Input(shape=(224,224,3))

    #initial convolution
    x = ZeroPadding2D((3, 3), name='conv1_zeropadding')(img_input)
    x = Conv2D(nb_filter, 7, 2, name='conv1', use_bias=False)(x)
    x = BatchNormalization(epsilon=1.1e-5, axis=3, name='conv1_bn')(x)
    x = Activation('relu', name='relu1')(x)
    x = ZeroPadding2D((1, 1), name='pool1_zeropadding')(x)
    x = MaxPooling2D((3, 3), strides=(2, 2), name='pool1')(x)

    #Dense block and Transition layer
    for block_id in range(nb_dense_block-1):
        stage = block_id+2 # start from 2
        x, nb_filter = dense_block(x, stage, nb_layers[block_id], nb_filter, growth_rate,
        dropout_rate=dropout_rate, name='Dense')

        x = transition_block(x, stage, nb_filter, compression=compression, dropout_rate=dropout_rate, name='Trans')
        nb_filter *=compression
   
    final_stage = stage + 1
    x, nb_filter=dense_block(x, final_stage, nb_layers[-1], nb_filter, growth_rate,
        dropout_rate=dropout_rate, name='Dense')

    # top layer
    x = BatchNormalization(name= 'final_conv_bn')(x)
    x = Activation('relu', name='final_act')(x)
    x = GlobalAveragePooling2D(name='final_pooling')(x)
    x = Dense(classes, activation='softmax', name='fc')(x)

    model=models.Model(img_input, x, name='DenseNet121')

    return model
   
def main():

    model = DenseNet(reduction=.5)
    model.summary()

if __name__=='__main__':
    main()