from __future__ import print_function
import os
import glob
import scipy
import tensorflow as tf
import numpy as np
from PIL import Image
import skimage.io as io
import matplotlib.pyplot as plt
class Arguments(object):
data_path = 'results_celebA/preprocessed/' #path to CelebA dataset
save_path = 'results_celebA' #path to save preprocessed image folder
preproc_foldername = 'preprocessed' #folder name for preprocessed images
image_size = 64 #images are resized to image_size value
num_images = 202590 #the number of training images
batch_size = 64 #batch size
dim_z = 100 #the dimension of z variable (the generator input dimension)
n_g_filters = 64 #the number of the generator filters (gets multiplied between layers)
n_f_filters = 64 #the number of the discriminator filters (gets multiplied between layers)
n_epoch = 25 #the number of epochs
lr = 0.0002 #learning rate
beta1 = 0.5 #beta_1 parameter of Adam optimizer
beta2 = 0.99 #beta_2 parameter of Adam optimizer
args = Arguments()
class Dataset(object):
def __init__(self, data_path, num_imgs, target_imgsize):
self.data_path = data_path
self.num_imgs = num_imgs
self.target_imgsize = target_imgsize
def normalize_np_image(self, image):
return (image / 255.0 - 0.5) / 0.5
def denormalize_np_image(self, image):
return (image * 0.5 + 0.5) * 255
def get_input(self, image_path):
image = np.array(Image.open(image_path)).astype(np.float32)
return self.normalize_np_image(image)
def get_imagelist(self, data_path, celebA=False):
if celebA == True:
imgs_path = os.path.join(data_path, 'img_align_celeba/*.jpg')
else:
imgs_path = os.path.join(data_path, '*.jpg')
all_namelist = glob.glob(imgs_path, recursive=True)
return all_namelist[:self.num_imgs]
def load_and_preprocess_image(self, image_path):
image = Image.open(image_path)
j = (image.size[0] - 100) // 2
i = (image.size[1] - 100) // 2
image = image.crop([j, i, j + 100, i + 100])
image = image.resize([self.target_imgsize, self.target_imgsize], Image.BILINEAR)
image = np.array(image.convert('RGB')).astype(np.float32)
image = self.normalize_np_image(image)
return image
#reads data, preprocesses and saves to another folder with the given path.
def preprocess_and_save_images(self, dir_name, save_path=''):
preproc_folder_path = os.path.join(save_path, dir_name)
if not os.path.exists(preproc_folder_path):
os.makedirs(preproc_folder_path)
imgs_path = os.path.join(self.data_path, 'img_align_celeba/*.jpg')
print('Saving and preprocessing images ...')
for num, imgname in enumerate(glob.iglob(imgs_path, recursive=True)):
cur_image = self.load_and_preprocess_image(imgname)
cur_image = Image.fromarray(np.uint8(self.denormalize_np_image(cur_image)))
cur_image.save(preproc_folder_path + '/preprocessed_image_%d.jpg' %(num))
self.data_path= preproc_folder_path
def get_nextbatch(self, batch_size):
assert (batch_size > 0),"Give a valid batch size"
cur_idx = 0
image_namelist = self.get_imagelist(self.data_path)
while cur_idx + batch_size <= self.num_imgs:
cur_namelist = image_namelist[cur_idx:cur_idx + batch_size]
cur_batch = [self.get_input(image_path) for image_path in cur_namelist]
cur_batch = np.array(cur_batch).astype(np.float32)
cur_idx += batch_size
yield cur_batch
def show_image(self, image, normalized=True):
if not type(image).__module__ == np.__name__:
image = image.numpy()
if normalized:
npimg = (image * 0.5) + 0.5
npimg.astype(np.uint8)
plt.imshow(npimg, interpolation='nearest')
def generator(x, args, reuse=False):
with tf.device('/gpu:0'):
with tf.variable_scope("generator", reuse=reuse):
#Layer Block 1
with tf.variable_scope("layer1"):
deconv1 = tf.layers.conv2d_transpose(inputs=x,
filters= args.n_g_filters*8,
kernel_size=4,
strides=1,
padding='valid',
use_bias=False,
name='deconv')
batch_norm1=tf.layers.batch_normalization(deconv1,
name = 'batch_norm')
relu1 = tf.nn.relu(batch_norm1, name='relu')
#Layer Block 2
with tf.variable_scope("layer2"):
deconv2 = tf.layers.conv2d_transpose(inputs=relu1,
filters=args.n_g_filters*4,
kernel_size=4,
strides=2,
padding='same',
use_bias=False,
name='deconv')
batch_norm2 = tf.layers.batch_normalization(deconv2,
name = 'batch_norm')
relu2 = tf.nn.relu(batch_norm2, name='relu')
#Layer Block 3
with tf.variable_scope("layer3"):
deconv3 = tf.layers.conv2d_transpose(inputs=relu2,
filters=args.n_g_filters*2,
kernel_size=4,
strides=2,
padding='same',
use_bias = False,
name='deconv')
batch_norm3 = tf.layers.batch_normalization(deconv3,
name = 'batch_norm')
relu3 = tf.nn.relu(batch_norm3, name='relu')
#Layer Block 4
with tf.variable_scope("layer4"):
deconv4 = tf.layers.conv2d_transpose(inputs=relu3,
filters=args.n_g_filters,
kernel_size=4,
strides=2,
padding='same',
use_bias=False,
name='deconv')
batch_norm4 = tf.layers.batch_normalization(deconv4,
name = 'batch_norm')
relu4 = tf.nn.relu(batch_norm4, name='relu')
#Output Layer
with tf.variable_scope("last_layer"):
logit = tf.layers.conv2d_transpose(inputs=relu4,
filters=3,
kernel_size=4,
strides=2,
padding='same',
use_bias=False,
name='logit')
output = tf.nn.tanh(logit)
return output, logit
def discriminator(x, args, reuse=False):
with tf.device('/gpu:0'):
with tf.variable_scope("discriminator", reuse=reuse):
with tf.variable_scope("layer1"):
conv1 = tf.layers.conv2d(inputs=x,
filters=args.n_f_filters,
kernel_size=4,
strides=2,
padding='same',
use_bias=False,
name='conv')
relu1 = tf.nn.leaky_relu(conv1, alpha=0.2, name='relu')
with tf.variable_scope("layer2"):
conv2 = tf.layers.conv2d(inputs=relu1,
filters=args.n_f_filters*2,
kernel_size=4,
strides=2,
padding='same',
use_bias=False,
name='conv')
batch_norm2 = tf.layers.batch_normalization(conv2,name='batch_norm')
relu2 = tf.nn.leaky_relu(batch_norm2, alpha=0.2, name='relu')
with tf.variable_scope("layer3"):
conv3 = tf.layers.conv2d(inputs=relu2,
filters=args.n_f_filters*4,
kernel_size=4,
strides=2,
padding='same',
use_bias=False,
name='conv')
batch_norm3 = tf.layers.batch_normalization(conv3, name='batch_norm')
relu3 = tf.nn.leaky_relu(batch_norm3, name='relu')
with tf.variable_scope("layer4"):
conv4 = tf.layers.conv2d(inputs=relu3,
filters=args.n_f_filters*8,
kernel_size=4,
strides=2,
padding='same',
use_bias=False,
name='conv')
batch_norm4 = tf.layers.batch_normalization(conv4, name='batch_norm')
relu4 = tf.nn.leaky_relu(batch_norm4, alpha=0.2, name='relu')
with tf.variable_scope("last_layer"):
logit = tf.layers.conv2d(inputs=relu4,
filters=1,
kernel_size=4,
strides=1,
padding='valid',
use_bias=False,
name='conv')
output = tf.nn.sigmoid(logit)
return output, logit
def sample_z(dim_z, num_batch):
mu = 0
sigma = 1
s = np.random.normal(mu, sigma, num_batch*dim_z)
samples = s.reshape(num_batch, 1, 1, dim_z)
return samples
def generator_loss(fake):
loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.ones_like(fake), logits=fake))
return loss
def discriminator_loss(real, fake):
real_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.ones_like(real), logits=real))
fake_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.zeros_like(fake), logits=fake))
loss = real_loss + fake_loss
return loss
def get_losses(d_real_logits, d_fake_logits):
# d_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_real_logits,labels=tf.ones_like(d_real_logits)) + tf.nn.sigmoid_cross_entropy_with_logits(logits=d_fake_logits,labels=tf.zeros_like(d_fake_logits)))
# g_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_fake_logits,labels=tf.ones_like(d_fake_logits)))
d_loss = discriminator_loss(real=d_real_logits, fake=d_fake_logits)
g_loss = generator_loss(fake=d_fake_logits)
return d_loss, g_loss
def get_optimizers(learning_rate, beta1, beta2):
d_optimizer = tf.train.AdamOptimizer(learning_rate, beta1, beta2)
g_optimizer = tf.train.AdamOptimizer(learning_rate, beta1, beta2)
return d_optimizer, g_optimizer
def optimize(d_optimizer, g_optimizer, d_loss, g_loss):
d_step = d_optimizer.minimize(d_loss)
g_step = g_optimizer.minimize(g_loss)
return d_step, g_step
LOGDIR = "logs_basic_dcgan"
def merge_images(image_batch, size):
h,w = image_batch.shape[1], image_batch.shape[2]
c = image_batch.shape[3]
img = np.zeros((int(h*size[0]), w*size[1], c))
for idx, im in enumerate(image_batch):
i = idx % size[1]
j = idx // size[1]
img[j*h:j*h+h, i*w:i*w+w,:] = im
return img
itr_fh = open('basic_gan_itr.txt', 'a+')
def train(args):
tf.reset_default_graph()
data_loader = Dataset(args.data_path, args.num_images, args.image_size)
#data_loader.preprocess_and_save_images('preprocessed', 'results_celebA') #preprocess the images once
X = tf.placeholder(tf.float32, shape=[args.batch_size, args.image_size , args.image_size, 3])
Z = tf.placeholder(tf.float32, shape=[args.batch_size, 1, 1, args.dim_z])
G_sample, _ = generator(Z, args)
D_real, D_real_logits = discriminator(X, args)
D_fake, D_fake_logits = discriminator(G_sample, args, reuse=True)
d_loss, g_loss = get_losses(D_real_logits, D_fake_logits)
d_optimizer, g_optimizer = get_optimizers(args.lr, args.beta1, args.beta2)
d_step, g_step = optimize(d_optimizer, g_optimizer, d_loss, g_loss)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for epoch in range(args.n_epoch):
for itr, real_batch in enumerate(data_loader.get_nextbatch(args.batch_size)):
print('itr is %d, and epoch is %d' %(itr, epoch))
itr_fh.write("epoch: " + str(epoch) + " itr: " + str(itr) + "\n")
Z_sample = sample_z(args.dim_z, args.batch_size)
sess.run(d_step, feed_dict={X:real_batch , Z:Z_sample})
sess.run(g_step, feed_dict={X:real_batch , Z:Z_sample})
sample = sess.run(G_sample, feed_dict={Z:Z_sample})
print("sample size is: ", sample.shape)
if itr==3164: #num_images/batch_size
im_merged = merge_images(sample[:16], [4,4])
plt.imsave('sample_gan_images/im_merged_epoch_%d.png' %(epoch), im_merged )
scipy.misc.imsave('sample_gan_images/im_epoch_%d_itr_%d.png' %(epoch,itr), sample[1])
##merged_summary = sess.run(merged_summary, feed_dict={X:real_batch , Z:Z_sample})
###writer = tf.summary.FileWriter(LOGDIR)
###writer.add_summary(merged_summary, itr)
###d_loss_summary = tf.summary.scalar("Discriminator_Total_Loss", d_loss) #how to store loss at the end of each epoch not at each iteration?
###g_loss_summary = tf.summary.scalar("Generator_Total_Loss", g_loss)
###merged_summary = tf.summary.merge_all()
###writer.add_graph(sess.graph)
###saver.save(sess, save_path='logs_basic_dcgan/gan.ckpt')
train(args)
Comments