style-transfer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Style Transfer with Deep Learning using the VGG-19 network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from six.moves import urllib\n",
    "from scipy.io import loadmat\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.misc import imresize\n",
    "%matplotlib inline\n",
    "\n",
    "\n",
    "def download_hook(count, block_size, total_size):\n",
    "        if count % 20 == 0 or count * block_size == total_size:\n",
    "            percentage = 100.0 * count * block_size / total_size\n",
    "            barstring = [\"=\" for _ in range(int(percentage / 2.0))] + [\">\"] + [\".\" for _ in range(50 - int(percentage / 2.0))]\n",
    "            barstring = \"[\" + \"\".join(barstring) + \"]\"\n",
    "            outstring = '%02.02f%% (%02.02f of %02.02f MB)\\t\\t' + barstring\n",
    "            print(outstring % (percentage, count * block_size / 1024.0 / 1024.0, total_size / 1024.0 / 1024.0), end='\\r')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The VGG-19 model is quite large, so be a little patient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"http://www.vlfeat.org/matconvnet/models/imagenet-vgg-verydeep-19.mat\"\n",
    "fname = \"vgg-19.mat\"\n",
    "if not os.path.exists(fname):\n",
    "    print(\"Downloading ...\")\n",
    "    filepath, _ = urllib.request.urlretrieve(path, filename=fname, reporthook=download_hook)\n",
    "    print(\"Done.\")\n",
    "\n",
    "    \n",
    "if not os.path.exists(\"content.jpg\"):    \n",
    "    urllib.request.urlretrieve(\"\", filename=\"content.jpg\") # Attribution: Gage Skidmore\n",
    "    urllib.request.urlretrieve(\"\", filename=\"style.jpg\")    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Find a description (in the form of python code) of the loaded model at\n",
    "https://github.com/chiphuyen/stanford-tensorflow-tutorials/blob/master/assignments/style_transfer/vgg_model.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "original_layers = loadmat(fname)[\"layers\"][0]\n",
    "original_layers.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The downloaded file contains the VGG-19 model, consisting of 43 trained layers. You can access the neccessary information as follows:\n",
    "* The name of layer `i` can be found by accessing `original_layers[i][0][0][0][0]`\n",
    "* The layer weight matrix can be found by accessing `original_layers[i][0][2][0][0]`\n",
    "* The bias can be found by accessing `original_layers[i][0][2][0][1]`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_layer_name(i):\n",
    "    return original_layers[i][0][0][0][0]\n",
    "\n",
    "def get_layer_weights(i):\n",
    "    return original_layers[i][0][0][2][0][0]\n",
    "\n",
    "def get_layer_bias(i):\n",
    "    return original_layers[i][0][0][2][0][1]\n",
    "    \n",
    "def get_layer_params(i):\n",
    "    return (get_layer_weights(i), get_layer_bias(i))\n",
    "\n",
    "layer_names = [get_layer_name(i) for i in range(len(original_layers))]\n",
    "\n",
    "def get_layer_by_name(name):\n",
    "    return layer_names.index(name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get an inuition of what the VGG-19 network looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\", \".join(layer_names))\n",
    "conv_layers = [ln for ln in layer_names if ln.startswith(\"conv\")]\n",
    "pool_layers = [ln for ln in layer_names if ln.startswith(\"pool\")]\n",
    "\n",
    "print(conv_layers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The network consists of a series of convolutional layers making sense of the input content, finally the network makes a decision about what is being depicted in the input image by moving the content through some fully connected layers followed by a softmax function.\n",
    "\n",
    "We now have an understanding of what the VGG-19 network consists of, and we also are able to access its weights and biases. What we need to do next is move the model to tensorflow. We only need to rebuild those parts that we're going to use later on--that is, we can ignore the fully connected layers, as they are only needed for making guesses about the kind of object depicted in the input image. Everything that's really interesting happens in the convolutional layers.\n",
    "\n",
    "This means: We are not trying to rebuild the complete model. We will take the convolutional layers that are needed to reason about the image style and image content, and we will only work with those layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_activated_convlayer(prev, i):\n",
    "    layer_index = get_layer_by_name(conv_layers[i])\n",
    "    W, b = get_layer_params(layer_index)\n",
    "    W = tf.constant(W)\n",
    "    b = tf.constant(np.reshape(b, (b.size)))\n",
    "    conv = tf.nn.conv2d(prev, filter=W, strides=[1,1,1,1], padding='SAME') + b\n",
    "    return tf.nn.relu(conv)\n",
    "\n",
    "def create_pool_layer(prev):\n",
    "    return tf.nn.avg_pool(prev, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')\n",
    "\n",
    "def get_next_convlayer(i, prev):\n",
    "    next_i = i + 1\n",
    "    return (next_i, create_activated_convlayer(prev, i))\n",
    "\n",
    "def get_next_convlayer_name(i):\n",
    "    return conv_layers[i]\n",
    "\n",
    "def get_last_convlayer_name(i):\n",
    "    return conv_layers[i - 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_\"For image synthesis we found that replacing the\n",
    "max-pooling operation by average pooling improves the gradient flow and one obtains slightly\n",
    "more appealing results, which is why the images shown were generated with average pooling.\"_\n",
    "\n",
    "Let's keep this in mind and rebuild the model parts that we're going to need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "\n",
    "model = {}\n",
    "\n",
    "content = plt.imread(\"content.jpg\")    \n",
    "style   = plt.imread(\"style.jpg\")   \n",
    "\n",
    "\n",
    "scale_down = 5\n",
    "\n",
    "\n",
    "height = content.shape[0] // scale_down\n",
    "width = content.shape[1] // scale_down\n",
    "index = 0\n",
    "\n",
    "model['in'] = tf.Variable(np.zeros((1, height, width, 3)), dtype=tf.float32)\n",
    "index, model[get_next_convlayer_name(index - 1)] = get_next_convlayer(index, model['in'])\n",
    "index, model[get_next_convlayer_name(index - 1)] = get_next_convlayer(index, model[get_last_convlayer_name(index)])\n",
    "model['avgpool1'] = create_pool_layer(model[get_last_convlayer_name(index)])\n",
    "index, model[get_next_convlayer_name(index - 1)] = get_next_convlayer(index, model['avgpool1'])\n",
    "index, model[get_next_convlayer_name(index - 1)] = get_next_convlayer(index, model[get_last_convlayer_name(index)])\n",
    "model['avgpool2'] = create_pool_layer(model[get_last_convlayer_name(index)])\n",
    "index, model[get_next_convlayer_name(index - 1)] = get_next_convlayer(index, model['avgpool2'])\n",
    "for i in range(3):\n",
    "    index, model[get_next_convlayer_name(index - 1)] = get_next_convlayer(index, model[get_last_convlayer_name(index)])\n",
    "model['avgpool3'] = create_pool_layer(model[get_last_convlayer_name(index)])\n",
    "index, model[get_next_convlayer_name(index - 1)] = get_next_convlayer(index, model['avgpool3'])\n",
    "for i in range(3):\n",
    "    index, model[get_next_convlayer_name(index - 1)] = get_next_convlayer(index, model[get_last_convlayer_name(index)])\n",
    "model['avgpool4'] = create_pool_layer(model[get_last_convlayer_name(index)])\n",
    "index, model[get_next_convlayer_name(index - 1)] = get_next_convlayer(index, model['avgpool4'])\n",
    "for i in range(3):\n",
    "    index, model[get_next_convlayer_name(index - 1)] = get_next_convlayer(index, model[get_last_convlayer_name(index)])\n",
    "model['avgpool5'] = create_pool_layer(model[get_last_convlayer_name(index)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For working with input images, the image's mean accross each color channel has to be subtracted, so our images have the same format like those images the VGG-19 model was trained with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#mr = np.mean(content[:,:,0])\n",
    "#mg = np.mean(content[:,:,1])\n",
    "#mb = np.mean(content[:,:,2])\n",
    "#        --- (Not sure whether to use the fixed numbers or the calculated means actually) ---\n",
    "#means = np.reshape([mr, mg, mb], (1,1,3))\n",
    "means = np.reshape([116.779, 123.68, 103.939], (1,1,3))\n",
    "\n",
    "def preprocess_image(img_in):\n",
    "    img = img_in.astype(\"float32\")\n",
    "    img = imresize(img, (height, width))\n",
    "    \n",
    "    img = img - means\n",
    "    img = img[np.newaxis]\n",
    "    return means, img\n",
    "    \n",
    "def unprocess_image(img_in):\n",
    "    img = img_in\n",
    "    img = img[0]\n",
    "    img = img + means\n",
    "    img = np.clip(img, 0, 255).astype('uint8')\n",
    "    return img\n",
    "    \n",
    "    \n",
    "means_c, processed_content = preprocess_image(content)\n",
    "means_s, processed_style   = preprocess_image(style)\n",
    "unprocessed = unprocess_image(processed_content)\n",
    "unprocessed_style = unprocess_image(processed_style)\n",
    "\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.axis('off')\n",
    "plt.imshow(unprocessed.astype('uint8'))\n",
    "plt.show()\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.axis('off')\n",
    "plt.imshow(unprocessed_style.astype('uint8'))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_\"For the images […] we matched the content representation on layer ‘conv4 2’ and the\n",
    "style representations on layers ‘conv1 1’, ‘conv2 1’, ‘conv3 1’, ‘conv4 1’ and ‘conv5 1’\"_\n",
    "\n",
    "\n",
    "\n",
    "Higher convolutional layers are good at capturing the overall content of an image. The lowest layers catch simple features (like horizontal and vertical edges or curves), whereas the layers in between capture features of medium complexity (like noses and eyes or more complex style patterns). Hence, we can take a high-level layer and see how it reacts when receiving the content image. The activation values can be interpreted as a representation of the content we would like to preserve. To get a representation of the style we would like to apply, we look how different layers react when receiving the style image as an input. \n",
    "\n",
    "\n",
    "\n",
    "_You may want to play around with the layers and weights; different choices will lead to differing results._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "content_layer = 'conv4_2'\n",
    "style_layers = ['conv1_1', 'conv2_1', 'conv3_1', 'conv4_1', 'conv5_1']\n",
    "style_weights = [.8, .8, .8, .8, .8]\n",
    "\n",
    "\n",
    "sess = tf.InteractiveSession()\n",
    "sess.run(tf.global_variables_initializer())\n",
    "\n",
    "content_features = sess.run(model[content_layer], feed_dict={model['in']: processed_content})\n",
    "style_features = sess.run([model[sl] for sl in style_layers], feed_dict={model['in']: processed_style})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the content loss, we will basically compare the current activations of the content layer with the saved values it produced when receiving the actual content image. The style loss is a bit more complicated and involves the calculation of _gram matrices_. Take a look at the original paper if you are more interested in this topic.\n",
    "\n",
    "https://arxiv.org/pdf/1508.06576.pdf\n",
    "\n",
    "We also calculate the _variation loss_. This will be added to the total loss and make sure that neighboring pixels are relatively similar, so as to avoid clutter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def content_loss(p):\n",
    "    size = np.prod(content_features.shape[1:])\n",
    "    return (1 / (2.0 * size)) * tf.reduce_sum(tf.pow((p - content_features), 2)) \n",
    "    \n",
    "\n",
    "def gram_matrix(features, n, m):\n",
    "    features_t = tf.reshape(f, (m, n))\n",
    "    return tf.matmul(tf.transpose(features_t), features_t)\n",
    "\n",
    "def style_loss(a, x):\n",
    "    n = a.shape[3]\n",
    "    m = a.shape[1] * a.shape[2]\n",
    "    a_matrix = gram_matrix(a, n, m)\n",
    "    g_matrix = gram_matrix(x, n, m)\n",
    "    return (1 / (4 * n**2 * m**2)) * tf.reduce_sum(tf.pow(g_matrix - a_matrix, 2))\n",
    "\n",
    "def var_loss(x):\n",
    "    h, w = x.get_shape().as_list()[1], x.get_shape().as_list()[2]\n",
    "    dx = tf.square(x[:, :h - 1, :w - 1, :] - x[:, :h - 1, 1:, :])\n",
    "    dy = tf.square(x[:, :h - 1, :w - 1, :] - x[:, 1:, :w - 1, :])\n",
    "    return tf.reduce_sum(tf.pow(dx + dy, 1.25))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The $\\alpha$-, $\\beta$- and $\\gamma$- values are used to balance the style, content and variation loss. If your results are not satisfying, these are some values you might want to adjust."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "e = [style_loss(sf, model[ln]) for sf, ln in zip(style_features, style_layers)]\n",
    "styleloss = sum([style_weights[l] * e[l] for l in range(len(style_layers))])\n",
    "contentloss = content_loss(model[content_layer])\n",
    "varloss = var_loss(model['in'])\n",
    "\n",
    "alpha = 1\n",
    "beta = 100\n",
    "gamma = 0.1\n",
    "\n",
    "total_loss = alpha * contentloss + beta * styleloss + gamma * varloss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The authors of the original paper proposed to assign white noise to the input of our network and let the network transform this noise into a combined representation of the desired content and style. However, we can make the task of restoring the image content easier by inputting a noisy representation of our content image to the network. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise_ratio = 0.7\n",
    "content_ratio = 1. - noise_ratio\n",
    "noise = np.random.uniform(-15, 15, processed_content.shape)\n",
    "input_image = (processed_content * content_ratio) + noise_ratio * noise\n",
    "\n",
    "unp = unprocess_image(input_image)\n",
    "plt.imshow(unp.astype(\"uint8\"))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It was possible to train the network using a large learning rate, but you might want to adjust these settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = tf.train.AdamOptimizer(1).minimize(total_loss)\n",
    "sess = tf.InteractiveSession()\n",
    "sess.run(tf.global_variables_initializer())\n",
    "sess.run(model['in'].assign(input_image))\n",
    "pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the Network\n",
    "\n",
    "It takes about 100 iterations until the created image will look somewhat like a merge of the content image with the desired style. Then, the output slowly becomes more visually appealing. I stopped training before the loss had converged, as the outputs mostly looked neat enough much earlier and the training would otherwise take quite long. I run the training steps on a Tesla K80 GPU for about 15-20 minutes for most content/style combinations until I was satisfied with the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(1000):\n",
    "    if i % 50 == 0:\n",
    "        m_in = sess.run(model['in'])\n",
    "        plt.imshow(unprocess_image(m_in, means_c).astype(\"uint8\"))\n",
    "        plt.show()\n",
    "        \n",
    "    _, ls = sess.run([optimizer, total_loss])\n",
    "        \n",
    "    if i % 10 == 0:        \n",
    "        print(i, ls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}