Delete GAN.py

GAN.py

@@ -1,870 +0,0 @@
-#!/usr/bin/env python
-# coding: utf-8
-
-# In[1]:
-
-
-import numpy as np
-import pandas as pd
-import matplotlib.pyplot as plt
-import seaborn as sns
-import os
-import pickle
-import time
-import random
-
-
-# In[8]:
-
-
-import PIL
-from PIL import Image
-import keras.backend as K
-import tensorflow as tf
-from tensorflow import keras
-from keras.optimizers import Adam
-from keras.models import Sequential
-from keras import layers,Model,Input
-from keras.layers import Lambda,Reshape,UpSampling2D,ReLU,add,ZeroPadding2D
-from keras.layers import Activation,BatchNormalization,Concatenate,concatenate
-from keras.layers import Dense,Conv2D,Flatten,Dropout,LeakyReLU
-from keras.preprocessing.image import ImageDataGenerator
-
-
-# ### Conditioning Augmentation Network
-
-# In[3]:
-
-
-# conditioned by the text.
-def conditioning_augmentation(x):
-    """The mean_logsigma passed as argument is converted into the text conditioning variable.
-
-    Args:
-        x: The output of the text embedding passed through a FC layer with LeakyReLU non-linearity.
-
-    Returns:
-        c: The text conditioning variable after computation.
-    """
-    mean = x[:, :128]
-    log_sigma = x[:, 128:]
-
-    stddev = tf.math.exp(log_sigma)
-    epsilon = K.random_normal(shape=K.constant((mean.shape[1], ), dtype='int32'))
-    c = mean + stddev * epsilon
-    return c
-
-def build_ca_network():
-    """Builds the conditioning augmentation network.
-    """
-    input_layer1 = Input(shape=(1024,)) #size of the vocabulary in the text data
-    mls = Dense(256)(input_layer1)
-    mls = LeakyReLU(alpha=0.2)(mls)
-    ca = Lambda(conditioning_augmentation)(mls)
-    return Model(inputs=[input_layer1], outputs=[ca]) 
-
-
-# ### Stage 1 Generator Network
-
-# In[4]:
-
-
-def UpSamplingBlock(x, num_kernels):
-    """An Upsample block with Upsampling2D, Conv2D, BatchNormalization and a ReLU activation.
-
-    Args:
-        x: The preceding layer as input.
-        num_kernels: Number of kernels for the Conv2D layer.
-
-    Returns:
-        x: The final activation layer after the Upsampling block.
-    """
-    x = UpSampling2D(size=(2,2))(x)
-    x = Conv2D(num_kernels, kernel_size=(3,3), padding='same', strides=1, use_bias=False,
-                kernel_initializer='he_uniform')(x)
-    x = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x) #prevent from mode collapse
-    x = ReLU()(x)
-    return x
-
-
-def build_stage1_generator():
-
-    input_layer1 = Input(shape=(1024,))
-    ca = Dense(256)(input_layer1)
-    ca = LeakyReLU(alpha=0.2)(ca)
-
-    # Obtain the conditioned text
-    c = Lambda(conditioning_augmentation)(ca)
-
-    input_layer2 = Input(shape=(100,))
-    concat = Concatenate(axis=1)([c, input_layer2]) 
-
-    x = Dense(16384, use_bias=False)(concat) 
-    x = ReLU()(x)
-    x = Reshape((4, 4, 1024), input_shape=(16384,))(x)
-
-    x = UpSamplingBlock(x, 512) 
-    x = UpSamplingBlock(x, 256)
-    x = UpSamplingBlock(x, 128)
-    x = UpSamplingBlock(x, 64)   # upsampled our image to 64*64*3 
-
-    x = Conv2D(3, kernel_size=3, padding='same', strides=1, use_bias=False,
-                kernel_initializer='he_uniform')(x)
-    x = Activation('tanh')(x)
-
-    stage1_gen = Model(inputs=[input_layer1, input_layer2], outputs=[x, ca]) 
-    return stage1_gen
-
-
-# In[5]:
-
-
-generator = build_stage1_generator()
-generator.summary()
-
-
-# ### Stage 1 Discriminator Network
-
-# In[9]:
-
-
-def ConvBlock(x, num_kernels, kernel_size=(4,4), strides=2, activation=True):
-    """A ConvBlock with a Conv2D, BatchNormalization and LeakyReLU activation.
-
-    Args:
-        x: The preceding layer as input.
-        num_kernels: Number of kernels for the Conv2D layer.
-
-    Returns:
-        x: The final activation layer after the ConvBlock block.
-    """
-    x = Conv2D(num_kernels, kernel_size=kernel_size, padding='same', strides=strides, use_bias=False,
-                kernel_initializer='he_uniform')(x)
-    x = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x)
-    
-    if activation:
-        x = LeakyReLU(alpha=0.2)(x)
-    return x
-
-
-def build_embedding_compressor():
-    """Build embedding compressor model
-    """
-    input_layer1 = Input(shape=(1024,)) 
-    x = Dense(128)(input_layer1)
-    x = ReLU()(x)
-
-    model = Model(inputs=[input_layer1], outputs=[x])
-    return model
-
-# the discriminator is fed with two inputs, the feature from Generator and the text embedding
-def build_stage1_discriminator():
-    """Builds the Stage 1 Discriminator that uses the 64x64 resolution images from the generator
-    and the compressed and spatially replicated embedding.
-
-    Returns:
-        Stage 1 Discriminator Model for StackGAN.
-    """
-    input_layer1 = Input(shape=(64, 64, 3))  
-
-    x = Conv2D(64, kernel_size=(4,4), strides=2, padding='same', use_bias=False,
-                kernel_initializer='he_uniform')(input_layer1)
-    x = LeakyReLU(alpha=0.2)(x)
-
-    x = ConvBlock(x, 128)
-    x = ConvBlock(x, 256)
-    x = ConvBlock(x, 512)
-
-    # Obtain the compressed and spatially replicated text embedding
-    input_layer2 = Input(shape=(4, 4, 128)) #2nd input to discriminator, text embedding
-    concat = concatenate([x, input_layer2])
-
-    x1 = Conv2D(512, kernel_size=(1,1), padding='same', strides=1, use_bias=False,
-                kernel_initializer='he_uniform')(concat)
-    x1 = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x)
-    x1 = LeakyReLU(alpha=0.2)(x)
-
-    # Flatten and add a FC layer to predict.
-    x1 = Flatten()(x1)
-    x1 = Dense(1)(x1)
-    x1 = Activation('sigmoid')(x1)
-
-    stage1_dis = Model(inputs=[input_layer1, input_layer2], outputs=[x1])  
-    return stage1_dis
-
-
-# In[10]:
-
-
-discriminator = build_stage1_discriminator()
-discriminator.summary()
-
-
-# ### Stage 1 Adversarial Model (Building a GAN)
-
-# In[11]:
-
-
-# Building GAN with Generator and Discriminator
-
-def build_adversarial(generator_model, discriminator_model):
-    """Stage 1 Adversarial model.
-
-    Args:
-        generator_model: Stage 1 Generator Model
-        discriminator_model: Stage 1 Discriminator Model
-
-    Returns:
-        Adversarial Model.
-    """
-    input_layer1 = Input(shape=(1024,))  
-    input_layer2 = Input(shape=(100,)) 
-    input_layer3 = Input(shape=(4, 4, 128)) 
-
-    x, ca = generator_model([input_layer1, input_layer2]) #text,noise
-
-    discriminator_model.trainable = False 
-
-    probabilities = discriminator_model([x, input_layer3]) 
-    adversarial_model = Model(inputs=[input_layer1, input_layer2, input_layer3], outputs=[probabilities, ca])
-    return adversarial_model
-
-
-# In[12]:
-
-
-ganstage1 = build_adversarial(generator, discriminator)
-ganstage1.summary()
-
-
-# ### Train Utilities
-
-# In[13]:
-
-
-def checkpoint_prefix():
-    checkpoint_dir = './training_checkpoints'
-    checkpoint_prefix = os.path.join(checkpoint_dir, 'ckpt')
-
-    return checkpoint_prefix
-
-def adversarial_loss(y_true, y_pred):
-    mean = y_pred[:, :128]
-    ls = y_pred[:, 128:]
-    loss = -ls + 0.5 * (-1 + tf.math.exp(2.0 * ls) + tf.math.square(mean))
-    loss = K.mean(loss)
-    return loss
-
-def normalize(input_image, real_image):
-    input_image = (input_image / 127.5) - 1
-    real_image = (real_image / 127.5) - 1
-
-    return input_image, real_image
-
-def load_class_ids_filenames(class_id_path, filename_path):
-    with open(class_id_path, 'rb') as file:
-        class_id = pickle.load(file, encoding='latin1')
-
-    with open(filename_path, 'rb') as file:
-        filename = pickle.load(file, encoding='latin1')
-
-    return class_id, filename
-
-def load_text_embeddings(text_embeddings):
-    with open(text_embeddings, 'rb') as file:
-        embeds = pickle.load(file, encoding='latin1')
-        embeds = np.array(embeds)
-
-    return embeds
-
-def load_bbox(data_path):
-    bbox_path = data_path + '/bounding_boxes.txt'
-    image_path = data_path + '/images.txt'
-    bbox_df = pd.read_csv(bbox_path, delim_whitespace=True, header=None).astype(int)
-    filename_df = pd.read_csv(image_path, delim_whitespace=True, header=None)
-
-    filenames = filename_df[1].tolist()
-    bbox_dict = {i[:-4]:[] for i in filenames[:2]}
-
-    for i in range(0, len(filenames)):
-        bbox = bbox_df.iloc[i][1:].tolist()
-        dict_key = filenames[i][:-4]
-        bbox_dict[dict_key] = bbox
-
-    return bbox_dict
-
-def load_images(image_path, bounding_box, size):
-    """Crops the image to the bounding box and then resizes it.
-    """
-    image = Image.open(image_path).convert('RGB')
-    w, h = image.size
-    if bounding_box is not None:
-        r = int(np.maximum(bounding_box[2], bounding_box[3]) * 0.75)
-        c_x = int((bounding_box[0] + bounding_box[2]) / 2)
-        c_y = int((bounding_box[1] + bounding_box[3]) / 2)
-        y1 = np.maximum(0, c_y - r)
-        y2 = np.minimum(h, c_y + r)
-        x1 = np.maximum(0, c_x - r)
-        x2 = np.minimum(w, c_x + r)
-        image = image.crop([x1, y1, x2, y2])
-
-    image = image.resize(size, PIL.Image.BILINEAR)
-    return image
-
-def load_data(filename_path, class_id_path, dataset_path, embeddings_path, size):
-    """Loads the Dataset.
-    """
-    data_dir = "D:/1-pipelined_topics/GAN_texttoimage/birds_implementation/birds"
-    train_dir = data_dir + "/train"
-    test_dir = data_dir + "/test"
-    embeddings_path_train = train_dir + "/char-CNN-RNN-embeddings.pickle"
-    embeddings_path_test = test_dir + "/char-CNN-RNN-embeddings.pickle"
-    filename_path_train = train_dir + "/filenames.pickle"
-    filename_path_test = test_dir + "/filenames.pickle"
-    class_id_path_train = train_dir + "/class_info.pickle"
-    class_id_path_test = test_dir + "/class_info.pickle"
-    dataset_path = "D:/1-pipelined_topics/GAN_texttoimage/birds_implementation/CUB_200_2011"
-    class_id, filenames = load_class_ids_filenames(class_id_path, filename_path)
-    embeddings = load_text_embeddings(embeddings_path)
-    bbox_dict = load_bbox(dataset_path)
-
-    x, y, embeds = [], [], []
-
-    for i, filename in enumerate(filenames):
-        bbox = bbox_dict[filename]
-
-        try:
-            image_path = f'{dataset_path}/images/{filename}.jpg'
-            image = load_images(image_path, bbox, size)
-            e = embeddings[i, :, :]
-            embed_index = np.random.randint(0, e.shape[0] - 1)
-            embed = e[embed_index, :]
-
-            x.append(np.array(image))
-            y.append(class_id[i])
-            embeds.append(embed)
-
-        except Exception as e:
-            print(f'{e}')
-    
-    x = np.array(x)
-    y = np.array(y)
-    embeds = np.array(embeds)
-    
-    return x, y, embeds
-
-def save_image(file, save_path):
-    """Saves the image at the specified file path.
-    """
-    image = plt.figure()
-    ax = image.add_subplot(1,1,1)
-    ax.imshow(file)
-    ax.axis("off")
-    plt.savefig(save_path)
-
-
-# In[28]:
-
-
-############################################################
-# StackGAN class
-############################################################
-
-class StackGanStage1(object):
-    """StackGAN Stage 1 class."""
-
-    data_dir = "D:/1-pipelined_topics/GAN_texttoimage/birds_implementation/birds"
-    train_dir = data_dir + "/train"
-    test_dir = data_dir + "/test"
-    embeddings_path_train = train_dir + "/char-CNN-RNN-embeddings.pickle"
-    embeddings_path_test = test_dir + "/char-CNN-RNN-embeddings.pickle"
-    filename_path_train = train_dir + "/filenames.pickle"
-    filename_path_test = test_dir + "/filenames.pickle"
-    class_id_path_train = train_dir + "/class_info.pickle"
-    class_id_path_test = test_dir + "/class_info.pickle"
-    dataset_path = "D:/1-pipelined_topics/GAN_texttoimage/birds_implementation/CUB_200_2011"
-    def __init__(self, epochs=500, z_dim=100, batch_size=64, enable_function=True, stage1_generator_lr=0.0002, stage1_discriminator_lr=0.0002):
-        self.epochs = epochs
-        self.z_dim = z_dim
-        self.enable_function = enable_function
-        self.stage1_generator_lr = stage1_generator_lr
-        self.stage1_discriminator_lr = stage1_discriminator_lr
-        self.image_size = 64
-        self.conditioning_dim = 128
-        self.batch_size = batch_size
-        
-        self.stage1_generator_optimizer = Adam(lr=stage1_generator_lr, beta_1=0.5, beta_2=0.999)
-        self.stage1_discriminator_optimizer = Adam(lr=stage1_discriminator_lr, beta_1=0.5, beta_2=0.999)
-        
-        self.stage1_generator = build_stage1_generator()
-        self.stage1_generator.compile(loss='mse', optimizer=self.stage1_generator_optimizer)
-
-        self.stage1_discriminator = build_stage1_discriminator()
-        self.stage1_discriminator.compile(loss='binary_crossentropy', optimizer=self.stage1_discriminator_optimizer)
-
-        self.ca_network = build_ca_network()
-        self.ca_network.compile(loss='binary_crossentropy', optimizer='Adam')
-
-        self.embedding_compressor = build_embedding_compressor()
-        self.embedding_compressor.compile(loss='binary_crossentropy', optimizer='Adam')
-
-        self.stage1_adversarial = build_adversarial(self.stage1_generator, self.stage1_discriminator)
-        self.stage1_adversarial.compile(loss=['binary_crossentropy', adversarial_loss], loss_weights=[1, 2.0], optimizer=self.stage1_generator_optimizer)
-
-        self.checkpoint1 = tf.train.Checkpoint(
-            generator_optimizer=self.stage1_generator_optimizer,
-            discriminator_optimizer=self.stage1_discriminator_optimizer,
-            generator=self.stage1_generator,
-            discriminator=self.stage1_discriminator)
-
-    def visualize_stage1(self):
-        """Running Tensorboard visualizations.
-        """
-        tb = TensorBoard(log_dir="logs/".format(time.time()))
-        tb.set_model(self.stage1_generator)
-        tb.set_model(self.stage1_discriminator)
-        tb.set_model(self.ca_network)
-        tb.set_model(self.embedding_compressor)
-
-    def train_stage1(self):
-        """Trains the stage1 StackGAN.
-        """
-        x_train, y_train, train_embeds = load_data(filename_path=filename_path_train, class_id_path=class_id_path_train,
-                                        dataset_path=dataset_path, embeddings_path=embeddings_path_train, size=(64, 64))
-
-        x_test, y_test, test_embeds = load_data(filename_path=filename_path_test, class_id_path=class_id_path_test, 
-      dataset_path=dataset_path, embeddings_path=embeddings_path_test, size=(64, 64))
-
-        real = np.ones((self.batch_size, 1), dtype='float') * 0.9
-        fake = np.zeros((self.batch_size, 1), dtype='float') * 0.1
-
-        for epoch in range(self.epochs):
-            print(f'Epoch: {epoch}')
-
-            gen_loss = []
-            dis_loss = []
-
-            num_batches = int(x_train.shape[0] / self.batch_size)
-
-            for i in range(num_batches):
-
-                latent_space = np.random.normal(0, 1, size=(self.batch_size, self.z_dim))
-                embedding_text = train_embeds[i * self.batch_size:(i + 1) * self.batch_size]
-                compressed_embedding = self.embedding_compressor.predict_on_batch(embedding_text)
-                compressed_embedding = np.reshape(compressed_embedding, (-1, 1, 1, 128))
-                compressed_embedding = np.tile(compressed_embedding, (1, 4, 4, 1))
-
-            image_batch = x_train[i * self.batch_size:(i+1) * self.batch_size]
-            image_batch = (image_batch - 127.5) / 127.5
-
-            gen_images, _ = self.stage1_generator.predict([embedding_text, latent_space])
-
-            discriminator_loss = self.stage1_discriminator.train_on_batch([image_batch, compressed_embedding], 
-                    np.reshape(real, (self.batch_size, 1)))
-
-            discriminator_loss_gen = self.stage1_discriminator.train_on_batch([gen_images, compressed_embedding],
-                    np.reshape(fake, (self.batch_size, 1)))
-
-            discriminator_loss_wrong = self.stage1_discriminator.train_on_batch([gen_images[: self.batch_size-1], compressed_embedding[1:]], 
-                    np.reshape(fake[1:], (self.batch_size-1, 1)))
-
-# Discriminator loss
-            d_loss = 0.5 * np.add(discriminator_loss, 0.5 * np.add(discriminator_loss_gen, discriminator_loss_wrong))
-            dis_loss.append(d_loss)
-
-            print(f'Discriminator Loss: {d_loss}')
-
-               # Generator loss
-            g_loss = self.stage1_adversarial.train_on_batch([embedding_text, latent_space, compressed_embedding],
-                    [K.ones((self.batch_size, 1)) * 0.9, K.ones((self.batch_size, 256)) * 0.9])
-
-            print(f'Generator Loss: {g_loss}')
-            gen_loss.append(g_loss)
-
-            if epoch % 5 == 0:
-                    latent_space = np.random.normal(0, 1, size=(self.batch_size, self.z_dim))
-                    embedding_batch = test_embeds[0 : self.batch_size]
-                    gen_images, _ = self.stage1_generator.predict_on_batch([embedding_batch, latent_space])
-
-                    for i, image in enumerate(gen_images[:10]):
-                        save_image(image, f'test/gen_1_{epoch}_{i}')
-
-            if epoch % 25 == 0:
-                self.stage1_generator.save_weights('weights/stage1_gen.h5')
-                self.stage1_discriminator.save_weights("weights/stage1_disc.h5")
-                self.ca_network.save_weights('weights/stage1_ca.h5')
-                self.embedding_compressor.save_weights('weights/stage1_embco.h5')
-                self.stage1_adversarial.save_weights('weights/stage1_adv.h5')      
-
-        self.stage1_generator.save_weights('weights/stage1_gen.h5')
-        self.stage1_discriminator.save_weights("weights/stage1_disc.h5")
-
-
-# In[ ]:
-
-
-stage1 = StackGanStage1()
-stage1.train_stage1()
-
-
-# ### Check test folder for gernerated images from Stage1 Generator
-# ### Let's Implement Stage 2 Generator
-
-# In[29]:
-
-
-############################################################
-# Stage 2 Generator Network
-############################################################
-
-def concat_along_dims(inputs):
-    """Joins the conditioned text with the encoded image along the dimensions.
-
-    Args:
-        inputs: consisting of conditioned text and encoded images as [c,x].
-
-    Returns:
-        Joint block along the dimensions.
-    """
-    c = inputs[0]
-    x = inputs[1]
-
-    c = K.expand_dims(c, axis=1)
-    c = K.expand_dims(c, axis=1)
-    c = K.tile(c, [1, 16, 16, 1])
-    return K.concatenate([c, x], axis = 3)
-
-def residual_block(input):
-    """Residual block with plain identity connections.
-
-    Args:
-        inputs: input layer or an encoded layer
-
-    Returns:
-        Layer with computed identity mapping.
-    """
-    x = Conv2D(512, kernel_size=(3,3), padding='same', use_bias=False,
-                kernel_initializer='he_uniform')(input)
-    x = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x)
-    x = ReLU()(x)
-    
-    x = Conv2D(512, kernel_size=(3,3), padding='same', use_bias=False,
-                kernel_initializer='he_uniform')(x)
-    x = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x)
-    
-    x = add([x, input])
-    x = ReLU()(x)
-
-    return x
-
-def build_stage2_generator():
-    """Build the Stage 2 Generator Network using the conditioning text and images from stage 1.
-
-    Returns:
-        Stage 2 Generator Model for StackGAN.
-    """
-    input_layer1 = Input(shape=(1024,))
-    input_images = Input(shape=(64, 64, 3))
-
-    # Conditioning Augmentation
-    ca = Dense(256)(input_layer1)
-    mls = LeakyReLU(alpha=0.2)(ca)
-    c = Lambda(conditioning_augmentation)(mls)
-
-    # Downsampling block
-    x = ZeroPadding2D(padding=(1,1))(input_images)
-    x = Conv2D(128, kernel_size=(3,3), strides=1, use_bias=False,
-                kernel_initializer='he_uniform')(x)
-    x = ReLU()(x)
-
-    x = ZeroPadding2D(padding=(1,1))(x)
-    x = Conv2D(256, kernel_size=(4,4), strides=2, use_bias=False,
-                kernel_initializer='he_uniform')(x)
-    x = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x)
-    x = ReLU()(x)
-
-    x = ZeroPadding2D(padding=(1,1))(x)
-    x = Conv2D(512, kernel_size=(4,4), strides=2, use_bias=False,
-                kernel_initializer='he_uniform')(x)
-    x = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x)
-    x = ReLU()(x)
-
-    # Concatenate text conditioning block with the encoded image
-    concat = concat_along_dims([c, x])
-
-    # Residual Blocks
-    x = ZeroPadding2D(padding=(1,1))(concat)
-    x = Conv2D(512, kernel_size=(3,3), use_bias=False, kernel_initializer='he_uniform')(x)
-    x = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x)
-    x = ReLU()(x)
-
-    x = residual_block(x)
-    x = residual_block(x)
-    x = residual_block(x)
-    x = residual_block(x)
-    
-    # Upsampling Blocks
-    x = UpSamplingBlock(x, 512)
-    x = UpSamplingBlock(x, 256)
-    x = UpSamplingBlock(x, 128)
-    x = UpSamplingBlock(x, 64)
-
-    x = Conv2D(3, kernel_size=(3,3), padding='same', use_bias=False, kernel_initializer='he_uniform')(x)
-    x = Activation('tanh')(x)
-    
-    stage2_gen = Model(inputs=[input_layer1, input_images], outputs=[x, mls])
-    return stage2_gen
-
-
-# In[30]:
-
-
-generator_stage2 = build_stage2_generator()
-generator_stage2.summary()
-
-
-# In[31]:
-
-
-############################################################
-# Stage 2 Discriminator Network
-############################################################
-
-def build_stage2_discriminator():
-    """Builds the Stage 2 Discriminator that uses the 256x256 resolution images from the generator
-    and the compressed and spatially replicated embeddings.
-
-    Returns:
-        Stage 2 Discriminator Model for StackGAN.
-    """
-    input_layer1 = Input(shape=(256, 256, 3))
-
-    x = Conv2D(64, kernel_size=(4,4), padding='same', strides=2, use_bias=False,
-                kernel_initializer='he_uniform')(input_layer1)
-    x = LeakyReLU(alpha=0.2)(x)
-
-    x = ConvBlock(x, 128)
-    x = ConvBlock(x, 256)
-    x = ConvBlock(x, 512)
-    x = ConvBlock(x, 1024)
-    x = ConvBlock(x, 2048)
-    x = ConvBlock(x, 1024, (1,1), 1)
-    x = ConvBlock(x, 512, (1,1), 1, False)
-
-    x1 = ConvBlock(x, 128, (1,1), 1)
-    x1 = ConvBlock(x1, 128, (3,3), 1)
-    x1 = ConvBlock(x1, 512, (3,3), 1, False)
-
-    x2 = add([x, x1])
-    x2 = LeakyReLU(alpha=0.2)(x2)
-
-    # Concatenate compressed and spatially replicated embedding
-    input_layer2 = Input(shape=(4, 4, 128))
-    concat = concatenate([x2, input_layer2])
-
-    x3 = Conv2D(512, kernel_size=(1,1), strides=1, padding='same', kernel_initializer='he_uniform')(concat)
-    x3 = BatchNormalization(gamma_initializer='ones', beta_initializer='zeros')(x3)
-    x3 = LeakyReLU(alpha=0.2)(x3)
-
-    # Flatten and add a FC layer
-    x3 = Flatten()(x3)
-    x3 = Dense(1)(x3)
-    x3 = Activation('sigmoid')(x3)
-
-    stage2_dis = Model(inputs=[input_layer1, input_layer2], outputs=[x3])
-    return stage2_dis
-
-
-# In[32]:
-
-
-discriminator_stage2 = build_stage2_discriminator()
-discriminator_stage2.summary()
-
-
-# In[33]:
-
-
-############################################################
-# Stage 2 Adversarial Model
-############################################################
-
-def stage2_adversarial_network(stage2_disc, stage2_gen, stage1_gen):
-    """Stage 2 Adversarial Network.
-
-    Args:
-        stage2_disc: Stage 2 Discriminator Model.
-        stage2_gen: Stage 2 Generator Model.
-        stage1_gen: Stage 1 Generator Model.
-
-    Returns:
-        Stage 2 Adversarial network.
-    """
-    conditioned_embedding = Input(shape=(1024, ))
-    latent_space = Input(shape=(100, ))
-    compressed_replicated = Input(shape=(4, 4, 128))
-    
-    #the discriminator is trained separately and stage1_gen already trained, and this is the reason why we freeze its layers by setting the property trainable=false
-    input_images, ca = stage1_gen([conditioned_embedding, latent_space])
-    stage2_disc.trainable = False
-    stage1_gen.trainable = False
-
-    images, ca2 = stage2_gen([conditioned_embedding, input_images])
-    probability = stage2_disc([images, compressed_replicated])
-
-    return Model(inputs=[conditioned_embedding, latent_space, compressed_replicated],
-        outputs=[probability, ca2])
-
-
-# In[34]:
-
-
-adversarial_stage2 = stage2_adversarial_network(discriminator_stage2, generator_stage2, generator)
-adversarial_stage2.summary()
-
-
-# In[35]:
-
-
-class StackGanStage2(object):
-    """StackGAN Stage 2 class.
-
-    Args:
-        epochs: Number of epochs
-        z_dim: Latent space dimensions
-        batch_size: Batch Size
-        enable_function: If True, training function is decorated with tf.function
-        stage2_generator_lr: Learning rate for stage 2 generator
-        stage2_discriminator_lr: Learning rate for stage 2 discriminator
-    """
-    def __init__(self, epochs=500, z_dim=100, batch_size=64, enable_function=True, stage2_generator_lr=0.0002, stage2_discriminator_lr=0.0002):
-        self.epochs = epochs
-        self.z_dim = z_dim
-        self.enable_function = enable_function
-        self.stage1_generator_lr = stage2_generator_lr
-        self.stage1_discriminator_lr = stage2_discriminator_lr
-        self.low_image_size = 64
-        self.high_image_size = 256
-        self.conditioning_dim = 128
-        self.batch_size = batch_size
-        self.stage2_generator_optimizer = Adam(lr=stage2_generator_lr, beta_1=0.5, beta_2=0.999)
-        self.stage2_discriminator_optimizer = Adam(lr=stage2_discriminator_lr, beta_1=0.5, beta_2=0.999)
-        self.stage1_generator = build_stage1_generator()
-        self.stage1_generator.compile(loss='binary_crossentropy', optimizer=self.stage2_generator_optimizer)
-        self.stage1_generator.load_weights('weights/stage1_gen.h5')
-        self.stage2_generator = build_stage2_generator()
-        self.stage2_generator.compile(loss='binary_crossentropy', optimizer=self.stage2_generator_optimizer)
-
-        self.stage2_discriminator = build_stage2_discriminator()
-        self.stage2_discriminator.compile(loss='binary_crossentropy', optimizer=self.stage2_discriminator_optimizer)
-
-        self.ca_network = build_ca_network()
-        self.ca_network.compile(loss='binary_crossentropy', optimizer='Adam')
-
-        self.embedding_compressor = build_embedding_compressor()
-        self.embedding_compressor.compile(loss='binary_crossentropy', optimizer='Adam')
-
-        self.stage2_adversarial = stage2_adversarial_network(self.stage2_discriminator, self.stage2_generator, self.stage1_generator)
-        self.stage2_adversarial.compile(loss=['binary_crossentropy', adversarial_loss], loss_weights=[1, 2.0], optimizer=self.stage2_generator_optimizer)	
-
-        self.checkpoint2 = tf.train.Checkpoint(
-            generator_optimizer=self.stage2_generator_optimizer,
-            discriminator_optimizer=self.stage2_discriminator_optimizer,
-            generator=self.stage2_generator,
-            discriminator=self.stage2_discriminator,
-            generator1=self.stage1_generator)
-
-    def visualize_stage2(self):
-        """Running Tensorboard visualizations.
-        """
-        tb = TensorBoard(log_dir="logs/".format(time.time()))
-        tb.set_model(self.stage2_generator)
-        tb.set_model(self.stage2_discriminator)
-
-    def train_stage2(self):
-        """Trains Stage 2 StackGAN.
-        """
-        x_high_train, y_high_train, high_train_embeds = load_data(filename_path=filename_path_train, class_id_path=class_id_path_train,
-            dataset_path=dataset_path, embeddings_path=embeddings_path_train, size=(256, 256))
-
-        x_high_test, y_high_test, high_test_embeds = load_data(filename_path=filename_path_test, class_id_path=class_id_path_test, 
-              dataset_path=dataset_path, embeddings_path=embeddings_path_test, size=(256, 256))
-
-        x_low_train, y_low_train, low_train_embeds = load_data(filename_path=filename_path_train, class_id_path=class_id_path_train,
-          dataset_path=dataset_path, embeddings_path=embeddings_path_train, size=(64, 64))
-
-        x_low_test, y_low_test, low_test_embeds = load_data(filename_path=filename_path_test, class_id_path=class_id_path_test, 
-          dataset_path=dataset_path, embeddings_path=embeddings_path_test, size=(64, 64))
-
-        real = np.ones((self.batch_size, 1), dtype='float') * 0.9
-        fake = np.zeros((self.batch_size, 1), dtype='float') * 0.1
-
-        for epoch in range(self.epochs):
-            print(f'Epoch: {epoch}')
-
-            gen_loss = []
-            disc_loss = []
-
-            num_batches = int(x_high_train.shape[0] / self.batch_size)
-
-            for i in range(num_batches):
-
-                latent_space = np.random.normal(0, 1, size=(self.batch_size, self.z_dim))
-                embedding_text = high_train_embeds[i * self.batch_size:(i + 1) * self.batch_size]
-                compressed_embedding = self.embedding_compressor.predict_on_batch(embedding_text)
-                compressed_embedding = np.reshape(compressed_embedding, (-1, 1, 1, self.conditioning_dim))
-                compressed_embedding = np.tile(compressed_embedding, (1, 4, 4, 1))
-
-                image_batch = x_high_train[i * self.batch_size:(i+1) * self.batch_size]
-                image_batch = (image_batch - 127.5) / 127.5
-                
-                low_res_fakes, _ = self.stage1_generator.predict([embedding_text, latent_space], verbose=3)
-                high_res_fakes, _ = self.stage2_generator.predict([embedding_text, low_res_fakes], verbose=3)
-
-                discriminator_loss = self.stage2_discriminator.train_on_batch([image_batch, compressed_embedding],
-                    np.reshape(real, (self.batch_size, 1)))
-
-                discriminator_loss_gen = self.stage2_discriminator.train_on_batch([high_res_fakes, compressed_embedding],
-                    np.reshape(fake, (self.batch_size, 1)))
-
-                discriminator_loss_fake = self.stage2_discriminator.train_on_batch([image_batch[:(self.batch_size-1)], compressed_embedding[1:]],
-                    np.reshape(fake[1:], (self.batch_size - 1, 1)))
-
-                d_loss = 0.5 * np.add(discriminator_loss, 0.5 * np.add(discriminator_loss_gen, discriminator_loss_fake))
-                disc_loss.append(d_loss)
-
-                print(f'Discriminator Loss: {d_loss}')
-
-                g_loss = self.stage2_adversarial.train_on_batch([embedding_text, latent_space, compressed_embedding],
-                    [K.ones((self.batch_size, 1)) * 0.9, K.ones((self.batch_size, 256)) * 0.9])
-                gen_loss.append(g_loss)
-
-                print(f'Generator Loss: {g_loss}')
-
-                if epoch % 5 == 0:
-                    latent_space = np.random.normal(0, 1, size=(self.batch_size, self.z_dim))
-                    embedding_batch = high_test_embeds[0 : self.batch_size]
-
-                    low_fake_images, _ = self.stage1_generator.predict([embedding_batch, latent_space], verbose=3)
-                    high_fake_images, _ = self.stage2_generator.predict([embedding_batch, low_fake_images], verbose=3)
-
-                    for i, image in enumerate(high_fake_images[:10]):
-                        save_image(image, f'results_stage2/gen_{epoch}_{i}.png')
-
-                if epoch % 10 == 0:
-                    self.stage2_generator.save_weights('weights/stage2_gen.h5')
-                    self.stage2_discriminator.save_weights("weights/stage2_disc.h5")
-                    self.ca_network.save_weights('weights/stage2_ca.h5')
-                    self.embedding_compressor.save_weights('weights/stage2_embco.h5')
-                    self.stage2_adversarial.save_weights('weights/stage2_adv.h5')
-
-        self.stage2_generator.save_weights('weights/stage2_gen.h5')
-        self.stage2_discriminator.save_weights("weights/stage2_disc.h5")
-
-
-# In[ ]:
-
-
-stage2 = StackGanStage2()
-stage2.train_stage2()
-