Sketch Cleanup


Intro

Exploring the task of sketch cleaning (also automatized linearts).

Based on Edgar Simo-Serra・Sketch Simplification


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')

Load Data


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')

Data Augmentation


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')

Train (CNN)


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

Train (GAN)


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])