text: # SPDX-License-Identifier: Apache-2.0
  
  import os
  import sys
  import unittest
  import mock_keras2onnx
  import onnx
  import numpy as np
  from mock_keras2onnx.proto import keras
  from keras.applications import VGG19
  from os.path import dirname, abspath
  sys.path.insert(0, os.path.join(dirname(abspath(__file__)), '../../keras2onnx_tests/'))
  from test_utils import run_onnx_runtime
  
  Activation = keras.layers.Activation
  Add = keras.layers.Add
  BatchNormalization = keras.layers.BatchNormalization
  Conv2D = keras.layers.Conv2D
  Dense = keras.layers.Dense
  Dropout = keras.layers.Dropout
  Embedding = keras.layers.Embedding
  Flatten = keras.layers.Flatten
  Input = keras.layers.Input
  LeakyReLU = keras.layers.LeakyReLU
  multiply = keras.layers.multiply
  Reshape = keras.layers.Reshape
  UpSampling2D = keras.layers.UpSampling2D
  ZeroPadding2D = keras.layers.ZeroPadding2D
  
  Sequential = keras.models.Sequential
  Model = keras.models.Model
  
  # From https://github.com/eriklindernoren/Keras-GAN/blob/master/srgan/srgan.py
  class SRGAN():
      def __init__(self):
          # Input shape
          self.channels = 3
          self.lr_height = 64                 # Low resolution height
          self.lr_width = 64                  # Low resolution width
          self.lr_shape = (self.lr_height, self.lr_width, self.channels)
          self.hr_height = self.lr_height*4   # High resolution height
          self.hr_width = self.lr_width*4     # High resolution width
          self.hr_shape = (self.hr_height, self.hr_width, self.channels)
  
          # Number of residual blocks in the generator
          self.n_residual_blocks = 16
  
          # 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 = 64
          self.df = 64
  
          # Build and compile the discriminator
          self.discriminator = self.build_discriminator()
  
          # Build the generator
          self.generator = self.build_generator()
  
          # High res. and low res. images
          img_hr = Input(shape=self.hr_shape)
          img_lr = Input(shape=self.lr_shape)
  
          # Generate high res. version from low res.
          fake_hr = self.generator(img_lr)
  
          # Extract image features of the generated img
          fake_features = self.vgg(fake_hr)
  
          # For the combined model we will only train the generator
          self.discriminator.trainable = False
  
          # Discriminator determines validity of generated high res. images
          validity = self.discriminator(fake_hr)
  
          self.combined = Model([img_lr, img_hr], [validity, fake_features])
  
      def build_vgg(self):
          """
          Builds a pre-trained VGG19 model that outputs image features extracted at the
          third block of the model
          """
          vgg = VGG19(weights="imagenet")
          # Set outputs to outputs of last conv. layer in block 3
          # See architecture at: https://github.com/keras-team/keras/blob/master/keras/applications/vgg19.py
          vgg.outputs = [vgg.layers[9].output]
  
          img = Input(shape=self.hr_shape)
  
          # Extract image features
          img_features = vgg(img)
  
          return Model(img, img_features)
  
      def build_generator(self):
  
          def residual_block(layer_input, filters):
              """Residual block described in paper"""
              d = Conv2D(filters, kernel_size=3, strides=1, padding='same')(layer_input)
              d = Activation('relu')(d)
              d = BatchNormalization(momentum=0.8)(d)
              d = Conv2D(filters, kernel_size=3, strides=1, padding='same')(d)
              d = BatchNormalization(momentum=0.8)(d)
              d = Add()([d, layer_input])
              return d
  
          def deconv2d(layer_input):
              """Layers used during upsampling"""
              u = UpSampling2D(size=2)(layer_input)
              u = Conv2D(256, kernel_size=3, strides=1, padding='same')(u)
              u = Activation('relu')(u)
              return u
  
          # Low resolution image input
          img_lr = Input(shape=self.lr_shape)
  
          # Pre-residual block
          c1 = Conv2D(64, kernel_size=9, strides=1, padding='same')(img_lr)
          c1 = Activation('relu')(c1)
  
          # Propogate through residual blocks
          r = residual_block(c1, self.gf)
          for _ in range(self.n_residual_blocks - 1):
              r = residual_block(r, self.gf)
  
          # Post-residual block
          c2 = Conv2D(64, kernel_size=3, strides=1, padding='same')(r)
          c2 = BatchNormalization(momentum=0.8)(c2)
          c2 = Add()([c2, c1])
  
          # Upsampling
          u1 = deconv2d(c2)
          u2 = deconv2d(u1)
  
          # Generate high resolution output
          gen_hr = Conv2D(self.channels, kernel_size=9, strides=1, padding='same', activation='tanh')(u2)
  
          return Model(img_lr, gen_hr)
  
      def build_discriminator(self):
  
          def d_block(layer_input, filters, strides=1, bn=True):
              """Discriminator layer"""
              d = Conv2D(filters, kernel_size=3, strides=strides, padding='same')(layer_input)
              d = LeakyReLU(alpha=0.2)(d)
              if bn:
                  d = BatchNormalization(momentum=0.8)(d)
              return d
  
          # Input img
          d0 = 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*2)
          d4 = d_block(d3, self.df*2, strides=2)
          d5 = d_block(d4, self.df*4)
          d6 = d_block(d5, self.df*4, strides=2)
          d7 = d_block(d6, self.df*8)
          d8 = d_block(d7, self.df*8, strides=2)
  
          d9 = Dense(self.df*16)(d8)
          d10 = LeakyReLU(alpha=0.2)(d9)
          validity = Dense(1, activation='sigmoid')(d10)
  
          return Model(d0, validity)
  
  
  class TestSRGAN(unittest.TestCase):
  
      def setUp(self):
          self.model_files = []
  
      def tearDown(self):
          for fl in self.model_files:
              os.remove(fl)
  
      def test_SRGAN(self):
          keras_model = SRGAN().combined
          x = np.random.rand(5, 64, 64, 3).astype(np.float32)
          y = np.random.rand(5, 256, 256, 3).astype(np.float32)
          expected = keras_model.predict([x, y])
          onnx_model = mock_keras2onnx.convert_keras(keras_model, keras_model.name)
          self.assertTrue(run_onnx_runtime(onnx_model.graph.name, onnx_model, {keras_model.input_names[0]: x, keras_model.input_names[1]: y}, expected, self.model_files, rtol=1.e-2, atol=1.e-3))
  
  
  if __name__ == "__main__":
      unittest.main()
  