GAN_model / app.py
Truptidand's picture
Update app.py
f57813b
#!/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
from tensorflow import keras
import keras.backend as K
import tensorflow as tf
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()