Spaces:
Runtime error
Runtime error
#!/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() | |