#-------------------------------------------------------------#
#   ResNet50的网络部分
#-------------------------------------------------------------#
import keras.backend as K
from keras import backend as K
from keras import initializers, layers, regularizers
from keras.engine import InputSpec, Layer
from keras.initializers import random_normal
from keras.layers import (Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed,
                          ZeroPadding2D)


class BatchNormalization(Layer):
    def __init__(self, epsilon=1e-3, axis=-1,
                 weights=None, beta_init='zero', gamma_init='one',
                 gamma_regularizer=None, beta_regularizer=None, **kwargs):

        self.supports_masking = True
        self.beta_init = initializers.get(beta_init)
        self.gamma_init = initializers.get(gamma_init)
        self.epsilon = epsilon
        self.axis = axis
        self.gamma_regularizer = regularizers.get(gamma_regularizer)
        self.beta_regularizer = regularizers.get(beta_regularizer)
        self.initial_weights = weights
        super(BatchNormalization, self).__init__(**kwargs)

    def build(self, input_shape):
        self.input_spec = [InputSpec(shape=input_shape)]
        shape = (input_shape[self.axis],)

        self.gamma = self.add_weight(shape,
                                     initializer=self.gamma_init,
                                     regularizer=self.gamma_regularizer,
                                     name='{}_gamma'.format(self.name),
                                     trainable=False)
        self.beta = self.add_weight(shape,
                                    initializer=self.beta_init,
                                    regularizer=self.beta_regularizer,
                                    name='{}_beta'.format(self.name),
                                    trainable=False)
        self.running_mean = self.add_weight(shape, initializer='zero',
                                            name='{}_running_mean'.format(self.name),
                                            trainable=False)
        self.running_std = self.add_weight(shape, initializer='one',
                                           name='{}_running_std'.format(self.name),
                                           trainable=False)

        if self.initial_weights is not None:
            self.set_weights(self.initial_weights)
            del self.initial_weights

        self.built = True

    def call(self, x, mask=None):

        assert self.built, 'Layer must be built before being called'
        input_shape = K.int_shape(x)

        reduction_axes = list(range(len(input_shape)))
        del reduction_axes[self.axis]
        broadcast_shape = [1] * len(input_shape)
        broadcast_shape[self.axis] = input_shape[self.axis]

        if sorted(reduction_axes) == range(K.ndim(x))[:-1]:
            x_normed = K.batch_normalization(
                x, self.running_mean, self.running_std,
                self.beta, self.gamma,
                epsilon=self.epsilon)
        else:
            broadcast_running_mean = K.reshape(self.running_mean, broadcast_shape)
            broadcast_running_std = K.reshape(self.running_std, broadcast_shape)
            broadcast_beta = K.reshape(self.beta, broadcast_shape)
            broadcast_gamma = K.reshape(self.gamma, broadcast_shape)
            x_normed = K.batch_normalization(
                x, broadcast_running_mean, broadcast_running_std,
                broadcast_beta, broadcast_gamma,
                epsilon=self.epsilon)

        return x_normed

    def get_config(self):
        config = {'epsilon': self.epsilon,
                  'axis': self.axis,
                  'gamma_regularizer': self.gamma_regularizer.get_config() if self.gamma_regularizer else None,
                  'beta_regularizer': self.beta_regularizer.get_config() if self.beta_regularizer else None}
        base_config = super(BatchNormalization, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


def identity_block(input_tensor, kernel_size, filters, stage, block):

    filters1, filters2, filters3 = filters

    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    x = Conv2D(filters1, (1, 1), kernel_initializer=random_normal(stddev=0.02), name=conv_name_base + '2a')(input_tensor)
    x = BatchNormalization(name=bn_name_base + '2a')(x)
    x = Activation('relu')(x)

    x = Conv2D(filters2, kernel_size, padding='same', kernel_initializer=random_normal(stddev=0.02), name=conv_name_base + '2b')(x)
    x = BatchNormalization(name=bn_name_base + '2b')(x)
    x = Activation('relu')(x)

    x = Conv2D(filters3, (1, 1), kernel_initializer=random_normal(stddev=0.02), name=conv_name_base + '2c')(x)
    x = BatchNormalization(name=bn_name_base + '2c')(x)

    x = layers.add([x, input_tensor])
    x = Activation('relu')(x)
    return x


def conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, 2)):

    filters1, filters2, filters3 = filters

    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    x = Conv2D(filters1, (1, 1), strides=strides, kernel_initializer=random_normal(stddev=0.02),
               name=conv_name_base + '2a')(input_tensor)
    x = BatchNormalization(name=bn_name_base + '2a')(x)
    x = Activation('relu')(x)

    x = Conv2D(filters2, kernel_size, padding='same', kernel_initializer=random_normal(stddev=0.02),
               name=conv_name_base + '2b')(x)
    x = BatchNormalization(name=bn_name_base + '2b')(x)
    x = Activation('relu')(x)

    x = Conv2D(filters3, (1, 1), kernel_initializer=random_normal(stddev=0.02), 
               name=conv_name_base + '2c')(x)
    x = BatchNormalization(name=bn_name_base + '2c')(x)

    shortcut = Conv2D(filters3, (1, 1), strides=strides, kernel_initializer=random_normal(stddev=0.02),
                      name=conv_name_base + '1')(input_tensor)
    shortcut = BatchNormalization(name=bn_name_base + '1')(shortcut)

    x = layers.add([x, shortcut])
    x = Activation('relu')(x)
    return x

def ResNet50(inputs):
    #-----------------------------------#
    #   假设输入进来的图片是600,600,3
    #-----------------------------------#
    img_input = inputs

    # 600,600,3 -> 300,300,64
    x = ZeroPadding2D((3, 3))(img_input)
    x = Conv2D(64, (7, 7), strides=(2, 2), name='conv1')(x)
    x = BatchNormalization(name='bn_conv1')(x)
    x = Activation('relu')(x)

    # 300,300,64 -> 150,150,64
    x = MaxPooling2D((3, 3), strides=(2, 2), padding="same")(x)

    # 150,150,64 -> 150,150,256
    x = conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1))
    x = identity_block(x, 3, [64, 64, 256], stage=2, block='b')
    x = identity_block(x, 3, [64, 64, 256], stage=2, block='c')

    # 150,150,256 -> 75,75,512
    x = conv_block(x, 3, [128, 128, 512], stage=3, block='a')
    x = identity_block(x, 3, [128, 128, 512], stage=3, block='b')
    x = identity_block(x, 3, [128, 128, 512], stage=3, block='c')
    x = identity_block(x, 3, [128, 128, 512], stage=3, block='d')

    # 75,75,512 -> 38,38,1024
    x = conv_block(x, 3, [256, 256, 1024], stage=4, block='a')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='b')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='c')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='d')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='e')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='f')

    # 最终获得一个38,38,1024的共享特征层
    return x

def identity_block_td(input_tensor, kernel_size, filters, stage, block):
    nb_filter1, nb_filter2, nb_filter3 = filters
    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    x = TimeDistributed(Conv2D(nb_filter1, (1, 1), kernel_initializer='normal'), name=conv_name_base + '2a')(input_tensor)
    x = TimeDistributed(BatchNormalization(), name=bn_name_base + '2a')(x)
    x = Activation('relu')(x)

    x = TimeDistributed(Conv2D(nb_filter2, (kernel_size, kernel_size), kernel_initializer='normal',padding='same'), name=conv_name_base + '2b')(x)
    x = TimeDistributed(BatchNormalization(), name=bn_name_base + '2b')(x)
    x = Activation('relu')(x)

    x = TimeDistributed(Conv2D(nb_filter3, (1, 1), kernel_initializer='normal'), name=conv_name_base + '2c')(x)
    x = TimeDistributed(BatchNormalization(), name=bn_name_base + '2c')(x)

    x = Add()([x, input_tensor])
    x = Activation('relu')(x)

    return x

def conv_block_td(input_tensor, kernel_size, filters, stage, block, strides=(2, 2)):
    nb_filter1, nb_filter2, nb_filter3 = filters
    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    x = TimeDistributed(Conv2D(nb_filter1, (1, 1), strides=strides, kernel_initializer='normal'), name=conv_name_base + '2a')(input_tensor)
    x = TimeDistributed(BatchNormalization(), name=bn_name_base + '2a')(x)
    x = Activation('relu')(x)

    x = TimeDistributed(Conv2D(nb_filter2, (kernel_size, kernel_size), padding='same', kernel_initializer='normal'), name=conv_name_base + '2b')(x)
    x = TimeDistributed(BatchNormalization(), name=bn_name_base + '2b')(x)
    x = Activation('relu')(x)

    x = TimeDistributed(Conv2D(nb_filter3, (1, 1), kernel_initializer='normal'), name=conv_name_base + '2c')(x)
    x = TimeDistributed(BatchNormalization(), name=bn_name_base + '2c')(x)

    shortcut = TimeDistributed(Conv2D(nb_filter3, (1, 1), strides=strides, kernel_initializer='normal'), name=conv_name_base + '1')(input_tensor)
    shortcut = TimeDistributed(BatchNormalization(), name=bn_name_base + '1')(shortcut)

    x = Add()([x, shortcut])
    x = Activation('relu')(x)
    return x

def classifier_layers(x):
    # num_rois, 14, 14, 1024 -> num_rois, 7, 7, 2048
    x = conv_block_td(x, 3, [512, 512, 2048], stage=5, block='a', strides=(2, 2))
    # num_rois, 7, 7, 2048 -> num_rois, 7, 7, 2048
    x = identity_block_td(x, 3, [512, 512, 2048], stage=5, block='b')
    # num_rois, 7, 7, 2048 -> num_rois, 7, 7, 2048
    x = identity_block_td(x, 3, [512, 512, 2048], stage=5, block='c')
    # num_rois, 7, 7, 2048 -> num_rois, 1, 1, 2048
    x = TimeDistributed(AveragePooling2D((7, 7)), name='avg_pool')(x)

    return x

[(109, 'keras.layers.add', 'layers.add', (['[x, input_tensor]'], {}), False, 'from keras import initializers, layers, regularizers\n'), (139, 'keras.layers.add', 'layers.add', (['[x, shortcut]'], {}), False, 'from keras import initializers, layers, regularizers\n'), (19, 'keras.initializers.get', 'initializers.get', (['beta_init'], {}), False, 'from keras import initializers, layers, regularizers\n'), (20, 'keras.initializers.get', 'initializers.get', (['gamma_init'], {}), False, 'from keras import initializers, layers, regularizers\n'), (23, 'keras.regularizers.get', 'regularizers.get', (['gamma_regularizer'], {}), False, 'from keras import initializers, layers, regularizers\n'), (24, 'keras.regularizers.get', 'regularizers.get', (['beta_regularizer'], {}), False, 'from keras import initializers, layers, regularizers\n'), (58, 'keras.backend.int_shape', 'K.int_shape', (['x'], {}), True, 'from keras import backend as K\n'), (100, 'keras.layers.Activation', 'Activation', (['"""relu"""'], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (104, 'keras.layers.Activation', 'Activation', (['"""relu"""'], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (110, 'keras.layers.Activation', 'Activation', (['"""relu"""'], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (124, 'keras.layers.Activation', 'Activation', (['"""relu"""'], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (129, 'keras.layers.Activation', 'Activation', (['"""relu"""'], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (140, 'keras.layers.Activation', 'Activation', (['"""relu"""'], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (150, 'keras.layers.ZeroPadding2D', 'ZeroPadding2D', (['(3, 3)'], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (151, 'keras.layers.Conv2D', 'Conv2D', (['(64)', '(7, 7)'], {'strides': '(2, 2)', 'name': '"""conv1"""'}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (153, 'keras.layers.Activation', 'Activation', (['"""relu"""'], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (156, 'keras.layers.MaxPooling2D', 'MaxPooling2D', (['(3, 3)'], {'strides': '(2, 2)', 'padding': '"""same"""'}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (187, 'keras.layers.Activation', 'Activation', (['"""relu"""'], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (191, 'keras.layers.Activation', 'Activation', (['"""relu"""'], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (196, 'keras.layers.Add', 'Add', ([], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (197, 'keras.layers.Activation', 'Activation', (['"""relu"""'], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (208, 'keras.layers.Activation', 'Activation', (['"""relu"""'], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (212, 'keras.layers.Activation', 'Activation', (['"""relu"""'], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (220, 'keras.layers.Add', 'Add', ([], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (221, 'keras.layers.Activation', 'Activation', (['"""relu"""'], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (29, 'keras.engine.InputSpec', 'InputSpec', ([], {'shape': 'input_shape'}), False, 'from keras.engine import InputSpec, Layer\n'), (66, 'keras.backend.batch_normalization', 'K.batch_normalization', (['x', 'self.running_mean', 'self.running_std', 'self.beta', 'self.gamma'], {'epsilon': 'self.epsilon'}), True, 'from keras import backend as K\n'), (71, 'keras.backend.reshape', 'K.reshape', (['self.running_mean', 'broadcast_shape'], {}), True, 'from keras import backend as K\n'), (72, 'keras.backend.reshape', 'K.reshape', (['self.running_std', 'broadcast_shape'], {}), True, 'from keras import backend as K\n'), (73, 'keras.backend.reshape', 'K.reshape', (['self.beta', 'broadcast_shape'], {}), True, 'from keras import backend as K\n'), (74, 'keras.backend.reshape', 'K.reshape', (['self.gamma', 'broadcast_shape'], {}), True, 'from keras import backend as K\n'), (75, 'keras.backend.batch_normalization', 'K.batch_normalization', (['x', 'broadcast_running_mean', 'broadcast_running_std', 'broadcast_beta', 'broadcast_gamma'], {'epsilon': 'self.epsilon'}), True, 'from keras import backend as K\n'), (185, 'keras.layers.Conv2D', 'Conv2D', (['nb_filter1', '(1, 1)'], {'kernel_initializer': '"""normal"""'}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (189, 'keras.layers.Conv2D', 'Conv2D', (['nb_filter2', '(kernel_size, kernel_size)'], {'kernel_initializer': '"""normal"""', 'padding': '"""same"""'}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (193, 'keras.layers.Conv2D', 'Conv2D', (['nb_filter3', '(1, 1)'], {'kernel_initializer': '"""normal"""'}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (206, 'keras.layers.Conv2D', 'Conv2D', (['nb_filter1', '(1, 1)'], {'strides': 'strides', 'kernel_initializer': '"""normal"""'}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (210, 'keras.layers.Conv2D', 'Conv2D', (['nb_filter2', '(kernel_size, kernel_size)'], {'padding': '"""same"""', 'kernel_initializer': '"""normal"""'}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (214, 'keras.layers.Conv2D', 'Conv2D', (['nb_filter3', '(1, 1)'], {'kernel_initializer': '"""normal"""'}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (217, 'keras.layers.Conv2D', 'Conv2D', (['nb_filter3', '(1, 1)'], {'strides': 'strides', 'kernel_initializer': '"""normal"""'}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (232, 'keras.layers.AveragePooling2D', 'AveragePooling2D', (['(7, 7)'], {}), False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), (98, 'keras.initializers.random_normal', 'random_normal', ([], {'stddev': '(0.02)'}), False, 'from keras.initializers import random_normal\n'), (102, 'keras.initializers.random_normal', 'random_normal', ([], {'stddev': '(0.02)'}), False, 'from keras.initializers import random_normal\n'), (106, 'keras.initializers.random_normal', 'random_normal', ([], {'stddev': '(0.02)'}), False, 'from keras.initializers import random_normal\n'), (121, 'keras.initializers.random_normal', 'random_normal', ([], {'stddev': '(0.02)'}), False, 'from keras.initializers import random_normal\n'), (126, 'keras.initializers.random_normal', 'random_normal', ([], {'stddev': '(0.02)'}), False, 'from keras.initializers import random_normal\n'), (131, 'keras.initializers.random_normal', 'random_normal', ([], {'stddev': '(0.02)'}), False, 'from keras.initializers import random_normal\n'), (135, 'keras.initializers.random_normal', 'random_normal', ([], {'stddev': '(0.02)'}), False, 'from keras.initializers import random_normal\n'), (65, 'keras.backend.ndim', 'K.ndim', (['x'], {}), True, 'from keras import backend as K\n')]