from __future__ import print_function

from __future__ import division

from __future__ import absolute_import

import skimage.io as io

io.use_plugin('pil')

from scipy.optimize import fmin_l_bfgs_b

import numpy as np

import time

import argparse

import warnings

from keras.models import Model

from keras.layers import Input

from keras.layers.convolutional import Convolution2D, AveragePooling2D, MaxPooling2D

from keras import backend as K

from keras.utils.data_utils import get_file

from keras.utils.layer_utils import convert_all_kernels_in_model

"""

Neural Style Transfer with Keras 2.0.5

Based on:

https://github.com/fchollet/keras/blob/master/examples/neural_style_transfer.py

-----------------------------------------------------------------------------------------------------------------------

"""

THEANO_WEIGHTS_PATH_NO_TOP = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg16_weights_th_dim_ordering_th_kernels_notop.h5'

TF_WEIGHTS_PATH_NO_TOP = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5'

TH_19_WEIGHTS_PATH_NO_TOP = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg19_weights_th_dim_ordering_th_kernels_notop.h5'

TF_19_WEIGHTS_PATH_NO_TOP = 'https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg19_weights_tf_dim_ordering_tf_kernels_notop.h5'

parser = argparse.ArgumentParser(description='Neural style transfer with Keras.')

parser.add_argument('base_image_path', metavar='base', type=str,

help='Path to the image to transform.')

parser.add_argument('syle_image_paths', metavar='ref', nargs='+', type=str,

help='Path to the style reference image.')

parser.add_argument('result_prefix', metavar='res_prefix', type=str,

help='Prefix for the saved results.')

parser.add_argument("--style_masks", type=str, default=None, nargs='+',

help='Masks for style images')

parser.add_argument("--content_mask", type=str, default=None,

help='Masks for the content image')

parser.add_argument("--color_mask", type=str, default=None,

help='Mask for color preservation')

parser.add_argument("--image_size", dest="img_size", default=400, type=int,

help='Minimum image size')

parser.add_argument("--content_weight", dest="content_weight", default=0.025, type=float,

help="Weight of content")

parser.add_argument("--style_weight", dest="style_weight", nargs='+', default=[1], type=float,

help="Weight of style, can be multiple for multiple styles")

parser.add_argument("--style_scale", dest="style_scale", default=1.0, type=float,

help="Scale the weighing of the style")

parser.add_argument("--total_variation_weight", dest="tv_weight", default=8.5e-5, type=float,

help="Total Variation weight")

parser.add_argument("--num_iter", dest="num_iter", default=10, type=int,

help="Number of iterations")

parser.add_argument("--model", default="vgg16", type=str,

help="Choices are 'vgg16' and 'vgg19'")

parser.add_argument("--content_loss_type", default=0, type=int,

help='Can be one of 0, 1 or 2. Readme contains the required information of each mode.')

parser.add_argument("--rescale_image", dest="rescale_image", default="False", type=str,

help="Rescale image after execution to original dimentions")

parser.add_argument("--rescale_method", dest="rescale_method", default="bilinear", type=str,

help="Rescale image algorithm")

parser.add_argument("--maintain_aspect_ratio", dest="maintain_aspect_ratio", default="True", type=str,

help="Maintain aspect ratio of loaded images")

parser.add_argument("--content_layer", dest="content_layer", default="conv5_2", type=str,

help="Content layer used for content loss.")

parser.add_argument("--init_image", dest="init_image", default="content", type=str,

help="Initial image used to generate the final image. Options are 'content', 'noise', or 'gray'")

parser.add_argument("--pool_type", dest="pool", default="max", type=str,

help='Pooling type. Can be "ave" for average pooling or "max" for max pooling')

parser.add_argument('--preserve_color', dest='color', default="False", type=str,

help='Preserve original color in image')

parser.add_argument('--min_improvement', default=0.0, type=float,

help='Defines minimum improvement required to continue script')

def str_to_bool(v):

return v.lower() in ("true", "yes", "t", "1")

''' Arguments '''

args = parser.parse_args()

base_image_path = args.base_image_path

style_reference_image_paths = args.syle_image_paths

result_prefix = args.result_prefix

style_image_paths = []

for style_image_path in style_reference_image_paths:

style_image_paths.append(style_image_path)

style_masks_present = args.style_masks is not None

mask_paths = []

if style_masks_present:

for mask_path in args.style_masks:

mask_paths.append(mask_path)

if style_masks_present:

assert len(style_image_paths) == len(mask_paths), "Wrong number of style masks provided.\n" \

"Number of style images = %d, \n" \

"Number of style mask paths = %d." % \

(len(style_image_paths), len(style_masks_present))

content_mask_present = args.content_mask is not None

content_mask_path = args.content_mask

color_mask_present = args.color_mask is not None

rescale_image = str_to_bool(args.rescale_image)

maintain_aspect_ratio = str_to_bool(args.maintain_aspect_ratio)

preserve_color = str_to_bool(args.color)

# these are the weights of the different loss components

content_weight = args.content_weight

total_variation_weight = args.tv_weight

style_weights = []

if len(style_image_paths) != len(args.style_weight):

print("Mismatch in number of style images provided and number of style weights provided. \n"

"Found %d style images and %d style weights. \n"

"Equally distributing weights to all other styles." % (len(style_image_paths), len(args.style_weight)))

weight_sum = sum(args.style_weight) * args.style_scale

count = len(style_image_paths)

for i in range(len(style_image_paths)):

style_weights.append(weight_sum / count)

else:

for style_weight in args.style_weight:

style_weights.append(style_weight * args.style_scale)

# Decide pooling function

pooltype = str(args.pool).lower()

assert pooltype in ["ave", "max"], 'Pooling argument is wrong. Needs to be either "ave" or "max".'

pooltype = 1 if pooltype == "ave" else 0

read_mode = "gray" if args.init_image == "gray" else "color"

# dimensions of the generated picture.

img_width = img_height = 0

img_WIDTH = img_HEIGHT = 0

aspect_ratio = 0

assert args.content_loss_type in [0, 1, 2], "Content Loss Type must be one of 0, 1 or 2"

# util function to open, resize and format pictures into appropriate tensors

def preprocess_image(image_path, load_dims=False, read_mode="color"):

return img

# util function to convert a tensor into a valid image

def deprocess_image(x):

return x

# util function to preserve image color

def original_color_transform(content, generated, mask=None):

return generated

def load_mask(mask_path, shape, return_mask_img=False):

return mask_tensor

def pooling_func(x):

return MaxPooling2D((2, 2), strides=(2, 2))(x)

# get tensor representations of our images

base_image = K.variable(preprocess_image(base_image_path, True, read_mode=read_mode))

style_reference_images = []

for style_path in style_image_paths:

style_reference_images.append(K.variable(preprocess_image(style_path)))

# this will contain our generated image

if K.image_dim_ordering() == 'th':

combination_image = K.placeholder((1, 3, img_width, img_height))

else:

combination_image = K.placeholder((1, img_width, img_height, 3))

image_tensors = [base_image]

for style_image_tensor in style_reference_images:

image_tensors.append(style_image_tensor)

image_tensors.append(combination_image)

nb_tensors = len(image_tensors)

nb_style_images = nb_tensors - 2 # Content and Output image not considered

# combine the various images into a single Keras tensor

input_tensor = K.concatenate(image_tensors, axis=0)

if K.image_dim_ordering() == "th":

shape = (nb_tensors, 3, img_width, img_height)

else:

shape = (nb_tensors, img_width, img_height, 3)

ip = Input(tensor=input_tensor, batch_shape=shape)

# build the VGG16 network with our 3 images as input

x = Convolution2D(64, (3, 3), activation='relu', name='conv1_1', padding='same')(ip)

x = Convolution2D(64, (3, 3), activation='relu', name='conv1_2', padding='same')(x)

x = pooling_func(x)

x = Convolution2D(128, (3, 3), activation='relu', name='conv2_1', padding='same')(x)

x = Convolution2D(128, (3, 3), activation='relu', name='conv2_2', padding='same')(x)

x = pooling_func(x)

x = Convolution2D(256, (3, 3), activation='relu', name='conv3_1', padding='same')(x)

x = Convolution2D(256, (3, 3), activation='relu', name='conv3_2', padding='same')(x)

x = Convolution2D(256, (3, 3), activation='relu', name='conv3_3', padding='same')(x)

if args.model == "vgg19":

x = Convolution2D(256, (3, 3), activation='relu', name='conv3_4', padding='same')(x)

x = pooling_func(x)

x = Convolution2D(512, (3, 3), activation='relu', name='conv4_1', padding='same')(x)

x = Convolution2D(512, (3, 3), activation='relu', name='conv4_2', padding='same')(x)

x = Convolution2D(512, (3, 3), activation='relu', name='conv4_3', padding='same')(x)

if args.model == "vgg19":

x = Convolution2D(512, (3, 3), activation='relu', name='conv4_4', padding='same')(x)

x = pooling_func(x)

x = Convolution2D(512, (3, 3), activation='relu', name='conv5_1', padding='same')(x)

x = Convolution2D(512, (3, 3), activation='relu', name='conv5_2', padding='same')(x)

x = Convolution2D(512, (3, 3), activation='relu', name='conv5_3', padding='same')(x)

if args.model == "vgg19":

x = Convolution2D(512, (3, 3), activation='relu', name='conv5_4', padding='same')(x)

x = pooling_func(x)

model = Model(ip, x)

if K.image_dim_ordering() == "th":

if args.model == "vgg19":

weights = get_file('vgg19_weights_th_dim_ordering_th_kernels_notop.h5', TH_19_WEIGHTS_PATH_NO_TOP, cache_subdir='models')

else:

weights = get_file('vgg16_weights_th_dim_ordering_th_kernels_notop.h5', THEANO_WEIGHTS_PATH_NO_TOP, cache_subdir='models')

else:

if args.model == "vgg19":

weights = get_file('vgg19_weights_tf_dim_ordering_tf_kernels_notop.h5', TF_19_WEIGHTS_PATH_NO_TOP, cache_subdir='models')

else:

weights = get_file('vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5', TF_WEIGHTS_PATH_NO_TOP, cache_subdir='models')

model.load_weights(weights)

if K.backend() == 'tensorflow' and K.image_dim_ordering() == "th":

warnings.warn('You are using the TensorFlow backend, yet you '

'are using the Theano '

'image dimension ordering convention '

'(`image_dim_ordering="th"`). '

'For best performance, set '

'`image_dim_ordering="tf"` in '

'your Keras config '

'at ~/.keras/keras.json.')

convert_all_kernels_in_model(model)

print('Model loaded.')

# get the symbolic outputs of each "key" layer (we gave them unique names).

outputs_dict = dict([(layer.name, layer.output) for layer in model.layers])

shape_dict = dict([(layer.name, layer.output_shape) for layer in model.layers])

# compute the neural style loss

# first we need to define 4 util functions

# the gram matrix of an image tensor (feature-wise outer product)

def gram_matrix(x):

return gram

# the "style loss" is designed to maintain

# the style of the reference image in the generated image.

# It is based on the gram matrices (which capture style) of

# feature maps from the style reference image

# and from the generated image

def style_loss(style, combination, mask_path=None, nb_channels=None):

return K.sum(K.square(S - C)) / (4. * (channels ** 2) * (size ** 2))

# an auxiliary loss function

# designed to maintain the "content" of the

# base image in the generated image

def content_loss(base, combination):

return multiplier * K.sum(K.square(combination - base))

# the 3rd loss function, total variation loss,

# designed to keep the generated image locally coherent

def total_variation_loss(x):

return K.sum(K.pow(a + b, 1.25))

# combine these loss functions into a single scalar

loss = K.variable(0.)

layer_features = outputs_dict[args.content_layer] # 'conv5_2' or 'conv4_2'

base_image_features = layer_features[0, :, :, :]

combination_features = layer_features[nb_tensors - 1, :, :, :]

loss = loss + content_weight * content_loss(base_image_features,

combination_features)

style_masks = []

if style_masks_present:

style_masks = mask_paths # If mask present, pass dictionary of masks to style loss

else:

style_masks = [None for _ in range(nb_style_images)] # If masks not present, pass None to the style loss

channel_index = 1 if K.image_dim_ordering() == "th" else -1

feature_layers = ['conv1_1', 'conv2_1', 'conv3_1', 'conv4_1', 'conv5_1']

for layer_name in feature_layers:

layer_features = outputs_dict[layer_name]

shape = shape_dict[layer_name]

combination_features = layer_features[nb_tensors - 1, :, :, :]

style_reference_features = layer_features[1:nb_tensors - 1, :, :, :]

sl = []

for j in range(nb_style_images):

sl.append(style_loss(style_reference_features[j], combination_features, style_masks[j], shape))

for j in range(nb_style_images):

loss = loss + (style_weights[j] / len(feature_layers)) * sl[j]

loss = loss + total_variation_weight * total_variation_loss(combination_image)

# get the gradients of the generated image wrt the loss

grads = K.gradients(loss, combination_image)

outputs = [loss]

if type(grads) in {list, tuple}:

outputs += grads

else:

outputs.append(grads)

f_outputs = K.function([combination_image], outputs)

def eval_loss_and_grads(x):

return loss_value, grad_values

# this Evaluator class makes it possible

# to compute loss and gradients in one pass

# while retrieving them via two separate functions,

# "loss" and "grads". This is done because scipy.optimize

# requires separate functions for loss and gradients,

# but computing them separately would be inefficient.

class Evaluator(object):

return grad_values

evaluator = Evaluator()

# run scipy-based optimization (L-BFGS) over the pixels of the generated image

# so as to minimize the neural style loss

if "content" in args.init_image or "gray" in args.init_image:

x = preprocess_image(base_image_path, True, read_mode=read_mode)

elif "noise" in args.init_image:

x = np.random.uniform(0, 255, (1, img_width, img_height, 3)) - 128.

if K.image_dim_ordering() == "th":

x = x.transpose((0, 3, 1, 2))

else:

print("Using initial image : ", args.init_image)

x = preprocess_image(args.init_image, read_mode=read_mode)

# We require original image if we are to preserve color in YCbCr mode

if preserve_color:

content = io(base_image_path, mode="YCbCr")

content = io(content, (img_width, img_height))

if color_mask_present:

if K.image_dim_ordering() == "th":

color_mask_shape = (None, None, img_width, img_height)

else:

color_mask_shape = (None, img_width, img_height, None)

color_mask = load_mask(args.color_mask, color_mask_shape, return_mask_img=True)

else:

color_mask = None

else:

color_mask = None

num_iter = args.num_iter

prev_min_val = -1

improvement_threshold = float(args.min_improvement)

for i in range(num_iter):

print("Starting iteration %d of %d" % ((i + 1), num_iter))

start_time = time.time()

x, min_val, info = fmin_l_bfgs_b(evaluator.loss, x.flatten(), fprime=evaluator.grads, maxfun=20)

if prev_min_val == -1:

prev_min_val = min_val

improvement = (prev_min_val - min_val) / prev_min_val * 100

print('Current loss value:', min_val, " Improvement : %0.3f" % improvement, "%")

prev_min_val = min_val

# save current generated image

img = deprocess_image(x.copy())

if preserve_color and content is not None:

img = original_color_transform(content, img, mask=color_mask)

if not rescale_image:

img_ht = int(img_width * aspect_ratio)

print("Rescaling Image to (%d, %d)" % (img_width, img_ht))

img = io(img, (img_width, img_ht), interp=args.rescale_method)

if rescale_image:

print("Rescaling Image to (%d, %d)" % (img_WIDTH, img_HEIGHT))

img = io(img, (img_WIDTH, img_HEIGHT), interp=args.rescale_method)

fname = result_prefix + '_at_iteration_%d.png' % (i + 1)

io(fname, img)

end_time = time.time()

print('Image saved as', fname)

print('Iteration %d completed in %ds' % (i + 1, end_time - start_time))

if improvement_threshold is not 0.0:

if improvement < improvement_threshold and improvement is not 0.0:

print("Improvement (%f) is less than improvement threshold (%f). Early stopping script." % (

improvement, improvement_threshold))

exit()