'''Neural doodle with Keras # Script Usage ## Arguments ``` --nlabels: # of regions (colors) in mask images --style-image: image to learn style from --style-mask: semantic labels for style image --target-mask: semantic labels for target image (your doodle) --content-image: optional image to learn content from --target-image-prefix: path prefix for generated target images ``` ## Example 1: doodle using a style image, style mask and target mask. ``` python neural_doodle.py --nlabels 4 --style-image Monet/style.png \ --style-mask Monet/style_mask.png --target-mask Monet/target_mask.png \ --target-image-prefix generated/monet ``` ## Example 2: doodle using a style image, style mask, target mask and an optional content image. ``` python neural_doodle.py --nlabels 4 --style-image Renoir/style.png \ --style-mask Renoir/style_mask.png --target-mask Renoir/target_mask.png \ --content-image Renoir/creek.jpg \ --target-image-prefix generated/renoir ``` # References - [Dmitry Ulyanov's blog on fast-neural-doodle] (http://dmitryulyanov.github.io/feed-forward-neural-doodle/) - [Torch code for fast-neural-doodle] (https://github.com/DmitryUlyanov/fast-neural-doodle) - [Torch code for online-neural-doodle] (https://github.com/DmitryUlyanov/online-neural-doodle) - [Paper Texture Networks: Feed-forward Synthesis of Textures and Stylized Images] (http://arxiv.org/abs/1603.03417) - [Discussion on parameter tuning] (https://github.com/keras-team/keras/issues/3705) # Resources Example images can be downloaded from https://github.com/DmitryUlyanov/fast-neural-doodle/tree/master/data ''' from __future__ import print_function import time import argparse import numpy as np from scipy.optimize import fmin_l_bfgs_b from keras import backend as K from keras.layers import Input, AveragePooling2D from keras.models import Model from keras.preprocessing.image import load_img, save_img, img_to_array from keras.applications import vgg19 # Command line arguments parser = argparse.ArgumentParser(description='Keras neural doodle example') parser.add_argument('--nlabels', type=int, help='number of semantic labels' ' (regions in differnet colors)' ' in style_mask/target_mask') parser.add_argument('--style-image', type=str, help='path to image to learn style from') parser.add_argument('--style-mask', type=str, help='path to semantic mask of style image') parser.add_argument('--target-mask', type=str, help='path to semantic mask of target image') parser.add_argument('--content-image', type=str, default=None, help='path to optional content image') parser.add_argument('--target-image-prefix', type=str, help='path prefix for generated results') args = parser.parse_args() style_img_path = args.style_image style_mask_path = args.style_mask target_mask_path = args.target_mask content_img_path = args.content_image target_img_prefix = args.target_image_prefix use_content_img = content_img_path is not None num_labels = args.nlabels num_colors = 3 # RGB # determine image sizes based on target_mask ref_img = img_to_array(load_img(target_mask_path)) img_nrows, img_ncols = ref_img.shape[:2] num_iterations = 50 total_variation_weight = 50. style_weight = 1. content_weight = 0.1 if use_content_img else 0 content_feature_layers = ['block5_conv2'] # To get better generation qualities, use more conv layers for style features style_feature_layers = ['block1_conv1', 'block2_conv1', 'block3_conv1', 'block4_conv1', 'block5_conv1'] # helper functions for reading/processing images def preprocess_image(image_path): img = load_img(image_path, target_size=(img_nrows, img_ncols)) img = img_to_array(img) img = np.expand_dims(img, axis=0) img = vgg19.preprocess_input(img) return img def deprocess_image(x): if K.image_data_format() == 'channels_first': x = x.reshape((3, img_nrows, img_ncols)) x = x.transpose((1, 2, 0)) else: x = x.reshape((img_nrows, img_ncols, 3)) # Remove zero-center by mean pixel x[:, :, 0] += 103.939 x[:, :, 1] += 116.779 x[:, :, 2] += 123.68 # 'BGR'->'RGB' x = x[:, :, ::-1] x = np.clip(x, 0, 255).astype('uint8') return x def kmeans(xs, k): assert xs.ndim == 2 try: from sklearn.cluster import k_means _, labels, _ = k_means(xs.astype('float64'), k) except ImportError: from scipy.cluster.vq import kmeans2 _, labels = kmeans2(xs, k, missing='raise') return labels def load_mask_labels(): '''Load both target and style masks. A mask image (nr x nc) with m labels/colors will be loaded as a 4D boolean tensor: (1, m, nr, nc) for 'channels_first' or (1, nr, nc, m) for 'channels_last' ''' target_mask_img = load_img(target_mask_path, target_size=(img_nrows, img_ncols)) target_mask_img = img_to_array(target_mask_img) style_mask_img = load_img(style_mask_path, target_size=(img_nrows, img_ncols)) style_mask_img = img_to_array(style_mask_img) if K.image_data_format() == 'channels_first': mask_vecs = np.vstack([style_mask_img.reshape((3, -1)).T, target_mask_img.reshape((3, -1)).T]) else: mask_vecs = np.vstack([style_mask_img.reshape((-1, 3)), target_mask_img.reshape((-1, 3))]) labels = kmeans(mask_vecs, num_labels) style_mask_label = labels[:img_nrows * img_ncols].reshape((img_nrows, img_ncols)) target_mask_label = labels[img_nrows * img_ncols:].reshape((img_nrows, img_ncols)) stack_axis = 0 if K.image_data_format() == 'channels_first' else -1 style_mask = np.stack([style_mask_label == r for r in range(num_labels)], axis=stack_axis) target_mask = np.stack([target_mask_label == r for r in range(num_labels)], axis=stack_axis) return (np.expand_dims(style_mask, axis=0), np.expand_dims(target_mask, axis=0)) # Create tensor variables for images if K.image_data_format() == 'channels_first': shape = (1, num_colors, img_nrows, img_ncols) else: shape = (1, img_nrows, img_ncols, num_colors) style_image = K.variable(preprocess_image(style_img_path)) target_image = K.placeholder(shape=shape) if use_content_img: content_image = K.variable(preprocess_image(content_img_path)) else: content_image = K.zeros(shape=shape) images = K.concatenate([style_image, target_image, content_image], axis=0) # Create tensor variables for masks raw_style_mask, raw_target_mask = load_mask_labels() style_mask = K.variable(raw_style_mask.astype('float32')) target_mask = K.variable(raw_target_mask.astype('float32')) masks = K.concatenate([style_mask, target_mask], axis=0) # index constants for images and tasks variables STYLE, TARGET, CONTENT = 0, 1, 2 # Build image model, mask model and use layer outputs as features # image model as VGG19 image_model = vgg19.VGG19(include_top=False, input_tensor=images) # mask model as a series of pooling mask_input = Input(tensor=masks, shape=(None, None, None), name='mask_input') x = mask_input for layer in image_model.layers[1:]: name = 'mask_%s' % layer.name if 'conv' in layer.name: x = AveragePooling2D((3, 3), padding='same', strides=( 1, 1), name=name)(x) elif 'pool' in layer.name: x = AveragePooling2D((2, 2), name=name)(x) mask_model = Model(mask_input, x) # Collect features from image_model and task_model image_features = {} mask_features = {} for img_layer, mask_layer in zip(image_model.layers, mask_model.layers): if 'conv' in img_layer.name: assert 'mask_' + img_layer.name == mask_layer.name layer_name = img_layer.name img_feat, mask_feat = img_layer.output, mask_layer.output image_features[layer_name] = img_feat mask_features[layer_name] = mask_feat # Define loss functions def gram_matrix(x): assert K.ndim(x) == 3 features = K.batch_flatten(x) gram = K.dot(features, K.transpose(features)) return gram def region_style_loss(style_image, target_image, style_mask, target_mask): '''Calculate style loss between style_image and target_image, for one common region specified by their (boolean) masks ''' assert 3 == K.ndim(style_image) == K.ndim(target_image) assert 2 == K.ndim(style_mask) == K.ndim(target_mask) if K.image_data_format() == 'channels_first': masked_style = style_image * style_mask masked_target = target_image * target_mask num_channels = K.shape(style_image)[0] else: masked_style = K.permute_dimensions( style_image, (2, 0, 1)) * style_mask masked_target = K.permute_dimensions( target_image, (2, 0, 1)) * target_mask num_channels = K.shape(style_image)[-1] num_channels = K.cast(num_channels, dtype='float32') s = gram_matrix(masked_style) / K.mean(style_mask) / num_channels c = gram_matrix(masked_target) / K.mean(target_mask) / num_channels return K.mean(K.square(s - c)) def style_loss(style_image, target_image, style_masks, target_masks): '''Calculate style loss between style_image and target_image, in all regions. ''' assert 3 == K.ndim(style_image) == K.ndim(target_image) assert 3 == K.ndim(style_masks) == K.ndim(target_masks) loss = K.variable(0) for i in range(num_labels): if K.image_data_format() == 'channels_first': style_mask = style_masks[i, :, :] target_mask = target_masks[i, :, :] else: style_mask = style_masks[:, :, i] target_mask = target_masks[:, :, i] loss = loss + region_style_loss(style_image, target_image, style_mask, target_mask) return loss def content_loss(content_image, target_image): return K.sum(K.square(target_image - content_image)) def total_variation_loss(x): assert 4 == K.ndim(x) if K.image_data_format() == 'channels_first': a = K.square(x[:, :, :img_nrows - 1, :img_ncols - 1] - x[:, :, 1:, :img_ncols - 1]) b = K.square(x[:, :, :img_nrows - 1, :img_ncols - 1] - x[:, :, :img_nrows - 1, 1:]) else: a = K.square(x[:, :img_nrows - 1, :img_ncols - 1, :] - x[:, 1:, :img_ncols - 1, :]) b = K.square(x[:, :img_nrows - 1, :img_ncols - 1, :] - x[:, :img_nrows - 1, 1:, :]) return K.sum(K.pow(a + b, 1.25)) # Overall loss is the weighted sum of content_loss, style_loss and tv_loss # Each individual loss uses features from image/mask models. loss = K.variable(0) for layer in content_feature_layers: content_feat = image_features[layer][CONTENT, :, :, :] target_feat = image_features[layer][TARGET, :, :, :] loss = loss + content_weight * content_loss(content_feat, target_feat) for layer in style_feature_layers: style_feat = image_features[layer][STYLE, :, :, :] target_feat = image_features[layer][TARGET, :, :, :] style_masks = mask_features[layer][STYLE, :, :, :] target_masks = mask_features[layer][TARGET, :, :, :] sl = style_loss(style_feat, target_feat, style_masks, target_masks) loss = loss + (style_weight / len(style_feature_layers)) * sl loss = loss + total_variation_weight * total_variation_loss(target_image) loss_grads = K.gradients(loss, target_image) # Evaluator class for computing efficiency outputs = [loss] if isinstance(loss_grads, (list, tuple)): outputs += loss_grads else: outputs.append(loss_grads) f_outputs = K.function([target_image], outputs) def eval_loss_and_grads(x): if K.image_data_format() == 'channels_first': x = x.reshape((1, 3, img_nrows, img_ncols)) else: x = x.reshape((1, img_nrows, img_ncols, 3)) outs = f_outputs([x]) loss_value = outs[0] if len(outs[1:]) == 1: grad_values = outs[1].flatten().astype('float64') else: grad_values = np.array(outs[1:]).flatten().astype('float64') return loss_value, grad_values class Evaluator(object): def __init__(self): self.loss_value = None self.grads_values = None def loss(self, x): assert self.loss_value is None loss_value, grad_values = eval_loss_and_grads(x) self.loss_value = loss_value self.grad_values = grad_values return self.loss_value def grads(self, x): assert self.loss_value is not None grad_values = np.copy(self.grad_values) self.loss_value = None self.grad_values = None return grad_values evaluator = Evaluator() # Generate images by iterative optimization if K.image_data_format() == 'channels_first': x = np.random.uniform(0, 255, (1, 3, img_nrows, img_ncols)) - 128. else: x = np.random.uniform(0, 255, (1, img_nrows, img_ncols, 3)) - 128. for i in range(num_iterations): print('Start of iteration', i, '/', num_iterations) start_time = time.time() x, min_val, info = fmin_l_bfgs_b(evaluator.loss, x.flatten(), fprime=evaluator.grads, maxfun=20) print('Current loss value:', min_val) # save current generated image img = deprocess_image(x.copy()) fname = target_img_prefix + '_at_iteration_%d.png' % i save_img(fname, img) end_time = time.time() print('Image saved as', fname) print('Iteration %d completed in %ds' % (i, end_time - start_time))