model1

from tensorflow import keras
from tensorflow.keras.models import Sequential
from tensorflow.keras import backend as K
from tensorflow.keras.layers import AveragePooling2D, Conv2D, Input, Dense, Dropout,MaxPooling2D , \
    Flatten, PReLU, BatchNormalization, Layer, Lambda, LeakyReLU, SpatialDropout2D, SeparableConv2D,GlobalAveragePooling2D
from tensorflow.keras.models import Model
import tensorflow as tf
input_shape = (img_width,img_height,3)
feature_dim = 100
#num_classes = 10

class CenterLossLayer(Layer):

    def __init__(self, class_num, feature_dim, alpha=0.5, **kwargs):
        super().__init__(**kwargs)
        self.class_num = class_num
        self.feature_dim = feature_dim
        self.alpha = alpha

    def build(self, input_shape):
        self.centers = self.add_weight(name='centers',
                                       shape=(self.class_num, self.feature_dim),
                                       initializer='uniform',
                                       trainable=False)
        super().build(input_shape)

    def call(self, x, mask=None):
        delta_centers = K.dot(K.transpose(x[1]), (K.dot(x[1], self.centers) - x[0]))  # 10x2
        center_counts = K.sum(K.transpose(x[1]), axis=1, keepdims=True) + 1  # 10x1
        delta_centers /= center_counts
        new_centers = self.centers - self.alpha * delta_centers
        self.centers.assign(new_centers)  # Chieko: something's wrong with add_update()
        self.result = x[0] - K.dot(x[1], self.centers)  # Chieko: recalculate the distance from center to each point
        self.result = K.sum(self.result ** 2, axis=1, keepdims=True)  # / K.dot(x[1], center_counts)
        # Chieko: N(x**2 + y**2)
        return self.result  # Nx1

    def compute_output_shape(self, input_shape):
        return K.int_shape(self.result)

    def get_config(self):
        config = super().get_config().copy()
        config.update({
            'centers': self.centers,
        })
        return config


### custom loss

def zero_loss(y_true, y_pred):
    return 0.5 * K.sum(y_pred, axis=0)
class ResnetIdentityBlock(tf.keras.Model):

  def __init__(self, filters):
    super(ResnetIdentityBlock, self).__init__(name='')
    self.filters = filters





  def call(self, input_tensor, training=False):

    x_box = Conv2D(self.filters, (3,3) , padding='same', kernel_initializer='he_normal' )(input_tensor)
    x_box = BatchNormalization()(x_box)
    x_box = LeakyReLU(0.1)(x_box)


    x_box = Conv2D(64, (3,3) , padding='same', kernel_initializer='he_normal' )(x_box)
    x_box = SpatialDropout2D(0.3)(x_box)
    x = LeakyReLU(0.1)(x_box)


    x += input_tensor
    return x

def resnet18(input_shape: tuple, num_classes: int, feature_dim) -> (Model, Model):
    x_box = inputs = Input(input_shape, name="inpnode")
    labels = Input((num_classes,), name='labels')

    x_box = Lambda(lambda x: x / 255)(x_box)





    x_box = Conv2D(32, (3,3) , padding='same', kernel_initializer='he_normal' )(x_box)
    x_box = BatchNormalization()(x_box)
    x_box = LeakyReLU(0.1)(x_box)


    for i in range(1):

      x = tf.identity(x_box)

      x_box = Conv2D(128, (3,3) , padding='same', kernel_initializer='he_normal' )(x_box)
      x_box = SpatialDropout2D(0.1)(x_box)
      x_box = LeakyReLU(0.1)(x_box)

      x_box = Conv2D(32, (3,3) , padding='same', kernel_initializer='he_normal' )(x_box)
      x_box = BatchNormalization()(x_box)
      x_box = LeakyReLU(0.1)(x_box) + x

      x_box = AveragePooling2D(2)(x_box)

    x_box = Conv2D(64, (3,3) , padding='same', kernel_initializer='he_normal' )(x_box)
    x_box = BatchNormalization()(x_box)
    x_box = LeakyReLU(0.1)(x_box)

    for i in range(1):
      x = tf.identity(x_box)

      x_box = Conv2D(128, (3,3) , padding='same', kernel_initializer='he_normal' )(x_box)
      x_box = SpatialDropout2D(0.1)(x_box)
      x_box = LeakyReLU(0.1)(x_box)

      x_box = Conv2D(64, (3,3) , padding='same', kernel_initializer='he_normal' )(x_box)
      x_box = BatchNormalization()(x_box)
      x_box = LeakyReLU(0.1)(x_box) + x

    x_box = AveragePooling2D(2)(x_box)

    x_box = Conv2D(128, (3,3), padding='same' , kernel_initializer='he_normal' )(x_box)
    x_box = BatchNormalization()(x_box)
    x_box = LeakyReLU(0.1)(x_box)

    for i in range(1):
      x = tf.identity(x_box)

      x_box = Conv2D(128, (3,3) , padding='same', kernel_initializer='he_normal' )(x_box)
      x_box = SpatialDropout2D(0.1)(x_box)
      x_box = LeakyReLU(0.1)(x_box)

      x_box = Conv2D(128, (3,3) , padding='same', kernel_initializer='he_normal' )(x_box)
      x_box = BatchNormalization()(x_box)
      x_box = LeakyReLU(0.1)(x_box) + x

    x_box = AveragePooling2D(2)(x_box)

    x_box = Conv2D(256, (3,3) , kernel_initializer='he_normal' )(x_box)
    x_box = BatchNormalization()(x_box)
    x_box = LeakyReLU(0.1)(x_box)

    x_box = AveragePooling2D(2)(x_box)


    #x_box = AveragePooling2D(2)(x_box)

    y_box = Flatten()(x_box)


    y_box = Dropout(0.2)(y_box)
    y_box =  Dense( 200, kernel_initializer='he_normal')(y_box)
    y_box =  Dropout(0.1) (y_box)
    y_box =  tf.keras.layers.PReLU() (y_box)
    y_box =  Dense( 200, kernel_initializer='he_normal')(y_box)
    y_box =  Dropout(0.1) (y_box)
    y_box =  tf.keras.layers.PReLU() (y_box)
    inter_dim =  Dense(feature_dim,  kernel_initializer='he_normal') (y_box)




    loss = CenterLossLayer(num_classes, feature_dim, name='center_loss')([inter_dim, labels])

    logits = Dense(num_classes, kernel_initializer='lecun_normal',
                   activation='softmax', name='logits')(inter_dim)

    model = Model(inputs=[inputs, labels], outputs=[logits, loss])
    infer_model = Model(inputs=inputs, outputs=inter_dim)
    return model, infer_model

train_model, infer_model = resnet18(input_shape, num_classes, feature_dim)