In [ ]:
import time
import numpy as np
import pdb
import seaborn as sns
import pandas as pd
import glob
import os
import sys
from os.path import join
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import animation
from keras.models import Sequential
from keras.models import Model
from keras.layers.core import Activation, Dense
from keras import backend as K
from keras import optimizers
sys.path.append(join(os.getcwd(), os.pardir))
from utils import image_processing
sns.set_style("dark")
sns.set_context("paper")
%matplotlib notebook
In [ ]:
RES_DIR = join(*[os.pardir]*2, 'data', 'sketch_dataset')
In [ ]:
def load_dataset_info(dirpath):
data = pd.DataFrame(columns=['dirpath', 'ref', 'sketch', 'sketch2', 'lineart'])
for f in glob.glob(join(dirpath, '**', "*.jpg"), recursive=True):
filepath, filename = f.rsplit('\\', 1)
img_id, category = filename.split('.')[0].split('_')
data.set_value(img_id, category, filename)
data.set_value(img_id, 'dirpath', filepath)
return data
In [ ]:
dataset_info = load_dataset_info(RES_DIR)
dataset_info.head()
In [ ]:
# load sketch
content_image = None
with Image.open(os.path.join(RES_DIR, 'superman.jpg')) as img:
img = img.resize((height, width))
content_image = np.asarray(img, dtype='float32')
plt.imshow(img.convert(mode='RGB'))
plt.show()
In [ ]:
f, axarr = plt.subplots(len(dataset_info), 3)
categories = {0:'ref', 1:'sketch', 2:'lineart'}
for row_idx, (img_id, row) in enumerate(dataset_info.iterrows()):
for i in range(3):
img = plt.imread(join(row['dirpath'], row[categories[i]]))
axarr[row_idx, i].imshow(img)
axarr[row_idx, i].axis('off')
In [ ]:
from skimage import data
from skimage import transform
from skimage import io
In [ ]:
def augment(imgs_info, n, suffix_folder=""):
for i in range(n):
for img_id, row in imgs_info.iterrows():
rotation = np.random.randint(360)
for cat in ['ref', 'sketch', 'sketch2', 'lineart']:
if not pd.isnull(row[cat]):
origin_img = plt.imread(join(row['dirpath'], row[cat]))
dest_img = transform.rotate(origin_img, rotation, mode='edge')
filename = "{}{}{}_{}.jpg".format(img_id, 'gen', i, cat)
io.imsave(join(row['dirpath'], suffix_folder, filename), dest_img)
In [ ]:
augment(dataset_info, 10, 'augmented')
In [ ]:
from keras.models import Sequential
from keras.layers import Convolution2D, MaxPooling2D, Conv2DTranspose
from keras.layers import Activation, Dropout, Flatten, Dense
from keras.preprocessing import image
from keras import optimizers
In [ ]:
img_shape = (128, 128, 3)
In [ ]:
X = image_processing.load_data(dataset_info.apply(lambda x : join(x['dirpath'], x['lineart']), axis=1).values, img_shape[:2])
y = image_processing.load_data(dataset_info.apply(lambda x : join(x['dirpath'], x['ref']), axis=1).values, img_shape[:2])
In [ ]:
X_train = X#(X/255)
y_train = y#(y/255)
print(X_train.shape)
print(y_train.shape)
In [ ]:
model = Sequential()
model.add(Convolution2D(32, (5, 5), strides=(2, 2), padding='same', input_shape=img_shape,
activation='relu'))
model.add(Convolution2D(64, (3, 3), padding='same',
activation='relu'))
model.add(Convolution2D(128, (5, 5), strides=(2, 2), padding='same',
activation='relu'))
model.add(Convolution2D(256, (3, 3), padding='same',
activation='relu'))
model.add(Convolution2D(256, (3, 3), padding='same',
activation='relu'))
model.add(Conv2DTranspose(128, (5, 5), strides=(2, 2), padding='same',
activation='relu'))
model.add(Convolution2D(64, (3, 3), padding='same',
activation='relu'))
model.add(Conv2DTranspose(32, (3, 3), strides=(2, 2), padding='same',
activation='relu'))
model.add(Convolution2D(3, (5, 5), padding='same',
activation='sigmoid'))
In [ ]:
model.summary()
In [ ]:
optimizer = optimizers.Adam(lr=0.001)
In [ ]:
model.compile(loss='mean_squared_error', optimizer=optimizer, metrics=['accuracy'])
In [ ]:
model.fit(y_train, y_train, batch_size=2, epochs=5)
In [ ]:
plt.imshow(X_train[3])
In [ ]:
plt.imshow(y_train[3])
In [ ]:
pred = model.predict(y_train[2].reshape((1, *img_shape)))
In [ ]:
pred.shape
In [ ]:
plt.imshow(pred[0])
In [ ]:
pred
In [ ]:
from keras.models import Sequential
from keras.layers.core import Activation, Dense
from keras import backend as K
from keras import optimizers
In [ ]:
img_shape = (128, 128, 3)
In [ ]:
X_train = image_processing.load_data(dataset_info.apply(lambda x : join(x['dirpath'], x['ref']), axis=1).values, img_shape[:2])
y_train = image_processing.load_data(dataset_info.apply(lambda x : join(x['dirpath'], x['lineart']), axis=1).values, img_shape[:2])
print(X_train.shape)
print(y_train.shape)
In [ ]:
def generator(img_shape):
model = Sequential()
model.add(Convolution2D(32, (5, 5), strides=(2, 2), padding='same', input_shape=img_shape,
activation='relu'))
model.add(Convolution2D(64, (3, 3), padding='same',
activation='relu'))
model.add(Convolution2D(128, (5, 5), strides=(2, 2), padding='same',
activation='relu'))
model.add(Convolution2D(256, (3, 3), padding='same',
activation='relu'))
model.add(Convolution2D(256, (3, 3), padding='same',
activation='relu'))
model.add(Conv2DTranspose(128, (5, 5), strides=(2, 2), padding='same',
activation='relu'))
model.add(Convolution2D(64, (3, 3), padding='same',
activation='relu'))
model.add(Conv2DTranspose(32, (3, 3), strides=(2, 2), padding='same',
activation='relu'))
model.add(Convolution2D(3, (5, 5), padding='same',
activation='sigmoid'))
return model
In [ ]:
def discriminator(img_shape):
model = Sequential()
model.add(Convolution2D(32, (5, 5), strides=(2, 2), padding='same', input_shape=img_shape,
activation='relu'))
model.add(Convolution2D(64, (3, 3), padding='same',
activation='relu'))
model.add(Convolution2D(128, (5, 5), strides=(2, 2), padding='same',
activation='relu'))
model.add(Convolution2D(256, (3, 3), padding='same',
activation='relu'))
model.add(Convolution2D(256, (3, 3), padding='same',
activation='relu'))
model.add(Conv2DTranspose(128, (5, 5), strides=(2, 2), padding='same',
activation='relu'))
model.add(Convolution2D(64, (3, 3), padding='same',
activation='relu'))
model.add(Conv2DTranspose(32, (3, 3), strides=(2, 2), padding='same',
activation='relu'))
model.add(Flatten())
model.add(Dense(512, activation='relu'))
#model.add(Dropout(0.5))
model.add(Dense(1, activation=K.sigmoid))
return model
In [ ]:
# init GAN components
d = discriminator(img_shape)
g = generator(img_shape)
In [ ]:
# discriminator model
optimizer = optimizers.RMSprop(lr=0.0008, clipvalue=1.0, decay=6e-8)
discriminator_model = d
discriminator_model.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=['accuracy'])
In [ ]:
# adversarial model
optimizer = optimizers.RMSprop(lr=0.0004, clipvalue=1.0, decay=3e-8)
adversarial_model = Sequential()
adversarial_model.add(g)
adversarial_model.add(d)
adversarial_model.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=['accuracy'])
In [ ]:
batch_size = 64
epochs = 100
In [ ]:
for step in range(epochs):
# generate data
# first we sample from the true distribution, then we generate some
# "fake" data by feeding noise to the generator
true_sample = np.reshape(gaussian_d.sample(batch_size), (batch_size, 1))
noise = generator_d.sample(batch_size)
fake_sample = g.predict(noise)
#pdb.set_trace()
# train discriminator
# feed true and fake samples with respective labels (1, 0) to the discriminator
x = np.reshape(np.concatenate((true_sample, fake_sample)), (batch_size*2, 1))
y = np.ones([batch_size*2, 1])
y[batch_size:, :] = 0
d_loss = discriminator_model.train_on_batch(x, y)
# train GAN
# feed noise to the model and expect true (1) response from discriminator,
# which is in turn fed with data generated by the generator
noise = np.reshape(generator_d.sample(batch_size), (batch_size, 1))
y = np.ones([batch_size, 1])
a_loss = adversarial_model.train_on_batch(noise, y)
log_mesg = "%d: [D loss: %f, acc: %f]" % (step, d_loss[0], d_loss[1])
log_mesg = "%s [A loss: %f, acc: %f]" % (log_mesg, a_loss[0], a_loss[1])