text: import os
  import sys
  import numpy as np
  from datetime import datetime
  
  import tensorflow as tf
  from keras.models import Model
  from keras.models import load_model
  from keras.optimizers import Adam
  from keras.layers import Input, Conv2D, UpSampling2D, Dropout, LeakyReLU, BatchNormalization, Activation, Lambda
  from keras.layers.merge import Concatenate
  from keras.applications import VGG16
  from keras import backend as K
  from keras.utils.multi_gpu_utils import multi_gpu_model
  from keras.layers import MaxPool2D
  
  from libs.pconv_layer import PConv2D
  
  
  class PConvUnet(object):
  
      def __init__(self, img_rows=512, img_cols=512, vgg_weights="imagenet", inference_only=False, net_name='default', gpus=1, vgg_device=None):
          """Create the PConvUnet. If variable image size, set img_rows and img_cols to None
          
          Args:
              img_rows (int): image height.
              img_cols (int): image width.
              vgg_weights (str): which weights to pass to the vgg network.
              inference_only (bool): initialize BN layers for inference.
              net_name (str): Name of this network (used in logging).
              gpus (int): How many GPUs to use for training.
              vgg_device (str): In case of training with multiple GPUs, specify which device to run VGG inference on.
                  e.g. if training on 8 GPUs, vgg inference could be off-loaded exclusively to one GPU, instead of
                  running on one of the GPUs which is also training the UNet.
          """
          
          # Settings
          self.img_rows = img_rows
          self.img_cols = img_cols
          self.img_overlap = 30
          self.inference_only = inference_only
          self.net_name = net_name
          self.gpus = gpus
          self.vgg_device = vgg_device
  
          # Scaling for VGG input
          self.mean = [0.485, 0.456, 0.406]
          self.std = [0.229, 0.224, 0.225]
  
          # Assertions
          assert self.img_rows >= 256, 'Height must be >256 pixels'
          assert self.img_cols >= 256, 'Width must be >256 pixels'
  
          # Set current epoch
          self.current_epoch = 0
          
          # VGG layers to extract features from (first maxpooling layers, see pp. 7 of paper)
          self.vgg_layers = [3, 6, 10]
  
          # Instantiate the vgg network
          if self.vgg_device:
              with tf.device(self.vgg_device):
                  self.vgg = self.build_vgg(vgg_weights)
          else:
              self.vgg = self.build_vgg(vgg_weights)
          
          # Create UNet-like model
          if self.gpus <= 1:
              self.model, inputs_mask = self.build_pconv_unet()
              self.compile_pconv_unet(self.model, inputs_mask)            
          else:
              with tf.device("/cpu:0"):
                  self.model, inputs_mask = self.build_pconv_unet()
              self.model = multi_gpu_model(self.model, gpus=self.gpus)
              self.compile_pconv_unet(self.model, inputs_mask)
          
      def build_vgg(self, weights="imagenet"):
          """
          Load pre-trained VGG16 from keras applications
          Extract features to be used in loss function from last conv layer, see architecture at:
          https://github.com/keras-team/keras/blob/master/keras/applications/vgg16.py
          """        
              
          # Input image to extract features from
          img = Input(shape=(self.img_rows, self.img_cols, 3))
  
          # Mean center and rescale by variance as in PyTorch
          processed = Lambda(lambda x: (x-self.mean) / self.std)(img)
          
          # If inference only, just return empty model        
          if self.inference_only:
              model = Model(inputs=img, outputs=[img for _ in range(len(self.vgg_layers))])
              model.trainable = False
              model.compile(loss='mse', optimizer='adam')
              return model
                  
          # Get the vgg network from Keras applications
          if weights in ['imagenet', None]:
              vgg = VGG16(weights=weights, include_top=False)
          else:
              vgg = VGG16(weights=None, include_top=False)
              vgg.load_weights(weights, by_name=True)
  
          # Output the first three pooling layers
          vgg.outputs = [vgg.layers[i].output for i in self.vgg_layers]        
          
          # Create model and compile
          model = Model(inputs=img, outputs=vgg(processed))
          model.trainable = False
          model.compile(loss='mse', optimizer='adam')
  
          return model
          
      def build_pconv_unet(self, train_bn=True):      
  
          # INPUTS
          inputs_img = Input((self.img_rows, self.img_cols, 3), name='inputs_img')
          inputs_mask = Input((self.img_rows, self.img_cols, 3), name='inputs_mask')
          
          # ENCODER
          def encoder_layer(img_in, mask_in, filters, kernel_size, bn=True):
              conv, mask = PConv2D(filters, kernel_size, strides=2, padding='same')([img_in, mask_in])
              if bn:
                  conv = BatchNormalization(name='EncBN'+str(encoder_layer.counter))(conv, training=train_bn)
              conv = Activation('relu')(conv)
              encoder_layer.counter += 1
              return conv, mask
          def encoder_layer_1by1(img_in, mask_in, filters, bn=True):
              conv, mask = PConv2D(filters, kernel_size=1, strides=1, padding='same')([img_in, mask_in])
              if bn:
                  conv = BatchNormalization(name='EncBN'+str(encoder_layer.counter))(conv, training=train_bn)
              conv = Activation('relu')(conv)
              encoder_layer.counter += 1
              return conv, mask
          encoder_layer.counter = 0
  
          e_conv1a, e_mask1a = encoder_layer_1by1(inputs_img, inputs_mask, 64, bn=False)
          e_conv1a, e_mask1a = encoder_layer(e_conv1a, e_mask1a, 64, 5)
          e_conv1b, e_mask1b = encoder_layer_1by1(inputs_img, inputs_mask, 64, bn=False)
          e_conv1b, e_mask1b = encoder_layer(e_conv1b, e_mask1b, 64, 3)
          e_conv1c, e_mask1c = encoder_layer(inputs_img, inputs_mask, 64, 1, bn=False)
          #e_conv1 = Concatenate(axis=3)([e_conv1a, e_conv1b, e_conv1c])
          #e_mask1 = Concatenate(axis=3)([e_mask1a, e_mask1b, e_mask1c])
          e_conv1d = MaxPool2D(2,None,padding='same')(inputs_img)
          e_mask1d = MaxPool2D(2,None,padding='same')(inputs_mask)
          e_conv1d, e_mask1d = encoder_layer_1by1(e_conv1d, e_mask1d, 64)
          e_conv1 = Concatenate(axis=3)([e_conv1a, e_conv1b, e_conv1c,e_conv1d])
          e_mask1 = Concatenate(axis=3)([e_mask1a, e_mask1b, e_mask1c,e_mask1d])
  
          def inception(inputs_img, inputs_mask, filters):
              e_conva, e_maska = encoder_layer_1by1(inputs_img, inputs_mask, 64)
              e_conva, e_maska = encoder_layer(e_conva, e_maska, filters, 5)
              e_convb, e_maskb = encoder_layer_1by1(inputs_img, inputs_mask, 64)
              e_convb, e_maskb = encoder_layer(e_convb, e_maskb, filters, 3)
              e_convc, e_maskc = encoder_layer(inputs_img, inputs_mask, filters, 1)
              #e_conv = Concatenate(axis=3)([e_conva, e_convb, e_convc])
              #e_mask = Concatenate(axis=3)([e_maska, e_maskb, e_maskc])
              e_convd = MaxPool2D(2, None, padding='same')(inputs_img)
              e_maskd = MaxPool2D(2, None, padding='same')(inputs_mask)
              e_convd, e_maskd = encoder_layer_1by1(e_convd, e_maskd, filters)
              e_conv = Concatenate(axis=3)([e_conva, e_convb, e_convc, e_convd])
              e_mask = Concatenate(axis=3)([e_maska, e_maskb, e_maskc, e_maskd])
              return e_conv, e_mask            
  
          e_conv2, e_mask2 = inception(e_conv1, e_mask1, 128)
          e_conv3, e_mask3 = inception(e_conv2, e_mask2, 256)
          e_conv4, e_mask4 = inception(e_conv3, e_mask3, 512)
          #e_conv5, e_mask5 = inception(e_conv4, e_mask4, 512)
          #e_conv6, e_mask6 = inception(e_conv5, e_mask5, 512)
          #e_conv7, e_mask7 = inception(e_conv6, e_mask6, 512)
          #e_conv8, e_mask8 = inception(e_conv7, e_mask7, 512)
          
          
          # DECODER Inception version [2020 WNY]
          def decoder_layer(img_in, mask_in, e_conv, e_mask, filters, bn=True):
              up_img = UpSampling2D(size=(2,2))(img_in)
              up_mask = UpSampling2D(size=(2,2))(mask_in)
              concat_img = Concatenate(axis=3)([e_conv,up_img])
              concat_mask = Concatenate(axis=3)([e_mask,up_mask])
              conva, maska = PConv2D(filters=64, kernel_size=1, padding='same')([concat_img, concat_mask])
              conva, maska = PConv2D(filters, kernel_size=5, padding='same')([conva, maska])
              convb, maskb = PConv2D(filters=64, kernel_size=1, padding='same')([concat_img, concat_mask])
              convb, maskb = PConv2D(filters, kernel_size=3, padding='same')([convb, maskb])
              convc, maskc = PConv2D(filters, kernel_size=1, padding='same')([concat_img, concat_mask])
              conv = Concatenate(axis=3)([conva, convb, convc])
              mask = Concatenate(axis=3)([maska, maskb, maskc])
  
              if bn:
                  conv = BatchNormalization()(conv)
              conv = LeakyReLU(alpha=0.2)(conv)
              return conv, mask
              
          #d_conv9, d_mask9 = decoder_layer(e_conv8, e_mask8, e_conv7, e_mask7, 512)
          #d_conv10, d_mask10 = decoder_layer(d_conv9, d_mask9, e_conv6, e_mask6, 512)
          #d_conv11, d_mask11 = decoder_layer(d_conv10, d_mask10, e_conv5, e_mask5, 512)
          #d_conv12, d_mask12 = decoder_layer(d_conv11, d_mask11, e_conv4, e_mask4, 512)
          #d_conv13, d_mask13 = decoder_layer(d_conv12, d_mask12, e_conv3, e_mask3, 256)
          #d_conv14, d_mask14 = decoder_layer(d_conv13, d_mask13, e_conv2, e_mask2, 128)
          #d_conv15, d_mask15 = decoder_layer(d_conv14, d_mask14, e_conv1, e_mask1, 64)
          #d_conv16, d_mask16 = decoder_layer(d_conv15, d_mask15, inputs_img, inputs_mask, 3, bn=False)
          #outputs = Conv2D(3, 1, activation = 'sigmoid', name='outputs_img')(d_conv16)
          d_conv1, d_mask1 = decoder_layer(e_conv4, e_mask4, e_conv3, e_mask3, 256)
          d_conv2, d_mask2 = decoder_layer(d_conv1, d_mask1, e_conv2, e_mask2, 128)
          d_conv3, d_mask3 = decoder_layer(d_conv2, d_mask2, e_conv1, e_mask1, 64)
          d_conv4, d_mask4 = decoder_layer(d_conv3, d_mask3, inputs_img, inputs_mask, 3, bn=False)        
          outputs = Conv2D(3, 1, activation = 'sigmoid', name='outputs_img')(d_conv4)
          
          # Setup the model inputs / outputs
          model = Model(inputs=[inputs_img, inputs_mask], outputs=outputs)
  
          return model, inputs_mask    
  
      def compile_pconv_unet(self, model, inputs_mask, lr=0.0002):
          model.compile(
              optimizer = Adam(lr=lr),
              loss=self.loss_total(inputs_mask),
              metrics=[self.PSNR]
          )
  
      def loss_total(self, mask):
          """
          Creates a loss function which sums all the loss components 
          and multiplies by their weights. See paper eq. 7.
          """
          def loss(y_true, y_pred):
  
              # Compute predicted image with non-hole pixels set to ground truth
              y_comp = mask * y_true + (1-mask) * y_pred
  
              # Compute the vgg features. 
              if self.vgg_device:
                  with tf.device(self.vgg_device):
                      vgg_out = self.vgg(y_pred)
                      vgg_gt = self.vgg(y_true)
                      vgg_comp = self.vgg(y_comp)
              else:
                  vgg_out = self.vgg(y_pred)
                  vgg_gt = self.vgg(y_true)
                  vgg_comp = self.vgg(y_comp)
              
              # Compute loss components
              l1 = self.loss_valid(mask, y_true, y_pred)
              l2 = self.loss_hole(mask, y_true, y_pred)
              l3 = self.loss_perceptual(vgg_out, vgg_gt, vgg_comp)
              l4 = self.loss_style(vgg_out, vgg_gt)
              l5 = self.loss_style(vgg_comp, vgg_gt)
              l6 = self.loss_tv(mask, y_comp)
  
              # Return loss function
              return l1 + 6*l2 + 0.05*l3 + 120*(l4+l5) + 0.1*l6
  
          return loss
  
      def loss_hole(self, mask, y_true, y_pred):
          """Pixel L1 loss within the hole / mask"""
          return self.l1((1-mask) * y_true, (1-mask) * y_pred)
      
      def loss_valid(self, mask, y_true, y_pred):
          """Pixel L1 loss outside the hole / mask"""
          return self.l1(mask * y_true, mask * y_pred)
      
      def loss_perceptual(self, vgg_out, vgg_gt, vgg_comp): 
          """Perceptual loss based on VGG16, see. eq. 3 in paper"""       
          loss = 0
          for o, c, g in zip(vgg_out, vgg_comp, vgg_gt):
              loss += self.l1(o, g) + self.l1(c, g)
          return loss
          
      def loss_style(self, output, vgg_gt):
          """Style loss based on output/computation, used for both eq. 4 & 5 in paper"""
          loss = 0
          for o, g in zip(output, vgg_gt):
              loss += self.l1(self.gram_matrix(o), self.gram_matrix(g))
          return loss
      
      def loss_tv(self, mask, y_comp):
          """Total variation loss, used for smoothing the hole region, see. eq. 6"""
  
          # Create dilated hole region using a 3x3 kernel of all 1s.
          kernel = K.ones(shape=(3, 3, mask.shape[3], mask.shape[3]))
          dilated_mask = K.conv2d(1-mask, kernel, data_format='channels_last', padding='same')
  
          # Cast values to be [0., 1.], and compute dilated hole region of y_comp
          dilated_mask = K.cast(K.greater(dilated_mask, 0), 'float32')
          P = dilated_mask * y_comp
  
          # Calculate total variation loss
          a = self.l1(P[:,1:,:,:], P[:,:-1,:,:])
          b = self.l1(P[:,:,1:,:], P[:,:,:-1,:])        
          return a+b
  
      def fit_generator(self, generator, *args, **kwargs):
          """Fit the U-Net to a (images, targets) generator
  
          Args:
              generator (generator): generator supplying input image & mask, as well as targets.
              *args: arguments to be passed to fit_generator
              **kwargs: keyword arguments to be passed to fit_generator
          """
          self.model.fit_generator(
              generator,
              *args, **kwargs
          )
          
      def summary(self):
          """Get summary of the UNet model"""
          print(self.model.summary())
  
      def load(self, filepath, train_bn=True, lr=0.0002):
  
          # Create UNet-like model
          self.model, inputs_mask = self.build_pconv_unet(train_bn)
          self.compile_pconv_unet(self.model, inputs_mask, lr) 
  
          # Load weights into model
          epoch = int(os.path.basename(filepath).split('.')[1].split('-')[0])
          assert epoch > 0, "Could not parse weight file. Should include the epoch"
          self.current_epoch = epoch
          self.model.load_weights(filepath)        
  
      @staticmethod
      def PSNR(y_true, y_pred):
          """
          PSNR is Peek Signal to Noise Ratio, see https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
          The equation is:
          PSNR = 20 * log10(MAX_I) - 10 * log10(MSE)
          
          Our input is scaled with be within the range -2.11 to 2.64 (imagenet value scaling). We use the difference between these
          two values (4.75) as MAX_I        
          """        
          #return 20 * K.log(4.75) / K.log(10.0) - 10.0 * K.log(K.mean(K.square(y_pred - y_true))) / K.log(10.0) 
          return - 10.0 * K.log(K.mean(K.square(y_pred - y_true))) / K.log(10.0) 
  
      @staticmethod
      def current_timestamp():
          return datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
      
      @staticmethod
      def l1(y_true, y_pred):
          """Calculate the L1 loss used in all loss calculations"""
          if K.ndim(y_true) == 4:
              return K.mean(K.abs(y_pred - y_true), axis=[1,2,3])
          elif K.ndim(y_true) == 3:
              return K.mean(K.abs(y_pred - y_true), axis=[1,2])
          else:
              raise NotImplementedError("Calculating L1 loss on 1D tensors? should not occur for this network")
      
      @staticmethod
      def gram_matrix(x, norm_by_channels=False):
          """Calculate gram matrix used in style loss"""
          
          # Assertions on input
          assert K.ndim(x) == 4, 'Input tensor should be a 4d (B, H, W, C) tensor'
          assert K.image_data_format() == 'channels_last', "Please use channels-last format"        
          
          # Permute channels and get resulting shape
          x = K.permute_dimensions(x, (0, 3, 1, 2))
          shape = K.shape(x)
          B, C, H, W = shape[0], shape[1], shape[2], shape[3]
          
          # Reshape x and do batch dot product
          features = K.reshape(x, K.stack([B, C, H*W]))
          gram = K.batch_dot(features, features, axes=2)
          
          # Normalize with channels, height and width
          gram = gram /  K.cast(C * H * W, x.dtype)
          
          return gram
      
      # Prediction functions
      ######################
      def predict(self, sample, **kwargs):
          """Run prediction using this model"""
          return self.model.predict(sample, **kwargs)
  