progresssive_MSG_capsgan.py

# -*- coding: utf-8 -*-
"""Final_Progresssive_CapsGAN_features.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/github/MahdiyarMM/Colab/blob/master/Final_Progresssive_CapsGAN_features.ipynb
"""

from google.colab import drive
drive.mount('/content/drive')

#Image lib
import cv2
# math libraries
import numpy as np
import scipy.misc
# ml libraries
import tensorflow as tf
from keras import layers, models, optimizers
from keras import backend as K
from keras.utils import to_categorical
from keras.datasets import mnist, cifar10
from keras.layers import Input, Dense, Reshape, Flatten, Dropout, Lambda, concatenate, Multiply
from keras.layers import BatchNormalization, Activation, ZeroPadding2D, Add, RepeatVector, AveragePooling2D
from keras.layers.advanced_activations import LeakyReLU
from keras.layers.convolutional import UpSampling2D, Conv2D
from keras.models import Sequential, Model
from keras.optimizers import Adam
from keras.preprocessing.image import ImageDataGenerator
from keras import callbacks
import datetime
# visualization
import skimage
from skimage import data, color, exposure
from skimage.transform import resize
import matplotlib.pyplot as plt
import math
from skimage.measure import compare_ssim as ssim
from skimage.measure import compare_psnr as psnr
from keras.applications import VGG19
from keras.backend import int_shape, tile


# sys and helpers
import sys
import os
import glob
from tqdm import tqdm

print('Modules imported.')

class DataLoader3_all():
    def __init__(self, dataset_name, img_res=(128, 128)):
        self.dataset_name = dataset_name
        self.img_res = img_res

    def load_data(self, batch_size=1, is_testing=False):
        data_type = "train" if not is_testing else "test"
        
        path = glob.glob('/content/drive/My Drive/CelebA_sample/CelebA/*')

        batch_images = np.random.choice(path, size=batch_size)

        imgs_hr = []
        imgs_lr = []
        imgs_32 = []
        imgs_64 = []
        for img_path in batch_images:
            imgl = self.imread(img_path)
            img = imgl[20:218-20,:]
            h, w = self.img_res
            low_h, low_w = int(h / 8), int(w / 8)

#            img_hr = scipy.misc.imresize(img, self.img_res)
 #           img_lr = scipy.misc.imresize(img, (low_h, low_w))
  #          img_32 = scipy.misc.imresize(img, (32, 32))
   #         img_64 = scipy.misc.imresize(img, (64, 64))

            img_hr = cv2.resize(img, self.img_res)
            img_lr = cv2.resize(img, (low_h, low_w))
            img_32 = cv2.resize(img, (32, 32))
            img_64 = cv2.resize(img, (64, 64))
            # If training => do random flip
            if not is_testing and np.random.random() < 0.5:
                img_hr = np.fliplr(img_hr)
                img_lr = np.fliplr(img_lr)
                img_32 = np.fliplr(img_32)
                img_64 = np.fliplr(img_64)

            imgs_hr.append(img_hr)
            imgs_lr.append(img_lr)
            imgs_32.append(img_32)
            imgs_64.append(img_64)

        imgs_hr = np.array(imgs_hr) / 127.5 - 1.
        imgs_lr = np.array(imgs_lr) / 127.5 - 1.
        imgs_32 = np.array(imgs_32) / 127.5 - 1.
        imgs_64 = np.array(imgs_64) / 127.5 - 1.

        return imgs_hr, imgs_lr , imgs_32, imgs_64


    def imread(self, path):
        return plt.imread(path).astype(np.float)
    
    def load_test(self,path):
        img = self.imread(path)
        img_hr = scipy.misc.imresize(img, (256,256))  

# Configure data loader
dataset_name = 'img_align_celeba'
hr_height = 128
hr_width = 128

data_loader = DataLoader3_all(dataset_name=dataset_name,
                              img_res=(hr_height, hr_width))

def all_psnr(imageA, imageB):
    psnrs = []
    for ii in range(len(imageA)):
        psnrs.append(psnr(imageA[ii],imageB[ii]))
    return np.mean(psnrs)

def all_ssim(imageA, imageB):
    psnrs = []
    for ii in range(len(imageA)):
        psnrs.append(ssim(imageA[ii],imageB[ii],multichannel =True))
    return np.mean(psnrs)

        
# squash function of capsule layers, borrowed from Xifeng Guo's implementation of Keras CapsNet `https://github.com/XifengGuo/CapsNet-Keras`
def squash(vectors, axis=-1):
    """
    The non-linear activation used in Capsule. It drives the length of a large vector to near 1 and small vector to 0
    :param vectors: some vectors to be squashed, N-dim tensor
    :param axis: the axis to squash
    :return: a Tensor with same shape as input vectors
    """
    s_squared_norm = K.sum(K.square(vectors), axis, keepdims=True)
    scale = s_squared_norm / (1 + s_squared_norm) / K.sqrt(s_squared_norm + K.epsilon())
    return scale * vectors

# device check
from tensorflow.python.client import device_lib
print('Devices:', device_lib.list_local_devices())

# GPU check
if not tf.test.gpu_device_name():
    print('No GPU found.')
else:
    print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))


def get_a_and_b(epoch):
  if epoch<5000:
    a=1
    b=1
  elif epoch<6500:
    a = (6500-epoch)/1500
    b = 1
  elif epcoh <11500:
    a=0
    b=1
  elif epoch <13000:
    a=0
    b = (13500-epoch)/1500
  else:
    a=0
    b=0
  return a,b

path = "/content/drive/My Drive/Keras_models"

def build_vgg(hr_shape):
    """
    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(hr_shape)

    # Extract image features
    img_features = vgg(img)

    return Model(img, img_features)

optimizer = Adam(0.0003, 0.5)
vgg = build_vgg((128,128,3))
vgg.trainable = False
vgg.compile(loss='mse',
            optimizer=optimizer,
            metrics=['accuracy'])

vgg64 = build_vgg((64,64,3))
vgg64.trainable = False
vgg64.compile(loss='mse',
            optimizer=optimizer,
            metrics=['accuracy'])

vgg32 = build_vgg((32,32,3))
vgg32.trainable = False
vgg32.compile(loss='mse',
            optimizer=optimizer,
            metrics=['accuracy'])

# Define custom layers
def tile_w():
  def tw(w,x):
    new_shape = (int_shape(x)[1],int_shape(x)[2])
    tiled = Lambda(lambda image: tf.image.resize_images(image,new_shape,method = tf.image.ResizeMethod.NEAREST_NEIGHBOR,align_corners = True, preserve_aspect_ratio = True))(w)
    if int_shape(x)[3] == 3:
      tiled = concatenate([tiled,tiled,tiled],axis=-1)
    else:
      for ii in range(int(np.log2(int_shape(x)[3]))):
        tiled = concatenate([tiled,tiled],axis=-1)
    return tiled
  return tw


def weighted_Add():
  def wa(input_1,input_2,w_val):
    #retruns input_1*w + input_2*(1-w)
    w = tile_w()(w_val,input_1)
    one_w = Lambda(lambda x: 1-x)(w)  
    i1_w1 = Multiply()([input_1,w])
    i2_w2 = Multiply()([input_2,one_w])
    added = Add()([i1_w1,i2_w2])
    return added
  return wa

def resize_img(new_shape):
  def res(img_lr_in):
    img_resized =  Lambda(lambda image: tf.image.resize_images(image,new_shape,method = tf.image.ResizeMethod.BICUBIC,align_corners = True, preserve_aspect_ratio = True))(img_lr_in) 
    return img_resized
  return res

def resize_img_NN(new_shape):
  def res(img_lr_in):
    img_resized =  Lambda(lambda image: tf.image.resize_images(image,new_shape,method = tf.image.ResizeMethod.NEAREST_NEIGHBOR,align_corners = True, preserve_aspect_ratio = True))(img_lr_in) 
    return img_resized
  return res

def Dynamic_routing():
  def dr(s_j,uhat):
    c = Activation('softmax')(s_j) 
    c = Dense(160)(c)
    x = Multiply()([uhat, c])
    s_j = LeakyReLU()(x)
    return s_j
  return dr

def Capslayer(casp_num,caps_dim):
  def CP(x):
    x1 = Conv2D(filters=64, kernel_size=9, strides=2, padding='valid', name='primarycap_conv2')(x)
    x2 = Flatten()(x1)
    x22 = Dense(casp_num*caps_dim)(x2)
    x3 = Reshape(target_shape=[-1, caps_dim])(x2)
    x4 = Lambda(squash, name='primarycap_squash')(x3)
    x5 = BatchNormalization(momentum=0.8)(x4)
    x6 = Flatten()(x5)
    return x6
  return CP

def Last_Caps(casp_num,caps_dim):
  def LC(x):
    uhat = Dense(160, kernel_initializer='he_normal', bias_initializer='zeros')(x)
    return uhat
  return LC

def Conv_LR(filters_num=128, kernel_size_num=9, strides_num=2):
  def CL(img):
    x1 = Conv2D(filters=filters_num, kernel_size=kernel_size_num, strides=strides_num, padding='same')(img)
    x11 = LeakyReLU()(x1)
    return x11
  return CL

def deconv2d(layer_input,num=256):
    """Layers used during upsampling"""
    u = UpSampling2D(size=2)(layer_input)
    u = Conv2D(num, kernel_size=3, strides=1, padding='same')(u)
    u = Activation('relu')(u)
    return u

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 process_block():
  def pb(input_layer):
    pd = Conv2D(64, kernel_size=3, strides=1, padding='same')(input_layer)
    x1 = residual_block(pd,64)
    x2 = deconv2d(x1,num=64)
    return x2
  return pb
  
'''  
def process_block():
  def pb(input_layer,vgg_features):
    pd = Conv2D(64, kernel_size=3, strides=1, padding='same')(input_layer)
    x1 = residual_block(pd,64)
    x2 = deconv2d(x1,num=64)
    vgg_features_cnn = Conv2D(64, kernel_size=3, strides=1, padding='same')(vgg_features)
    xconcat = concatenate([x2,vgg_features_cnn],axis=-1)
    return xconcat
  return pb
'''

def to_rgb():
  def tr(x):
    x2 = Conv2D(3, kernel_size=3, strides=1, padding='same',activation="tanh")(x)
    return x2
  return tr

def build_resize(in_shape,out_shape):
  img_in = Input(in_shape)
  img_out = Lambda(lambda image: tf.image.resize_images(image,out_shape,method = tf.image.ResizeMethod.NEAREST_NEIGHBOR,align_corners = True, preserve_aspect_ratio = True))(img_in)
  return Model(img_in,img_out)

# discriminator 1 structure
def build_discriminator_1():


    img = Input(shape=(32, 32, 3))
    
    x_16 = Conv_LR(strides_num = 1)(img)
    # original 'Dynamic Routing Between Capsules' paper does not include the batch norm layer after the first conv group
    x = BatchNormalization(momentum=0.8)(x_16)

    #fully connected layer befor Capsnet
    

    

    #Capsule Network
    x_caps = Capslayer(50,8)(x)
    uhat = Last_Caps(10,16)(x_caps)
    s_j_1 = Dynamic_routing()(uhat,uhat)
    s_j_2 = Dynamic_routing()(s_j_1,uhat) 
    s_j_3 = Dynamic_routing()(s_j_2,uhat) 
 
    #Patch
    pred = Dense(25*25, activation='sigmoid')(s_j_3)
    patch = Reshape(target_shape=[-1, 25], name='last_out_reshape')(pred)
    
    
    return Model([img], patch)



# build and compile the discriminator
discriminator_1 = build_discriminator_1()
print('DISCRIMINATOR:')
discriminator_1.summary()
sgd = optimizers.SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)
#discriminator_1.load_weights('"/content/drive/My Drive/Keras_models/images_MSG/wdisc5500.h5")
discriminator_1.compile(loss='binary_crossentropy', optimizer=sgd)

# discriminator 2 structure
def build_discriminator_2():


    img = Input(shape=(64, 64, 3))
    alpha = Input(shape=(1,1,1))
    x_32 = Conv_LR(strides_num = 1)(img)
    x_32_d = AveragePooling2D(pool_size=(2,2))(x_32)
    x_32_img = to_rgb()(x_32_d)
    downsampled = AveragePooling2D(pool_size=(2,2))(img)
    x_32_added = weighted_Add()(downsampled,x_32_img,alpha)
    patch = discriminator_1(x_32_added)
    
    
    return Model([img,alpha], patch)



# build and compile the discriminator
discriminator_2 = build_discriminator_2()
print('DISCRIMINATOR:')
discriminator_2.summary()
sgd = optimizers.SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)
#discriminator_1.load_weights('"/content/drive/My Drive/Keras_models/images_MSG/wdisc5500.h5")
discriminator_2.compile(loss='binary_crossentropy', optimizer=sgd)

# discriminator 3 structure
def build_discriminator_3():


    img = Input(shape=(128, 128, 3))
    alpha = Input(shape=(1,1,1))
    beta = Input(shape=(1,1,1))
    x_64 = Conv_LR(strides_num = 1)(img)
    x_64_d = AveragePooling2D(pool_size=(2,2))(x_64)
    x_64_img = to_rgb()(x_64_d)
    downsampled = AveragePooling2D(pool_size=(2,2))(img)
    x_64_added = weighted_Add()(downsampled,x_64_img,beta)
    patch = discriminator_2([x_64_added,alpha])
    
    
    return Model([img,alpha,beta], patch)



# build and compile the discriminator
discriminator_3 = build_discriminator_3()
print('DISCRIMINATOR:')
discriminator_3.summary()
sgd = optimizers.SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)

#discriminator_1.load_weights('"/content/drive/My Drive/Keras_models/images_MSG/wdisc5500.h5")
discriminator_3.compile(loss='binary_crossentropy', optimizer=sgd)

# discriminator 3 structure
def build_discriminator_3_b():


    img = Input(shape=(128, 128, 3))
    alpha = Input(shape=(1,1,1))
    beta = Input(shape=(1,1,1))
    x_64 = Conv_LR(strides_num = 1)(img)
    x_64_d = AveragePooling2D(pool_size=(2,2))(x_64)
    x_64_img = to_rgb()(x_64_d)
    downsampled = AveragePooling2D(pool_size=(2,2))(img)
    x_64_added = weighted_Add()(downsampled,x_64_img,beta)
    
    
    return Model([img,alpha,beta], x_64_added)

discriminator_3_b = build_discriminator_3_b()

def build_discriminator_3():
    img = Input(shape=(128, 128, 3))
    alpha = Input(shape=(1,1,1))
    beta = Input(shape=(1,1,1))
    x_64_added = discriminator_3_b([img,alpha,beta]) 
    patch = discriminator_2([x_64_added,alpha])

    return Model([img,alpha,beta], patch)


# build and compile the discriminator
discriminator_3 = build_discriminator_3()
print('DISCRIMINATOR:')
discriminator_3.summary()
sgd = optimizers.SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)

#discriminator_1.load_weights('"/content/drive/My Drive/Keras_models/images_MSG/wdisc5500.h5")
discriminator_3.compile(loss='binary_crossentropy', optimizer=sgd)

discriminator_3_b.save_weights("disc_b.h5")

#discriminator_3_b.save_weights("disc_b.h5")
discriminator_3_b.load_weights("disc_b.h5")
def build_discriminator_3():
    img = Input(shape=(128, 128, 3))
    alpha = Input(shape=(1,1,1))
    beta = Input(shape=(1,1,1))
    x_64_added = discriminator_3_b([img,alpha,beta]) 
    patch = discriminator_2([x_64_added,alpha])

    return Model([img,alpha,beta], patch)


# build and compile the discriminator
discriminator_3 = build_discriminator_3()
print('DISCRIMINATOR:')
discriminator_3.summary()
sgd = optimizers.SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)

#discriminator_1.load_weights('"/content/drive/My Drive/Keras_models/images_MSG/wdisc5500.h5")
discriminator_3.compile(loss='binary_crossentropy', optimizer=sgd)

discriminator_3_b.save_weights("disc_b.h5")

# generator 1  structure   
def build_generator_1():
  
    img_lr_in = Input(shape=(16,16,3))
    x_32 = process_block()(img_lr_in)
    x_32_img = to_rgb()(x_32)
    

    return Model([img_lr_in],[x_32_img])
    
    
# build and compile the generator
generator_1 = build_generator_1()
print('GENERATOR:')
generator_1.summary()
#generator_1.load_weights("/content/drive/My Drive/Keras_models/images_MSG/wgen5500.h5")
generator_1.compile(loss='binary_crossentropy', optimizer=Adam(0.0002, 0.5))

# generator 1  structure   
def build_generator_2():
    
    img_lr_in = Input(shape=(16,16,3))
    alpha = Input(shape=(1,1,1))
    img_32 = generator_1(img_lr_in)
    x_64 = process_block()(img_32)
    x_64_img = to_rgb()(x_64)
    img_64_upsampled = resize_img_NN((64,64))(img_32)
    x_64_added = weighted_Add()(img_64_upsampled,x_64_img,alpha)


    return Model([img_lr_in,alpha],[x_64_added])
    
    
# build and compile the generator
generator_2 = build_generator_2()
print('GENERATOR:')
generator_2.summary()
#generator_1.load_weights("/content/drive/My Drive/Keras_models/images_MSG/wgen5500.h5")
generator_2.compile(loss='binary_crossentropy', optimizer=Adam(0.0002, 0.5))

# generator 3  structure   
def build_generator_3():
    
    img_lr_in = Input(shape=(16,16,3))
    alpha = Input(shape=(1,1,1))
    beta = Input(shape=(1,1,1))
    img_64 = generator_2([img_lr_in,alpha])
    x_128 = process_block()(img_64)
    x_128_img = to_rgb()(x_128)
    img_128_upsampled = resize_img_NN((128,128))(img_64)
    x_128_added = weighted_Add()(img_128_upsampled,x_128_img,beta)


    return Model([img_lr_in,alpha,beta],[x_128_added])
    
    
# build and compile the generator
generator_3 = build_generator_3()
print('GENERATOR:')
generator_3.summary()
#generator_1.load_weights("/content/drive/My Drive/Keras_models/images_MSG/wgen5500.h5")
generator_3.compile(loss='binary_crossentropy', optimizer=Adam(0.0002, 0.5))

#creat combineds
#combined1
z = Input(shape=(16,16,3))
imgf = generator_1([z])
imgff = resize_img((128,128))(imgf)
fake_features = vgg(imgff)
discriminator_1.trainable = False
valid = discriminator_1([imgf])
combined_1 = Model([z], [valid,fake_features])
print('COMBINED:')
combined_1.summary()
combined_1.compile(loss=['binary_crossentropy', 'mse'], 
                 loss_weights=[1e-3, 1],
                 optimizer=Adam(0.002, 0.5))
#combined2
z = Input(shape=(16,16,3))
alpha = Input(shape=(1,1,1))

imgf = generator_2([z,alpha])
imgff = resize_img((128,128))(imgf)
fake_features = vgg(imgff)
discriminator_2.trainable = False
valid = discriminator_2([imgf,alpha])
combined_2 = Model([z,alpha], [valid,fake_features])
print('COMBINED:')
combined_2.summary()
combined_2.compile(loss=['binary_crossentropy', 'mse'], 
                 loss_weights=[1e-3, 1],
                 optimizer=Adam(0.002, 0.5))

#combined3
z = Input(shape=(16,16,3))
alpha = Input(shape=(1,1,1))
beta = Input(shape=(1,1,1))


imgf = generator_3([z,alpha,beta])
imgff = resize_img((128,128))(imgf)
fake_features = vgg(imgff)
discriminator_3.trainable = False
valid = discriminator_3([imgf,alpha,beta])
combined_3 = Model([z,alpha,beta], [valid,fake_features])
print('COMBINED:')
combined_3.summary()
combined_3.compile(loss=['binary_crossentropy', 'mse'], 
                 loss_weights=[1e-3, 1],
                 optimizer=Adam(0.002, 0.5))

def resize_img(new_shape):
  def res(img_lr_in):
    img_resized =  Lambda(lambda image: tf.image.resize_images(image,new_shape,method = tf.image.ResizeMethod.BICUBIC,align_corners = True, preserve_aspect_ratio = True))(img_lr_in) 
    return img_resized
  return res

def build_resize(input_shape = (32,32,3)):
  img_in = Input(input_shape)
  out = resize_img((128,128))(img_in)
  return Model(img_in,out)

resize32 = build_resize()
resize64 = build_resize(input_shape=(64,64,3))

# loss values for further plotting
D_L_REAL = []
D_L_FAKE = []
D_L = []
D_ACC = []
G_L = []

def train(dataset_title, epochs, batch_size=32, save_interval=50):
        start_time = datetime.datetime.now()
        half_batch = int(batch_size / 2)
        
        for epoch in range(0,50000):
            alpha_val,beta_val = get_a_and_b(epoch)
            if alpha_val==1 and beta_val==1:

              #  Train Discriminator
              # ---------------------

              # select a random half batch of images
              imgs_hr, imgs_lr, imgs_32, imgs_64= data_loader.load_data(half_batch)
              alpha = np.ones(shape = (half_batch,1,1,1))*alpha_val
              beta = np.ones(shape = (half_batch,1,1,1))*beta_val
              # generate a half batch of new images
              gen_imgs = generator_1.predict([imgs_lr])

              # train the discriminator by feeding both real and fake (generated) images one by one
              d_loss_real = discriminator_1.train_on_batch([imgs_32], np.random.uniform(low=0.8, high=1.2, size=(half_batch, 25,25))) # 0.9 for label smoothing
              d_loss_fake = discriminator_1.train_on_batch([gen_imgs], np.random.uniform(low=0.0, high=0.3, size=(half_batch, 25,25)))
              d_loss = 0.5 * np.add(d_loss_real, d_loss_fake)


              # ---------------------
              #  Train Generator
              # ---------------------


              imgs_hr, imgs_lr, imgs_32, imgs_64 = data_loader.load_data(batch_size)
              alpha = np.ones(shape = (batch_size,1,1,1))*alpha_val
              beta = np.ones(shape = (batch_size,1,1,1))*beta_val
              # the generator wants the discriminator to label the generated samples
              # as valid (ones)

              imgs_h = resize32.predict(imgs_32)
              image_features = vgg.predict(imgs_h)


              # train the generator
              g_loss = combined_1.train_on_batch([imgs_lr], [np.ones((batch_size,25,25)),image_features])
              elapsed_time = datetime.datetime.now() - start_time
              
              # Plot the progress
              print ("%d time: %s ,[D_acc.: %.2f%%] [G loss: %f] [alpha = %.2f beta = %.2f]" % (epoch,elapsed_time, 100*d_loss, g_loss[0],alpha_val,beta_val))
              D_L_REAL.append(d_loss_real)
              D_L_FAKE.append(d_loss_fake)
              D_L.append(d_loss)
              D_ACC.append(d_loss)
              G_L.append(g_loss)

              # if at save interval => save generated image samples
              if epoch % 10 == 0:
                  sample_images(epoch,alpha_val,beta_val)

              if epoch % (10*save_interval) == 0:    
                  generator_3.save_weights("images_P/wgen%d.h5" % (epoch))
                  discriminator_3.save_weights("images_P/wdis%d.h5" % (epoch))
                  sample_images(epoch,alpha_val,beta_val,drive = False)


            if alpha_val < 1 and beta_val==1:

              #  Train Discriminator
              # ---------------------

              # select a random half batch of images
              imgs_hr, imgs_lr, imgs_32, imgs_64= data_loader.load_data(half_batch)
              alpha = np.ones(shape = (half_batch,1,1,1))*alpha_val
              beta = np.ones(shape = (half_batch,1,1,1))*beta_val
              # generate a half batch of new images
              gen_imgs = generator_2.predict([imgs_lr,alpha])

              # train the discriminator by feeding both real and fake (generated) images one by one
              d_loss_real = discriminator_2.train_on_batch([imgs_64,alpha], np.random.uniform(low=0.8, high=1.2, size=(half_batch, 25,25))) # 0.9 for label smoothing
              d_loss_fake = discriminator_2.train_on_batch([gen_imgs,alpha], np.random.uniform(low=0.0, high=0.3, size=(half_batch, 25,25)))
              d_loss = 0.5 * np.add(d_loss_real, d_loss_fake)


              # ---------------------
              #  Train Generator
              # ---------------------


              imgs_hr, imgs_lr, imgs_32, imgs_64 = data_loader.load_data(batch_size)
              alpha = np.ones(shape = (batch_size,1,1,1))*alpha_val
              beta = np.ones(shape = (batch_size,1,1,1))*beta_val
              # the generator wants the discriminator to label the generated samples
              # as valid (ones)

              imgs_h = resize64.predict(imgs_64)
              image_features = vgg.predict(imgs_h)


              # train the generator
              g_loss = combined_2.train_on_batch([imgs_lr,alpha], [np.ones((batch_size,25,25)),image_features])
              elapsed_time = datetime.datetime.now() - start_time
              
              # Plot the progress
              print ("%d time: %s ,[D_acc.: %.2f%%] [G loss: %f] [alpha = %.2f beta = %.2f]" % (epoch,elapsed_time, 100*d_loss, g_loss[0],alpha_val,beta_val))
              D_L_REAL.append(d_loss_real)
              D_L_FAKE.append(d_loss_fake)
              D_L.append(d_loss)
              D_ACC.append(d_loss)
              G_L.append(g_loss)

              # if at save interval => save generated image samples
              if epoch % 10 == 0:
                  sample_images(epoch,alpha_val,beta_val)

              if epoch % (10*save_interval) == 0:    
                  generator_2.save_weights("images_P/wgen%d.h5" % (epoch))
                  discriminator_2.save_weights("images_P/wdis%d.h5" % (epoch))
                  sample_images(epoch,alpha_val,beta_val,drive = False)
                
            if alpha_val == 0 and beta_val < 1:

              #  Train Discriminator
              # ---------------------

              # select a random half batch of images
              imgs_hr, imgs_lr, imgs_32, imgs_64= data_loader.load_data(half_batch)
              alpha = np.ones(shape = (half_batch,1,1,1))*alpha_val
              beta = np.ones(shape = (half_batch,1,1,1))*beta_val
              # generate a half batch of new images
              gen_imgs = generator_3.predict([imgs_lr,alpha,beta])

              # train the discriminator by feeding both real and fake (generated) images one by one
              d_loss_real = discriminator_3.train_on_batch([imgs_hr,alpha,beta], np.random.uniform(low=0.8, high=1.2, size=(half_batch, 25,25))) # 0.9 for label smoothing
              d_loss_fake = discriminator_3.train_on_batch([gen_imgs,alpha,beta], np.random.uniform(low=0.0, high=0.3, size=(half_batch, 25,25)))
              d_loss = 0.5 * np.add(d_loss_real, d_loss_fake)


              # ---------------------
              #  Train Generator
              # ---------------------


              imgs_hr, imgs_lr, imgs_32, imgs_64 = data_loader.load_data(batch_size)
              alpha = np.ones(shape = (batch_size,1,1,1))*alpha_val
              beta = np.ones(shape = (batch_size,1,1,1))*beta_val
              # the generator wants the discriminator to label the generated samples
              # as valid (ones)

              imgs_h = imgs_hr
              image_features = vgg.predict(imgs_h)


              # train the generator
              g_loss = combined_3.train_on_batch([imgs_lr,alpha,beta], [np.ones((batch_size,25,25)),image_features])
              elapsed_time = datetime.datetime.now() - start_time
              
              # Plot the progress
              print ("%d time: %s ,[D_acc.: %.2f%%] [G loss: %f] [alpha = %.2f beta = %.2f]" % (epoch,elapsed_time, 100*d_loss, g_loss[0],alpha_val,beta_val))
              D_L_REAL.append(d_loss_real)
              D_L_FAKE.append(d_loss_fake)
              D_L.append(d_loss)
              D_ACC.append(d_loss)
              G_L.append(g_loss)

              # if at save interval => save generated image samples
              if epoch % 10 == 0:
                  sample_images(epoch,alpha_val,beta_val)

              if epoch % (10*save_interval) == 0:    
                  generator_3.save_weights("images_P/wgen%d.h5" % (epoch))
                  discriminator_3.save_weights("images_P/wdis%d.h5" % (epoch))
                  sample_images(epoch,alpha_val,beta_val,drive = False)

def sample_images(epoch,alpha_val,beta_val,drive = False):
    os.makedirs('images_P/%s' % dataset_name, exist_ok=True)
    r, c = 2, 3

    imgs_hr, imgs_lr,imgs_32, imgs_64= data_loader.load_data(batch_size=2, is_testing=True)
    alpha = np.ones(shape = (2,1,1,1))*alpha_val
    beta = np.ones(shape = (2,1,1,1))*beta_val

    if alpha_val==1 and beta_val==1:
      imgs_h = (imgs_32)
      fake_hr = generator_1.predict([imgs_lr])
    if alpha_val < 1 and beta_val==1:
      imgs_h = (imgs_64)
      fake_hr = generator_2.predict([imgs_lr,alpha])
    if alpha_val==0 and beta_val<1:
      imgs_h = imgs_hr
      fake_hr = generator_3.predict([imgs_lr,alpha,beta])
    
    


    imgs_hr_inp = imgs_hr
    imgs_lr_inp = imgs_lr
    # Rescale images 0 - 1
    print(np.max(fake_hr))
    print(np.min(fake_hr))
    imgs_lr = 0.5 * imgs_lr + 0.5
    fake_hr= np.clip(0.5 * fake_hr + 0.5,0,1)
    imgs_hr = np.clip(0.5 * imgs_h + 0.5,0,1)
    # Save generated images and the high resolution originals
    titles = ['Original','128x128']
    fig, axs = plt.subplots(r, c)
    cnt = 0
    for row in range(r):
        for col, image in enumerate([imgs_hr,fake_hr]):
            axs[row, col+1].imshow(imgs_lr[row])
            axs[row, col].imshow(fake_hr[row])
            axs[row, col-1].imshow( imgs_hr[row])
            axs[row, col].set_title('Super resolution')
            axs[row, col+1].set_title('input')
            axs[row, col-1].set_title("original")
            axs[row, col-1].axis('off')
            axs[row, col].axis('off')
            axs[row, col+1].axis('off')
        cnt += 1
    txt = ("psnr = %f - ssim = %f \n alpha = %f - beta = %f" % (all_psnr(fake_hr,imgs_hr), all_ssim(fake_hr,imgs_hr),alpha_val,beta_val))
    fig.text(.5, .05, txt, ha='center')
    if drive:
      fig.savefig("/content/drive/My Drive/Keras_models/images_MSG/%s/%d.png" % (dataset_name, epoch))
    else:
      fig.savefig("images_P/%s/%d.png" % (dataset_name, epoch))
    plt.close()


    
    print(all_psnr(fake_hr,imgs_hr))
    print(all_ssim(fake_hr,imgs_hr))

    return imgs_lr_inp,imgs_hr_inp

history = train('cifar10', epochs=30000, batch_size=32, save_interval=50)
#generator.save('mnist_model.h5')
#generator.save('cifar10_model.h5')


plt.plot(D_L)
plt.title('Discriminator results (MNIST)')
plt.xlabel('Epochs')
plt.ylabel('Discriminator Loss (blue), Discriminator Accuracy (orange)')
plt.legend(['Discriminator Loss', 'Discriminator Accuracy'])
plt.show()



plt.plot(G_L)
plt.title('Generator results (MNIST)')
plt.xlabel('Epochs')
plt.ylabel('Generator Loss (blue)')
plt.legend('Generator Loss')
plt.show()


plt.plot(D_L)
plt.title('Discriminator results (CIFAR10)')
plt.xlabel('Epochs')
plt.ylabel('Discriminator Loss (blue), Discriminator Accuracy (orange)')
plt.legend(['Discriminator Loss', 'Discriminator Accuracy'])
plt.show()

plt.plot(G_L)
plt.title('Generator results (CIFAR10)')
plt.xlabel('Epochs')
plt.ylabel('Generator Loss (blue)')
plt.legend('Generator Loss')
plt.show()

from google.colab import files
files.download('model.json')

#!rm -r images_MSG