From 38008e85d8a3af1d25fbe2ecee43a455a7660a40 Mon Sep 17 00:00:00 2001 From: Jacob Gildenblat Date: Fri, 4 Mar 2016 16:19:39 +0200 Subject: [PATCH] DCGAN --- dcgan.py | 160 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 160 insertions(+) create mode 100644 dcgan.py diff --git a/dcgan.py b/dcgan.py new file mode 100644 index 0000000..9cc6ba8 --- /dev/null +++ b/dcgan.py @@ -0,0 +1,160 @@ +from keras.models import Sequential +from keras.layers import Dense +from keras.layers import Reshape +from keras.layers.core import Activation +from keras.layers.advanced_activations import LeakyReLU +from keras.layers.normalization import BatchNormalization +from keras.layers.convolutional import UpSampling2D +from keras.layers.convolutional import Convolution2D +from keras.layers.core import Flatten +from keras.optimizers import Adam +from keras import backend as K +import numpy as np +import sys, glob +import cv2 +import os +import argparse + +def generator_model(): + model = Sequential() + model.add(Dense(input_dim=100, output_dim=1024*4*4)) + model.add(BatchNormalization()) + model.add(Activation('relu')) + model.add(Reshape(dims=(1024, 4, 4))) + model.add(UpSampling2D(size=(2, 2))) + model.add(Convolution2D(512, 5, 5, border_mode='same')) + model.add(BatchNormalization()) + model.add(Activation('relu')) + model.add(UpSampling2D(size=(2, 2))) + model.add(Convolution2D(256, 5, 5, border_mode='same')) + model.add(BatchNormalization()) + model.add(Activation('relu')) + model.add(UpSampling2D(size=(2, 2))) + model.add(Convolution2D(128, 5, 5, border_mode='same')) + model.add(BatchNormalization()) + model.add(Activation('relu')) + model.add(UpSampling2D(size=(2, 2))) + model.add(Convolution2D(3, 5, 5, border_mode='same')) + model.add(Activation('tanh')) + return model + +def descriminator_model(): + model = Sequential() + model.add(Convolution2D(128, 5, 5, subsample=(2, 2), input_shape=(3, 64, 64), border_mode = 'same')) + model.add(LeakyReLU(0.2)) + model.add(BatchNormalization()) + model.add(Convolution2D(256, 5, 5, subsample=(2, 2), border_mode = 'same')) + model.add(BatchNormalization()) + model.add(LeakyReLU(0.2)) + model.add(Convolution2D(512, 5, 5, subsample=(2, 2), border_mode = 'same')) + model.add(BatchNormalization()) + model.add(LeakyReLU(0.2)) + model.add(Convolution2D(1024, 5, 5, subsample=(2, 2), border_mode = 'same')) + model.add(BatchNormalization()) + model.add(LeakyReLU(0.2)) + model.add(Flatten()) + model.add(Dense(output_dim=1)) + model.add(Activation('sigmoid')) + return model + +def generator_containing_descriminator(generator, descriminator): + model = Sequential() + model.add(generator) + descriminator.trainable = False + model.add(descriminator) + return model + +def load_image(path): + img = cv2.imread(path, 1) + img = np.float32(cv2.resize(img, (64, 64))) / 127.5 - 1 + img = np.rollaxis(img, 2, 0) + return img + +def get_batches(paths, batch_size): + for i in range(len(paths)/batch_size): + yield i, [load_image(path) for path in paths[i*batch_size : (i + 1) * batch_size]] + +def train(path, BATCH_SIZE): + print "Loading paths.." + paths = glob.glob(os.path.join(path, "*.jpg")) + print "Got paths.." + + descriminator = descriminator_model() + generator = generator_model() + descriminator_on_generator = generator_containing_descriminator(generator, descriminator) + adam=Adam(lr=0.0002, beta_1=0.5, beta_2=0.999, epsilon=1e-08) + generator.compile(loss='binary_crossentropy', optimizer=adam) + descriminator_on_generator.compile(loss='binary_crossentropy', optimizer=adam) + descriminator.trainable = True + descriminator.compile(loss='binary_crossentropy', optimizer=adam) + + for epoch in range(5): + print "Epoch is", epoch + print "Number of batches", len(paths) / BATCH_SIZE + for index, image_batch in get_batches(paths, batch_size=BATCH_SIZE): + noise = np.zeros((BATCH_SIZE, 100)) + for i in range(BATCH_SIZE): + noise[i, : ] = np.random.uniform(-1, 1, 100) + + print 'Generating images..' + generated_images = generator.predict(noise) + print 'Generated..' + for i, img in enumerate(generated_images): + rolled = np.rollaxis(img, 0, 3) + cv2.imwrite(str(i) + ".jpg", np.uint8(255 * 0.5 * (rolled + 1.0))) + + X = np.concatenate((image_batch, generated_images)) + + y = [1] * BATCH_SIZE + [0] * BATCH_SIZE + print "Batch", index, "Training descriminator.." + d_loss = descriminator.train_on_batch(X, y) + + for j in range(1): + noise = np.zeros((BATCH_SIZE, 100)) + for i in range(BATCH_SIZE): + noise[i, : ] = np.random.uniform(-1, 1, 100) + + + print "Training generator.." + g_loss = descriminator_on_generator.train_on_batch(noise, [1] * BATCH_SIZE) + print "Generator loss", g_loss, "Descriminator loss", d_loss, "Total:", g_loss[0] + d_loss[0] + + if index % 10 == 9: + print 'Saving weights..' + generator.save_weights('generator', True) + descriminator.save_weights('descriminator', True) + + +def generate(): + BATCH_SIZE = 32 + generator = generator_model() + adam=Adam(lr=0.0002, beta_1=0.5, beta_2=0.999, epsilon=1e-08) + generator.compile(loss='binary_crossentropy', optimizer=adam) + generator.load_weights('generator') + + noise = np.zeros((BATCH_SIZE, 100)) + for i in range(BATCH_SIZE): + noise[i, : ] = np.random.uniform(0, 1, 100) + + print 'Generating images..' + generated_images = [np.rollaxis(img, 0, 3) for img in generator.predict(noise)] + for index, img in enumerate(generated_images): + cv2.imwrite("{}.jpg".format(index), np.uint8(255 * 0.5 * (img + 1.0))) + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--mode", type = str) + parser.add_argument("--path", type = str) + parser.add_argument("--batch_size", type = int, default = 128) + args = parser.parse_args() + return args + +if __name__ == "__main__": + args = get_args() + path = args.path + + if args.mode == "train": + train(path = args.path, BATCH_SIZE = args.batch_size) + elif args.mode == "generate": + generate() \ No newline at end of file