#-------------------------------------------------------------#
#   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, 8, 109, 37), 'keras.layers.add', 'layers.add', ({(109, 19, 109, 36): '[x, input_tensor]'}, {}), '([x, input_tensor])', False, 'from keras import initializers, layers, regularizers\n'), ((139, 8, 139, 33), 'keras.layers.add', 'layers.add', ({(139, 19, 139, 32): '[x, shortcut]'}, {}), '([x, shortcut])', False, 'from keras import initializers, layers, regularizers\n'), ((19, 25, 19, 52), 'keras.initializers.get', 'initializers.get', ({(19, 42, 19, 51): 'beta_init'}, {}), '(beta_init)', False, 'from keras import initializers, layers, regularizers\n'), ((20, 26, 20, 54), 'keras.initializers.get', 'initializers.get', ({(20, 43, 20, 53): 'gamma_init'}, {}), '(gamma_init)', False, 'from keras import initializers, layers, regularizers\n'), ((23, 33, 23, 68), 'keras.regularizers.get', 'regularizers.get', ({(23, 50, 23, 67): 'gamma_regularizer'}, {}), '(gamma_regularizer)', False, 'from keras import initializers, layers, regularizers\n'), ((24, 32, 24, 66), 'keras.regularizers.get', 'regularizers.get', ({(24, 49, 24, 65): 'beta_regularizer'}, {}), '(beta_regularizer)', False, 'from keras import initializers, layers, regularizers\n'), ((58, 22, 58, 36), 'keras.backend.int_shape', 'K.int_shape', ({(58, 34, 58, 35): 'x'}, {}), '(x)', True, 'from keras import backend as K\n'), ((100, 8, 100, 26), 'keras.layers.Activation', 'Activation', ({(100, 19, 100, 25): '"""relu"""'}, {}), "('relu')", False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((104, 8, 104, 26), 'keras.layers.Activation', 'Activation', ({(104, 19, 104, 25): '"""relu"""'}, {}), "('relu')", False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((110, 8, 110, 26), 'keras.layers.Activation', 'Activation', ({(110, 19, 110, 25): '"""relu"""'}, {}), "('relu')", False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((124, 8, 124, 26), 'keras.layers.Activation', 'Activation', ({(124, 19, 124, 25): '"""relu"""'}, {}), "('relu')", False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((129, 8, 129, 26), 'keras.layers.Activation', 'Activation', ({(129, 19, 129, 25): '"""relu"""'}, {}), "('relu')", False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((140, 8, 140, 26), 'keras.layers.Activation', 'Activation', ({(140, 19, 140, 25): '"""relu"""'}, {}), "('relu')", False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((150, 8, 150, 29), 'keras.layers.ZeroPadding2D', 'ZeroPadding2D', ({(150, 22, 150, 28): '(3, 3)'}, {}), '((3, 3))', False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((151, 8, 151, 56), 'keras.layers.Conv2D', 'Conv2D', (), '', False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((153, 8, 153, 26), 'keras.layers.Activation', 'Activation', ({(153, 19, 153, 25): '"""relu"""'}, {}), "('relu')", False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((156, 8, 156, 60), 'keras.layers.MaxPooling2D', 'MaxPooling2D', (), '', False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((187, 8, 187, 26), 'keras.layers.Activation', 'Activation', ({(187, 19, 187, 25): '"""relu"""'}, {}), "('relu')", False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((191, 8, 191, 26), 'keras.layers.Activation', 'Activation', ({(191, 19, 191, 25): '"""relu"""'}, {}), "('relu')", False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((196, 8, 196, 13), 'keras.layers.Add', 'Add', ({}, {}), '()', False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((197, 8, 197, 26), 'keras.layers.Activation', 'Activation', ({(197, 19, 197, 25): '"""relu"""'}, {}), "('relu')", False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((208, 8, 208, 26), 'keras.layers.Activation', 'Activation', ({(208, 19, 208, 25): '"""relu"""'}, {}), "('relu')", False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((212, 8, 212, 26), 'keras.layers.Activation', 'Activation', ({(212, 19, 212, 25): '"""relu"""'}, {}), "('relu')", False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((220, 8, 220, 13), 'keras.layers.Add', 'Add', ({}, {}), '()', False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((221, 8, 221, 26), 'keras.layers.Activation', 'Activation', ({(221, 19, 221, 25): '"""relu"""'}, {}), "('relu')", False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((29, 27, 29, 55), 'keras.engine.InputSpec', 'InputSpec', (), '', False, 'from keras.engine import InputSpec, Layer\n'), ((66, 23, 69, 37), 'keras.backend.batch_normalization', 'K.batch_normalization', (), '', True, 'from keras import backend as K\n'), ((71, 37, 71, 82), 'keras.backend.reshape', 'K.reshape', ({(71, 47, 71, 64): 'self.running_mean', (71, 66, 71, 81): 'broadcast_shape'}, {}), '(self.running_mean, broadcast_shape)', True, 'from keras import backend as K\n'), ((72, 36, 72, 80), 'keras.backend.reshape', 'K.reshape', ({(72, 46, 72, 62): 'self.running_std', (72, 64, 72, 79): 'broadcast_shape'}, {}), '(self.running_std, broadcast_shape)', True, 'from keras import backend as K\n'), ((73, 29, 73, 66), 'keras.backend.reshape', 'K.reshape', ({(73, 39, 73, 48): 'self.beta', (73, 50, 73, 65): 'broadcast_shape'}, {}), '(self.beta, broadcast_shape)', True, 'from keras import backend as K\n'), ((74, 30, 74, 68), 'keras.backend.reshape', 'K.reshape', ({(74, 40, 74, 50): 'self.gamma', (74, 52, 74, 67): 'broadcast_shape'}, {}), '(self.gamma, broadcast_shape)', True, 'from keras import backend as K\n'), ((75, 23, 78, 37), 'keras.backend.batch_normalization', 'K.batch_normalization', (), '', True, 'from keras import backend as K\n'), ((185, 24, 185, 79), 'keras.layers.Conv2D', 'Conv2D', (), '', False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((189, 24, 189, 114), 'keras.layers.Conv2D', 'Conv2D', (), '', False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((193, 24, 193, 79), 'keras.layers.Conv2D', 'Conv2D', (), '', False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((206, 24, 206, 96), 'keras.layers.Conv2D', 'Conv2D', (), '', False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((210, 24, 210, 115), 'keras.layers.Conv2D', 'Conv2D', (), '', False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((214, 24, 214, 79), 'keras.layers.Conv2D', 'Conv2D', (), '', False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((217, 31, 217, 103), 'keras.layers.Conv2D', 'Conv2D', (), '', False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((232, 24, 232, 48), 'keras.layers.AveragePooling2D', 'AveragePooling2D', ({(232, 41, 232, 47): '(7, 7)'}, {}), '((7, 7))', False, 'from keras.layers import Activation, Add, AveragePooling2D, Conv2D, MaxPooling2D, TimeDistributed, ZeroPadding2D\n'), ((98, 52, 98, 78), 'keras.initializers.random_normal', 'random_normal', (), '', False, 'from keras.initializers import random_normal\n'), ((102, 73, 102, 99), 'keras.initializers.random_normal', 'random_normal', (), '', False, 'from keras.initializers import random_normal\n'), ((106, 52, 106, 78), 'keras.initializers.random_normal', 'random_normal', (), '', False, 'from keras.initializers import random_normal\n'), ((121, 69, 121, 95), 'keras.initializers.random_normal', 'random_normal', (), '', False, 'from keras.initializers import random_normal\n'), ((126, 73, 126, 99), 'keras.initializers.random_normal', 'random_normal', (), '', False, 'from keras.initializers import random_normal\n'), ((131, 52, 131, 78), 'keras.initializers.random_normal', 'random_normal', (), '', False, 'from keras.initializers import random_normal\n'), ((135, 76, 135, 102), 'keras.initializers.random_normal', 'random_normal', (), '', False, 'from keras.initializers import random_normal\n'), ((65, 43, 65, 52), 'keras.backend.ndim', 'K.ndim', ({(65, 50, 65, 51): 'x'}, {}), '(x)', True, 'from keras import backend as K\n')]