text: from tensorflow import keras
  import tensorflow as tf
  
  
  class FastSRGAN(object):
      """SRGAN for fast super resolution."""
  
      def __init__(self, args):
          """
          Initializes the Mobile SRGAN class.
          Args:
              args: CLI arguments that dictate how to build the model.
          Returns:
              None
          """
          self.hr_height = args.hr_size
          self.hr_width = args.hr_size
          # self.lr_height = self.hr_height // 4  # Low resolution height
          # self.lr_width = self.hr_width // 4  # Low resolution width
          self.lr_height = args.lr_size  # Low resolution height
          self.lr_width = args.lr_size  # Low resolution width
          self.lr_shape = (self.lr_height, self.lr_width, 3)
          self.hr_shape = (self.hr_height, self.hr_width, 3)
          self.iterations = 0
  
          # Number of inverted residual blocks in the mobilenet generator
          self.n_residual_blocks = 6
  
          # Define a learning rate decay schedule.
          self.gen_schedule = keras.optimizers.schedules.ExponentialDecay(
              args.lr,
              decay_steps=100000,
              decay_rate=0.1,
              staircase=True
          )
  
          self.disc_schedule = keras.optimizers.schedules.ExponentialDecay(
              args.lr * 5,  # TTUR - Two Time Scale Updates
              decay_steps=100000,
              decay_rate=0.1,
              staircase=True
          )
  
          self.gen_optimizer = keras.optimizers.Adam(learning_rate=self.gen_schedule)
          self.disc_optimizer = keras.optimizers.Adam(learning_rate=self.disc_schedule)
  
          # We use a pre-trained VGG19 model to extract image features from the high resolution
          # and the generated high resolution images and minimize the mse between them
          self.vgg = self.build_vgg()
          self.vgg.trainable = False
  
          # Calculate output shape of D (PatchGAN)
          patch = int(self.hr_height / 2 ** 4)
          self.disc_patch = (patch, patch, 1)
  
          # Number of filters in the first layer of G and D
          self.gf = 32  # Realtime Image Enhancement GAN Galteri et al.
          self.df = 32
  
          # Build and compile the discriminator
          self.discriminator = self.build_discriminator()
  
          # Build and compile the generator for pretraining.
          self.generator = self.build_generator()
  
      @tf.function
      def content_loss(self, hr, sr):
          sr = keras.applications.vgg19.preprocess_input(((sr + 1.0) * 255) / 2.0)
          hr = keras.applications.vgg19.preprocess_input(((hr + 1.0) * 255) / 2.0)
          sr_features = self.vgg(sr) / 12.75
          hr_features = self.vgg(hr) / 12.75
          return tf.keras.losses.MeanSquaredError()(hr_features, sr_features)
  
      def build_vgg(self):
          """
          Builds a pre-trained VGG19 model that outputs image features extracted at the
          third block of the model
          """
          # Get the vgg network. Extract features from Block 5, last convolution.
          # vgg = keras.applications.VGG19(weights="imagenet", input_shape=self.hr_shape, include_top=False)
          vgg = keras.applications.VGG19(weights="./pretrained/vgg19_weights_tf_dim_ordering_tf_kernels_notop.h5", input_shape=self.hr_shape, include_top=False)
          vgg.trainable = False
          for layer in vgg.layers:
              layer.trainable = False
  
          # Create model and compile
          model = keras.models.Model(inputs=vgg.input, outputs=vgg.get_layer("block5_conv4").output)
  
          return model
  
      def build_generator(self):
          """Build the generator that will do the Super Resolution task.
          Based on the Mobilenet design. Idea from Galteri et al."""
  
          def _make_divisible(v, divisor, min_value=None):
              if min_value is None:
                  min_value = divisor
              new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
              # Make sure that round down does not go down by more than 10%.
              if new_v < 0.9 * v:
                  new_v += divisor
              return new_v
  
          def residual_block(inputs, filters, block_id, expansion=6, stride=1, alpha=1.0):
              """Inverted Residual block that uses depth wise convolutions for parameter efficiency.
              Args:
                  inputs: The input feature map.
                  filters: Number of filters in each convolution in the block.
                  block_id: An integer specifier for the id of the block in the graph.
                  expansion: Channel expansion factor.
                  stride: The stride of the convolution.
                  alpha: Depth expansion factor.
              Returns:
                  x: The output of the inverted residual block.
              """
              channel_axis = 1 if keras.backend.image_data_format() == 'channels_first' else -1
  
              in_channels = keras.backend.int_shape(inputs)[channel_axis]
              pointwise_conv_filters = int(filters * alpha)
              pointwise_filters = _make_divisible(pointwise_conv_filters, 8)
              x = inputs
              prefix = 'block_{}_'.format(block_id)
  
              if block_id:
                  # Expand
                  x = keras.layers.Conv2D(expansion * in_channels,
                                          kernel_size=1,
                                          padding='same',
                                          use_bias=True,
                                          activation=None,
                                          name=prefix + 'expand')(x)
                  x = keras.layers.BatchNormalization(axis=channel_axis,
                                                      epsilon=1e-3,
                                                      momentum=0.999,
                                                      name=prefix + 'expand_BN')(x)
                  x = keras.layers.Activation('relu', name=prefix + 'expand_relu')(x)
              else:
                  prefix = 'expanded_conv_'
  
              # Depthwise
              x = keras.layers.DepthwiseConv2D(kernel_size=3,
                                               strides=stride,
                                               activation=None,
                                               use_bias=True,
                                               padding='same' if stride == 1 else 'valid',
                                               name=prefix + 'depthwise')(x)
              x = keras.layers.BatchNormalization(axis=channel_axis,
                                                  epsilon=1e-3,
                                                  momentum=0.999,
                                                  name=prefix + 'depthwise_BN')(x)
  
              x = keras.layers.Activation('relu', name=prefix + 'depthwise_relu')(x)
  
              # Project
              x = keras.layers.Conv2D(pointwise_filters,
                                      kernel_size=1,
                                      padding='same',
                                      use_bias=True,
                                      activation=None,
                                      name=prefix + 'project')(x)
              x = keras.layers.BatchNormalization(axis=channel_axis,
                                                  epsilon=1e-3,
                                                  momentum=0.999,
                                                  name=prefix + 'project_BN')(x)
  
              if in_channels == pointwise_filters and stride == 1:
                  return keras.layers.Add(name=prefix + 'add')([inputs, x])
              return x
  
          def deconv2d(layer_input):
              """Upsampling layer to increase height and width of the input.
              Uses PixelShuffle for upsampling.
              Args:
                  layer_input: The input tensor to upsample.
              Returns:
                  u: Upsampled input by a factor of 2.
              """
              u = keras.layers.UpSampling2D(size=2, interpolation='bilinear')(layer_input)
              u = keras.layers.Conv2D(self.gf, kernel_size=3, strides=1, padding='same')(u)
              u = keras.layers.PReLU(shared_axes=[1, 2])(u)
              return u
  
          # Low resolution image input
          img_lr = keras.Input(shape=self.lr_shape)
  
          # Pre-residual block
          c1 = keras.layers.Conv2D(self.gf, kernel_size=3, strides=1, padding='same')(img_lr)
          c1 = keras.layers.BatchNormalization()(c1)
          c1 = keras.layers.PReLU(shared_axes=[1, 2])(c1)
  
          # Propogate through residual blocks
          r = residual_block(c1, self.gf, 0)
          for idx in range(1, self.n_residual_blocks):
              r = residual_block(r, self.gf, idx)
  
          # Post-residual block
          c2 = keras.layers.Conv2D(self.gf, kernel_size=3, strides=1, padding='same')(r)
          c2 = keras.layers.BatchNormalization()(c2)
          c2 = keras.layers.Add()([c2, c1])
  
          # Upsampling
          # u1 = deconv2d(c2)
          # u2 = deconv2d(u1)
  
          # Generate high resolution output
          # gen_hr = keras.layers.Conv2D(3, kernel_size=3, strides=1, padding='same', activation='tanh')(u2)
          gen_hr = keras.layers.Conv2D(3, kernel_size=3, strides=1, padding='same', activation='tanh')(c2)
  
          return keras.models.Model(img_lr, gen_hr)
  
      def build_discriminator(self):
          """Builds a discriminator network based on the SRGAN design."""
  
          def d_block(layer_input, filters, strides=1, bn=True):
              """Discriminator layer block.
              Args:
                  layer_input: Input feature map for the convolutional block.
                  filters: Number of filters in the convolution.
                  strides: The stride of the convolution.
                  bn: Whether to use batch norm or not.
              """
              d = keras.layers.Conv2D(filters, kernel_size=3, strides=strides, padding='same')(layer_input)
              if bn:
                  d = keras.layers.BatchNormalization(momentum=0.8)(d)
              d = keras.layers.LeakyReLU(alpha=0.2)(d)
  
              return d
  
          # Input img
          d0 = keras.layers.Input(shape=self.hr_shape)
  
          d1 = d_block(d0, self.df, bn=False)
          d2 = d_block(d1, self.df, strides=2)
          d3 = d_block(d2, self.df)
          d4 = d_block(d3, self.df, strides=2)
          d5 = d_block(d4, self.df * 2)
          d6 = d_block(d5, self.df * 2, strides=2)
          d7 = d_block(d6, self.df * 2)
          d8 = d_block(d7, self.df * 2, strides=2)
  
          validity = keras.layers.Conv2D(1, kernel_size=1, strides=1, activation='sigmoid', padding='same')(d8)
  
          return keras.models.Model(d0, validity)
  