Neural_Style_Transfer_V2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RC5cT2j6YmaU"
   },
   "source": [
    "# Neural Style Transfer using Keras and TensorFlow\n",
    "\n",
    "A very interesting and entertaining application of machine learning is Neural Style Transfer. First described in [this paper](https://arxiv.org/abs/1508.06576) by Leon Gatys, Alexander Ecker, and Matthias Bethge, it is an algorithm that is able to create novel works of art from a simple photograph. Essentially, what it does is imitate the style of a certain painting and replicate it in a photo of your choosing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RC5cT2j6YmaU"
   },
   "source": [
    "## What this project aims to do\n",
    "\n",
    "This project aims to provide a simple and intuitive overview of how the algorithm works, based entirely on the original paper as well as Professor Andrew Ng's ideas in his [deeplearning.ai course](https://www.coursera.org/learn/convolutional-neural-networks). It also aims to provide a simple implementation of the paper, using the popular Deep Learning libraries Keras and Tensorflow. This is partly out of frustration, as I found it hard to come across a simple example that implemented this algorithm from scratch using Tensorflow. The implementation is fully documented and easily modifiable.\n",
    "\n",
    "If you are already familiar with how the algorithm works, feel free to skip ahead to the **Implementation** section. \n",
    "\n",
    "If you simply want to try generating your own images, scroll down until the end of the notebook, to **Let's Generate Some Art!**."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RC5cT2j6YmaU"
   },
   "source": [
    "## How it works\n",
    "\n",
    "So, how does an _algorithm_ generate art? Isn't that supposed to be an innately human quality?\n",
    "\n",
    "Sort of.\n",
    "\n",
    "### Defining the problem\n",
    "\n",
    "What we aim to do is this: given a content image $I_c$, we would like to generate a new image $I_g$ which replicates the style of some style image, $I_s$. For example:\n",
    "\n",
    "| Content image $I_c$ | Style image $I_s$ | Generated image $I_g$ |\n",
    "| ------------- |-------------| --------|\n",
    "| <img src=\"caracas.jpg\" width=\"300\"/> | <img src=\"starry-night.jpg\" width=\"300\"/>| <img src=\"generated.jpg\" width=\"300\" />|\n",
    "\n",
    "\n",
    "To be able to achieve this, we will be using **machine learning** and **function minimization**.\n",
    "\n",
    "### Setting up the loss\n",
    "\n",
    "What we want to do is define a function that can take as input our three images $I_c$, $I_s$, and $I_g$, and tell us how different our generated image is from our content image, as well as how different the style of our generated image is from our style image. If the number returned by our function is high, this means that the generated image is far from looking like our content and style images. On the other hand, if the number returned is low, it means that the generated image is similar in content and style to our other images, respectively.\n",
    "\n",
    "#### Representing the images\n",
    "\n",
    "To do this, we first need a way of representing these images in a useful format. We do this by representing our three images as three-dimensional matrices of shape $(height, width, channels)$ that contain numbers ranging from $0$ to $255$ (a way of representing colour in a computer). If we are working with colour images, then the $channels$ dimension will be $3$, as we will have red, green, and blue channels. If we are working with grayscale images, we should only have one channel.\n",
    "\n",
    "However, to the computer, that is just a bunch of numbers. What would be ideal is to have a representation of what the images _actually_ are. A dog? A house? A person? To do this, we will use a pre-trained state of the art **convolutional neural network**.\n",
    "\n",
    "A convolutional neural network (CNN for short) is essentially a sequence of matrix operations on an image that eventually gives us a numerical representation of what is in it. The details of how they work will not be included here for brevity, but a quick primer can be found [here](http://cs231n.github.io/convolutional-networks/). The network can have many layers, and you can think of each layer as intuitively recognizing the features present in an image, progressing from the lowest-level features (for example, where vertical and horizontal edges are located in the image) to higher-level features (elaborate combinations or patterns of those lower level features, like a dog, or a person, or a house).\n",
    "\n",
    "_An example of a CNN, taken from Andrew Ng's deeplearning.ai course. Each \"rectangle\" is a 3-D matrix of numbers, representing an intermediate layer inside the neural network._\n",
    "![An example of a CNN](https://cdn-images-1.medium.com/max/1600/1*jqKHgwZ8alM3K_JRYO_l4w.png)\n",
    "\n",
    "By passing our images as input to a CNN, we obtain several intermediate \"layers\", which are new matrices that you can intuitively think of as representations of the features present in an image. The intermediate layers in these matrices are called \"activations\", because you can think of each number in them as \"neurons\" that \"activate\" (i.e, have a high numerical value) when a specific feature is detected in that part of the image.\n",
    "\n",
    "We can then use these activations as better representations of our images! Then, we define a function that takes in these three matrices as inputs, and gives us useful information about how far off we are from obtaining our desired result. To do this, we will first define two separate functions, which we will then use to compose the final one.\n",
    "\n",
    "#### Content loss\n",
    "\n",
    "The first function, which we will call $C(I_g, I_c)$, will tell us how far off our generated image $I_g$ is from the content image $I_c$. The definition is relatively simple:\n",
    "\n",
    "$$ C(I_g, I_c) = \\sum | I_g - I_c |^2 $$\n",
    "\n",
    "We will take the absolute squared difference of both matrices, and then sum over each and every element of the resulting matrix. This will give us a final positive number, which represents how far apart the generated image $I_g$ is from the content image $I_c$ - the higher the number, the greater the difference between both matrices, and therefore the more different they are.\n",
    "\n",
    "It is possible to use both the original image matrices $I_g$ and $I_c$ here, or an intermediate representation of these images after passing them through a CNN and obtaining their activations. The latter would simply require us to pass the matrices as inputs to our CNN first, and then pick one of the intermediate layers as the representations.\n",
    "\n",
    "*_Note_*: when implementing the NST algorithm, using the squared difference between both image matrices _without_ passing them through the network first seems to work just as well. In some circumstances, the absolute squared difference alone may be faster than using the CNN and then computing the squared difference.\n",
    "\n",
    "#### Style loss\n",
    "\n",
    "The second function, which we will call $S(I_g, I_s)$, will tell us how different the styles of our generated image and our style image are. The definition of this function is a bit more complex.\n",
    "\n",
    "First, we need a way to represent the _style_ of an image, not its content. It turns out that a fairly good representation of the style of a certain image is taking the [**gram matrix**](https://en.wikipedia.org/wiki/Gramian_matrix) of each of its activation matrices after passing it through a CNN. The gram matrix of a matrix $X$ is defined as the multiplication of itself by its transpose.\n",
    "\n",
    "$$ X_{gram} = X^T X $$\n",
    "\n",
    "To take the gram matrix of the activations, we first unroll the original activation matrix from a (height, width, channels) shape into a (height * width, channels) shape, putting all of the neurons in each channel on a single row and obtaining a 2-D matrix. The resulting matrix after taking the gram matrix of the unrolled activation is one that (roughly) indicates how correlated the different activations in the image are, and how prevalent are certain features in the image. The intuition behind this is the following:\n",
    "\n",
    "- If there are certain features present in the image, then individual \"neurons\" in the activation matrix will have been \"activated\", meaning they have a relatively large value, or are at least greater than 0.\n",
    "- If we then transpose the matrix, and multiply it by itself, those \"activated\" neurons will produce large numbers if multiplied with other \"activated\" neurons (other features that have been detected in the image), and produce smaller numbers if they are multiplied with \"de-activated\" neurons (features that are not present, and thus have a value of 0 or close to 0).\n",
    "- The resulting matrix will therefore indicate whether features are ocurring together, or not ocurring together, by producing high values or low values in positions along the matrix - loosely, the correlation between different features in the image. This gives us a representation of the **style** of the image.\n",
    "\n",
    "With this representation, we can then similarly define a function to represent the difference in style as follows:\n",
    "\n",
    "1. Compute the intermediate activations corresponding to the images $I_g$ and $I_s$ by running them through the CNN.\n",
    "2. For the activations corresponding to each layer $l$, we compute the reduced sum of the squared difference of the two activations $A_{l_g}$ and $A_{l_s}$:\n",
    "$$ D_l = \\sum | A_{l_g} - A_{l_s} |^2 $$\n",
    "\n",
    "3. We then sum up all the squared differences computed for each layer.\n",
    "\n",
    "Thus, our function is essentially\n",
    "\n",
    "$$ S(I_g, I_s) = \\sum_{l = 0}^{n} D_l(A_{l_g}, A_{l_s}) $$\n",
    "\n",
    "Where $l$ represents each layer of the CNN, $n$ is the total number of layers chosen, and $D_l$ is the function described above.\n",
    "\n",
    "In practice, it is useful to pick layers spanning the whole length of the CNN, due to the nature of the features that each layer detects. The first layers typically detect low-level features, such as brush strokes, whilst the later layers detect higher-level features, like shapes and patterns - thus, by computing the squared differences across all of these layers, we can identify style differences not only in low-level texture details, but also in higher-level patterns.\n",
    "\n",
    "\n",
    "#### Overall loss\n",
    "\n",
    "We can now define the overall \"loss\" of our three images - how they differ in content and style. The overall function is defined as the sum of the functions defined previously:\n",
    "\n",
    "$$ J(I_g, I_c, I_s) = (\\alpha * C(I_g, I_s)) + (\\beta * S(I_g, I_s)) $$\n",
    "\n",
    "Where the constants $\\alpha$ and $\\beta$ are weightings for each respective loss. If we want the generated image to give more importance to the style over representing the original content truthfully, we would give $\\beta$ a higher value than $\\alpha$, for example.\n",
    "\n",
    "### Minimizing our loss function\n",
    "\n",
    "Now that we have a function $J$ that tells us how different in style and content our generated image is from the target images (the bigger the output of the function, the more different in style and content is our generated image, and vice-versa), we have to minimize it: attempt to find a value for the individual pixels of our image $I_g$ for which our loss function $J$ gives the smallest value. How do we do this?\n",
    "\n",
    "The essential idea to minimize our function is described by the gradient descent algorithm:\n",
    "\n",
    "1. Differentiate our function $J$ with respect to each number in our image matrix $I_g$ to obtain a gradient. Think of each number in our matrix as an individual variable.\n",
    "2. Once we obtain a new matrix containing all of our respective partial gradients, we update each component in our image matrix by the negative of its corresponding gradient multiplied by a constant called the **learning rate**. For instance, once we differentiate each variable in our image matrix to obtain a matrix of gradients $G$, we update our image matrix with the following rule:\n",
    "$$ I_g := I_g - (r * G) $$\n",
    "where $r$ is our learning rate.\n",
    "3. Continuously carry out this iterative process until we arrive at a good enough minimum for our function $J$.\n",
    "\n",
    "This is a _very_ simplified explanation of gradient descent, for the sake of brevity. A more thorough explanation can be found [here](https://towardsdatascience.com/gradient-descent-simply-explained-1d2baa65c757).\n",
    "\n",
    "We are now ready to start implementing this algorithm!\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RC5cT2j6YmaU"
   },
   "source": [
    "--------\n",
    "\n",
    "## Implementation\n",
    "\n",
    "For the implementation, we will be using the Keras library to load the pre-trained CNN, and the Tensorflow library to define and minimize our loss functions. Keras comes pre-packaged in Tensorflow.\n",
    "\n",
    "### Setup\n",
    "\n",
    "First, we will import everything we need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CMdVxdPpYmaV"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.pyplot import imshow\n",
    "import numpy as np\n",
    "import scipy.misc\n",
    "import tensorflow as tf\n",
    "import tensorflow.contrib.eager as tfe\n",
    "import os\n",
    "\n",
    "# tell Jupyter to display all our plots inline in the notebook\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Eager execution\n",
    "\n",
    "As a start, we enable TensorFlow's Eager Execution mode. This will allow us to evaluate the value of tensors instantly instead of having to build a computational graph and run it. This mode is particularly useful to us because it allows us to seamlessly work with Keras models out of the box in our loss functions, and saves us the pain of managing multiple graphs and sessions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.enable_eager_execution()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EtNBqQYeYmaY"
   },
   "source": [
    "Now that we've imported everything, we can start defining the functions and variables that we will use (described in the previous section). For our pre-trained CNN, we'll be using the [VGG19 network](http://www.robots.ox.ac.uk/~vgg/research/very_deep/), made by Oxford University's Visual Geometry Group. The network is state of the art, and performs very well on the ImageNet standard benchmark. The paper can be found [here](https://arxiv.org/abs/1409.1556).\n",
    "\n",
    "The functions are:\n",
    "\n",
    "- `compute_total_cost_vgg19`: the function that we will be minimizing in order to generate our image. Computes the overall \"cost\", or \"loss\", of our model.\n",
    "- `get_vgg19_model`: returns the VGG19 architechture as a Keras Model, with preloaded weights trained on ImageNet.\n",
    "- `compute_vgg19_gradients`: does exactly what it says - computes the gradients of our loss function (which uses the VGG19 model, hence the name) which we then use to optimize our loss function. We explicitly compute the gradients this way in order to be able to apply gradient clipping if necessary.\n",
    "\n",
    "And our variable for image size:\n",
    "\n",
    "- `DEFAULT_INPUT_SIZE`: the size of the images we will be feeding in, as well as the size of the image generated.\n",
    "- `STYLE_LAYERS_VGG19`: the layers of the VGG19 net that we'll be using to compute our style loss function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yYuIVhQGgeBA"
   },
   "outputs": [],
   "source": [
    "# The size of the input images, as well as the image generated.\n",
    "# Increasing the height and width will increase the resolution of the generated image,\n",
    "# but it will take more time to produce.\n",
    "DEFAULT_INPUT_SIZE = (400, 400, 3)\n",
    "\n",
    "# The layers from the VGG net that we'll use to compute the style loss.\n",
    "# We also apply a corresponding weighting, if we want to give more importance to certain layers.\n",
    "# An equal weighting works well, but others can be tried.\n",
    "STYLE_LAYERS_VGG19 = [\n",
    "    ('block1_conv1', 0.2),\n",
    "    ('block2_conv1', 0.2),\n",
    "    ('block3_conv1', 0.2),\n",
    "    ('block4_conv1', 0.2),\n",
    "    ('block5_conv1', 0.2)\n",
    "]\n",
    "\n",
    "def get_vgg19_model():\n",
    "    \"\"\"\n",
    "    Fetches the VGG19 model ready for use in our loss functions.\n",
    "    \n",
    "    Returns:\n",
    "    tf.keras.models.Model() -- A Keras model that takes in an image as input, and returns the layers that\n",
    "                              were defined in the STYLE_LAYERS_VGG19 constant as output.\n",
    "    \"\"\"\n",
    "    vgg19 = tf.keras.applications.vgg19.VGG19(input_shape=DEFAULT_INPUT_SIZE, include_top=False, weights=\"imagenet\")\n",
    "    vgg19.trainable = False\n",
    "    style_outputs = [vgg19.get_layer(name).output for name, weight in STYLE_LAYERS_VGG19]\n",
    "    return tf.keras.models.Model(vgg19.input, style_outputs)\n",
    "\n",
    "def compute_content_cost(a_C, a_G, model):\n",
    "    \"\"\"\n",
    "    Computes the content cost.\n",
    "\n",
    "    Arguments:\n",
    "    a_C -- tensor of dimension (1, n_H, n_W, n_C) representing the image C (original image).\n",
    "    a_G -- tensor of dimension (1, n_H, n_W, n_C) representing the image G (generated image).\n",
    "    model -- Keras model that outputs the intermediate layer activations.\n",
    "\n",
    "    Returns:\n",
    "    J_content -- scalar representing the cost. The number is higher if the images are different.\n",
    "    \"\"\"\n",
    "    # run through CNN. We'll be using the fourth layer to compute the content cost.\n",
    "    a_C_act, a_G_act = model(a_C)[3], model(a_G)[3]\n",
    "    \n",
    "    # get dimensions\n",
    "    num_batch, height, width, channels = a_C_act.numpy().shape\n",
    "\n",
    "    # unroll into a 2D matrix, with each channel taking up an entire row.\n",
    "    # We don't really need it here, but we will need it to compute the gram matrices, so it's only logical.\n",
    "    a_C_unroll = tf.reshape(tf.transpose(a_C_act, perm=[0,3,2,1]), shape=[height * width, channels, num_batch])\n",
    "    a_G_unroll = tf.reshape(tf.transpose(a_G_act, perm=[0,3,2,1]), shape=[height * width, channels, num_batch])\n",
    "\n",
    "    # compute the sum and squared difference over the resulting matrix.\n",
    "    # we're using a regularizing constant, but this shouldn't really matter that much.\n",
    "    J_content = (1.0 / (4 * height * width * channels)) * tf.reduce_sum(tf.squared_difference(a_C_unroll, a_G_unroll))\n",
    "\n",
    "    return J_content\n",
    "\n",
    "def gram_matrix(A):\n",
    "    \"\"\"\n",
    "    Computes the \"style\" matrix of an image.\n",
    "\n",
    "    Argument:\n",
    "    A -- matrix of shape (n_C, n_H*n_W)\n",
    "\n",
    "    Returns:\n",
    "    gram -- Gram matrix of A, of shape (n_C, n_C)\n",
    "    \"\"\"\n",
    "    gram = tf.matmul(A, tf.transpose(A))\n",
    "    return gram\n",
    "\n",
    "def compute_layer_style_cost(a_S, a_G):\n",
    "    \"\"\"\n",
    "    Computes the style cost, i.e how similar the styles of a certain layer\n",
    "    of the images are. This is done by finding the correlation between different channels\n",
    "    in the activations of the image.\n",
    "\n",
    "    Arguments:\n",
    "    a_S -- tensor of dimension (1, n_H, n_W, n_C), hidden layer activations representing style of the image S\n",
    "    a_G -- tensor of dimension (1, n_H, n_W, n_C), hidden layer activations representing style of the image G\n",
    "\n",
    "    Returns:\n",
    "    cost -- scalar value, style cost.\n",
    "    \"\"\"\n",
    "    # get dimensions\n",
    "    num_batch, height, width, channels = a_S.numpy().shape\n",
    "\n",
    "    # unroll into a 2D matrix, where each channel takes up a whole row.\n",
    "    # this enables us to compute the gram matrices and get a representation of the style of the images.\n",
    "    a_S_unroll = tf.reshape(tf.transpose(a_S, perm=[3,2,1,0]), shape=[channels, width * height])\n",
    "    a_G_unroll = tf.reshape(tf.transpose(a_G, perm=[3,2,1,0]), shape=[channels, width * height])\n",
    "\n",
    "    # compute gram matrices\n",
    "    a_S_gram = gram_matrix(a_S_unroll)\n",
    "    a_G_gram = gram_matrix(a_G_unroll)\n",
    "\n",
    "    # regularization constant\n",
    "    reg_const = (1.0 / (4 * ((height * width) ** 2) * (channels ** 2)))\n",
    "    \n",
    "    # compute the style cost for the specific layer\n",
    "    cost = reg_const * tf.reduce_sum(tf.squared_difference(a_S_gram, a_G_gram))\n",
    "\n",
    "    return cost\n",
    "\n",
    "\n",
    "def compute_style_cost(a_G, a_S, model, STYLE_LAYERS, verbosity=0):\n",
    "    \"\"\"\n",
    "    Computes the overall style cost from several chosen layers\n",
    "\n",
    "    Arguments:\n",
    "    a_G -- Tensor of dimension (1, n_H, n_W, n_C) representing our generated image.\n",
    "    a_S -- Tensor of dimension (1, n_H, n_W, n_C) representing our style image.\n",
    "    model -- our Keras model of the CNN.\n",
    "    STYLE_LAYERS -- A python list containing:\n",
    "                        - the names of the layers we would like to extract style from\n",
    "                        - a weighting coefficient for each of them\n",
    "    verbosity -- Flag that indicates whether we should print each layer's style cost.\n",
    "\n",
    "    Returns:\n",
    "    J_style -- scalar, sum over the style costs of each layer\n",
    "    \n",
    "    \"\"\"\n",
    "    # initialize the overall style cost\n",
    "    J_style = 0.\n",
    "\n",
    "    # get our intermediate layer activations\n",
    "    a_G_acts = model(a_G)\n",
    "    a_S_acts = model(a_S)\n",
    "\n",
    "    # loop over the layers we want to use for the cost\n",
    "    for index, output in enumerate(a_G_acts):\n",
    "        \n",
    "        # get name and relative weighting of layer\n",
    "        name, weight = STYLE_LAYERS[index]\n",
    "        \n",
    "        # get the corresponding activations for each image for the current layer\n",
    "        a_G_act = output\n",
    "        a_S_act = a_S_acts[index]\n",
    "\n",
    "        # Compute style_cost for the current layer\n",
    "        J_style_layer = compute_layer_style_cost(a_S_act, a_G_act)\n",
    "\n",
    "        if verbosity >= 2:\n",
    "            print(\"Layer {}, cost {}\".format(name, J_style_layer))\n",
    "\n",
    "        # Add coeff * J_style_layer of this layer to overall style cost\n",
    "        J_style += weight * J_style_layer\n",
    "\n",
    "    return J_style\n",
    "\n",
    "\n",
    "def compute_total_cost_vgg19(a_G, a_C, a_S, model, alpha = 10, beta = 40, verbosity=0):\n",
    "    \"\"\"\n",
    "    Computes the overall cost, or loss, of our image.\n",
    "    \n",
    "    Arguments:\n",
    "    a_G -- Tensor of shape (1, n_H, n_W, n_C) representing our generated image.\n",
    "    a_C -- Tensor of shape (1, n_H, n_W, n_C) representing our content image.\n",
    "    a_S -- Tensor of shape (1, n_H, n_W, n_C) representing our style image.\n",
    "    model -- A Keras CNN model.\n",
    "    alpha -- The weighting for the content cost.\n",
    "    beta -- The weighting for the style cost.\n",
    "    verbosity -- Flag to indicate logging level (0, 1, or 2).\n",
    "    \n",
    "    Returns:\n",
    "    cost -- Scalar with overall cost/loss.\n",
    "    \"\"\"\n",
    "\n",
    "    content_cost = compute_content_cost(a_C, a_G, model)\n",
    "    style_cost = compute_style_cost(a_G, a_S, model, STYLE_LAYERS_VGG19, verbosity=verbosity)\n",
    "    cost = (alpha * content_cost) + (beta * style_cost)\n",
    "\n",
    "    if verbosity >= 1:\n",
    "        print(\"Content cost -- {}; Style cost -- {}\".format(content_cost, style_cost))\n",
    "\n",
    "    return cost\n",
    "\n",
    "def compute_vgg19_gradients(a_G, a_C, a_S, model, alpha = 10, beta = 40):\n",
    "    \"\"\"\n",
    "    Computes the gradients of a_G relative to our overall loss function. Takes in the same parameters as compute_total_cost_vgg19\n",
    "    \"\"\"\n",
    "    # use Tensorflow's gradient tape feature, that allows us to \"record\" certain operations on a tape\n",
    "    # and then \"rewind\" it to backpropagate and obtain gradients\n",
    "    with tf.GradientTape() as tape:\n",
    "        cost = compute_total_cost_vgg19(a_G, a_C, a_S, model, alpha = alpha, beta = beta)\n",
    "    # return the gradients with respect to a_G\n",
    "    return tape.gradient(cost, a_G)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "x0LvY8kYYmac"
   },
   "source": [
    "### Preprocessing our data\n",
    "\n",
    "The VGG19 paper states that their model uses certain preprocessing before feeding images into the network. For our model to work, we must also carry out this preprocessing.\n",
    "\n",
    "Furthermore, if we want our generated image to be visually pleasing as well, we also have to deprocess it when we save it, carrying out the reverse process. Hence, we define two helper functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NBofXg-IYmac"
   },
   "outputs": [],
   "source": [
    "def vgg19_preprocess(img):\n",
    "    \"\"\"\n",
    "    Preprocesses the image matrix _img_ as was done in the VGG paper. Means are subtracted from each pixel corresponding\n",
    "    to their channels.\n",
    "    \"\"\"\n",
    "    MEANS = np.array([123.68, 116.779, 103.939]).reshape((1,1,1,3)).astype(np.float32)\n",
    "    # Substract the mean to match the expected input of VGG19\n",
    "    img = img - MEANS\n",
    "\n",
    "    return img\n",
    "\n",
    "\n",
    "def vgg19_deprocess(img):\n",
    "    \"\"\"\n",
    "    Deprocesses the image matrix _img_. Means are added to each pixel corresponding\n",
    "    to their channels.\n",
    "    \"\"\"\n",
    "    MEANS = np.array([123.68, 116.779, 103.939]).reshape((1,1,1,3)).astype(np.float32)\n",
    "    # Substract the mean to match the expected input of VGG19\n",
    "    img = img + MEANS\n",
    "\n",
    "    return img"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "by_QEoy_Ymae"
   },
   "source": [
    "We also define some utility functions for when we save and load the images we'll use. We don't have to worry much about them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eqTzQYt8Ymaf"
   },
   "outputs": [],
   "source": [
    "def checkdir(path):\n",
    "    # If the path exists, do nothing. If not, create a directory to store the image.\n",
    "    if os.path.isdir(path):\n",
    "        return 0\n",
    "    else:\n",
    "        os.mkdir(path)\n",
    "\n",
    "\n",
    "def save_image(raw_image, path, name, epoch, deprocessing=None):\n",
    "    \"\"\"\n",
    "    Saves the image _raw_image_ (not been deprocessed yet) to the directory _path_, with the name \n",
    "    _name_. The deprocessing argument accepts a function to use for deprocessing the image.\n",
    "    \"\"\"\n",
    "    # use the specified deprocessing function when saving our image.\n",
    "    # if there is no deprocessing function, don't do anything to it.\n",
    "    if deprocessing is not None:\n",
    "        raw_image = deprocessing(raw_image)\n",
    "    \n",
    "    # reshape from (1, DEFAULT_INPUT_SIZE) to (DEFAULT_INPUT_SIZE) for saving\n",
    "    image = raw_image.numpy().reshape(DEFAULT_INPUT_SIZE)\n",
    "\n",
    "    # check directory exists\n",
    "    checkdir(path)\n",
    "\n",
    "    # save our image\n",
    "    scipy.misc.imsave(\"{}/{}_at_epoch_{}.jpg\".format(path, name, str(epoch)), image)\n",
    "    \n",
    "    \n",
    "def load_image(path, preprocessing=None, print_image=False):\n",
    "    \"\"\"\n",
    "    Loads the image and converts it to a Numpy array of shape (1, DEFAULT_INPUT_SIZE). Also carries out\n",
    "    the preprocessing necessary to run through the net, according to the preprocessing argument which accepts a function.\n",
    "    Can also print the image after to confirm that it has been done correctly if print_image is True.\n",
    "    \"\"\"\n",
    "    # read image and resize it to the size we need\n",
    "    image = scipy.misc.imread(path)\n",
    "    image = scipy.misc.imresize(image, size=(DEFAULT_INPUT_SIZE[0], DEFAULT_INPUT_SIZE[1]))\n",
    "    \n",
    "    # reshape into the shape our net accepts\n",
    "    image = image.reshape((1, ) + DEFAULT_INPUT_SIZE).astype(np.float32)\n",
    "\n",
    "    if print_image:\n",
    "        imshow(image.reshape(DEFAULT_INPUT_SIZE))\n",
    "        plt.show()\n",
    "        \n",
    "    # carry out preprocessing\n",
    "    image = preprocessing(image)\n",
    "\n",
    "    return image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lGpgbd70Ymah"
   },
   "source": [
    "### Initializing our image\n",
    "\n",
    "To initialize the image we will be minimizing our cost function over, we use the content image (the image that will be re-painted in the style of another) to generate a new image that is not so far off from the original one - this way, our algorithm can converge faster and not worry too much about replicating the content in the photo as well as applying a new style. We do this by mixing in some random noise to the original image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XuP-Ibd6Ymah"
   },
   "outputs": [],
   "source": [
    "def create_noise_image(content_image, noise_ratio):\n",
    "    \"\"\"\n",
    "    Creates a noise image from the content_image\n",
    "    \n",
    "    Arguments:\n",
    "    content_image -- Tensor of shape (1, n_H, n_W, n_C) representing the ALREADY PRE-PROCESSED image.\n",
    "    noise_ratio -- The degree of noise to add to the image, ranging from 0 to 1. \n",
    "    \n",
    "    Returns:\n",
    "    input_image -- A tfe.Variable containing our generated image.\n",
    "    \"\"\"\n",
    "    # Generate a random noise_image\n",
    "    noise_image = np.random.uniform(-20, 20, (1, ) + DEFAULT_INPUT_SIZE).astype('float32')\n",
    "\n",
    "    # Set the input_image to be a weighted average of the content_image and a noise_image\n",
    "    input_image = (noise_image * noise_ratio) + (content_image * (1 - noise_ratio))\n",
    "    \n",
    "    # create our TensorFlow variable that will be optimized during each step\n",
    "    input_image = tfe.Variable(input_image, name=\"generated_image\", dtype=tf.float32)\n",
    "\n",
    "    return input_image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ePjPwfiVYmaj"
   },
   "source": [
    "### Training\n",
    "\n",
    "We are finally ready to train our model! Our train function will iteratively use the [Adam optimizer](https://arxiv.org/abs/1412.6980) (a fancy version of gradient descent that converges faster) to minimize the cost function defined above, and slowly adjust the pixel values in our generated image to create our art. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Em7duuQEYmal"
   },
   "outputs": [],
   "source": [
    "def train(content_image, style_image, model=None, deprocessing=None, cost_func=compute_total_cost_vgg19, epochs=1500, learning_rate=2, save_every=20, save_dir=\"\", filename=\"\", verbosity=0):\n",
    "    \"\"\"\n",
    "    Performs gradient descent on our loss function and optimizes our generated image variable\n",
    "    \n",
    "    Arguments:\n",
    "    content_image -- Tensor representing our content image, shape (1, n_H, n_W, n_C). Already preprocessed.\n",
    "    style_image -- Tensor representing our style image, shape (1, n_H, n_W, n_C). Already preprocessed.\n",
    "    model -- the Keras CNN model to use in our cost function.\n",
    "    deprocessing -- The function used to de-process an image before saving it.\n",
    "    cost_func -- The cost function to minimize.\n",
    "    epochs -- Number of iterations to run the algorithm for. A number between 1500 - 2000 works well.\n",
    "    learning_rate -- Learning rate for the Adam optimizer. 1.2 yields good results, 2 converges faster (but may overshoot towards the end).\n",
    "    save_every -- number that indicates the steps to progressively save the generated image. For example, =20 means\n",
    "                  the generated image will be saved to disk every 20 iterations.\n",
    "    save_dir -- Directory to save the generated images.\n",
    "    filename -- the filename to save the generated image as. Will have a \"at_epoch_xx.jpg\" appended at the end.\n",
    "    verbosity -- Indicates the level of logging output. Can be 0, 1, or 2.\n",
    "    \"\"\"\n",
    "    # assign our content and style images to representative variables.\n",
    "    a_C, a_S = content_image, style_image\n",
    "\n",
    "    # create our generated image\n",
    "    a_G = create_noise_image(a_C, 0.6)\n",
    "\n",
    "    # we will be using the Adam optimizer to minimize our function\n",
    "    optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n",
    "\n",
    "    # iterate over the number of epochs\n",
    "    for epoch in range(epochs):\n",
    "\n",
    "        # compute our gradients with respect to the a_G variable.\n",
    "        grads = compute_vgg19_gradients(a_G, a_C, a_S, model)\n",
    "        \n",
    "        # uncomment this line and comment the subsequent one if you wish to use gradient clipping.\n",
    "        #capped_grads = tf.clip_by_value(grads, -20., 20.)\n",
    "        capped_grads = grads\n",
    "        \n",
    "        # apply gradients and perform one step of minimization to our image.\n",
    "        optimizer.apply_gradients([(capped_grads, a_G)])\n",
    "\n",
    "        # print our loss\n",
    "        print(\"Loss: {} --- Step {}\".format(cost_func(a_G, a_C, a_S, model, verbosity=verbosity), epoch))\n",
    "\n",
    "        # every save_every epochs, we will save the current image in the save_dir directory using the filename \"filename\".\n",
    "        if epoch % save_every == 0:\n",
    "            save_image(a_G, save_dir, filename, epoch, deprocessing=deprocessing)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rPevrKTwYmap"
   },
   "source": [
    "## Let's generate some art!\n",
    "\n",
    "Now, we just need to load the images we want to use and provide the adequate starting configuration.\n",
    "\n",
    "Modify the code below to select the images you wish to use, and then run!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 62917
    },
    "colab_type": "code",
    "id": "E40UNZpfYmap",
    "outputId": "cf5e905a-74e8-4fbe-d700-d83469a30b0d"
   },
   "outputs": [],
   "source": [
    "\n",
    "# START HERE\n",
    "# replace these constants with the parameters you want to use\n",
    "\n",
    "# filename of the content image. Make sure it is in the same directory as this notebook.\n",
    "CONTENT_IMAGE = \"blue-marble.jpg\"\n",
    "# style image.\n",
    "STYLE_IMAGE = \"starry-night.jpg\"\n",
    "# directory where program will save the generated images.\n",
    "SAVE_DIR = \"blue_marble_vg\"\n",
    "# filename to give the generated images\n",
    "FILENAME = \"blue_marble_vg\"\n",
    "\n",
    "\n",
    "content_img = load_image(CONTENT_IMAGE, print_image=False, preprocessing=vgg19_preprocess)\n",
    "style_img = load_image(STYLE_IMAGE, print_image=False, preprocessing=vgg19_preprocess)\n",
    "\n",
    "# initialize model\n",
    "model = get_vgg19_model()\n",
    "\n",
    "# making sure we don't modify the weights in our model\n",
    "for layer in model.layers:\n",
    "    layer.trainable = False\n",
    "\n",
    "# train! \n",
    "# you can try playing around with the \n",
    "train(content_img, \n",
    "      style_img, \n",
    "      model=model, \n",
    "      cost_func=compute_total_cost_vgg19, \n",
    "      deprocessing=vgg19_deprocess,\n",
    "      save_every=20,\n",
    "      save_dir=SAVE_DIR, \n",
    "      filename=FILENAME, \n",
    "      learning_rate=1.2,\n",
    "      verbosity=1,\n",
    "      epochs=1850\n",
    "     )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nuN7snImYmau"
   },
   "source": [
    "## Running in Google Colaboratory\n",
    "\n",
    "The ideal way to run this notebook is using Google's Colaboratory, which provides a free-to-use GPU that can _greatly_ accelerate training. Go to [Google Colab](https://colab.research.google.com/) and select \"Upload a notebook\", and then select this .ipynb file from your computer. \n",
    "\n",
    "Once you've loaded it, select \"Change runtime type\" from the \"Runtime\" menu, and click on the option to use a GPU. The availability of the GPU may vary, but you should eventually manage to connect to it. After this, make sure you upload the image files that you want to use to Google Colab using the \"Files\" tab in the sidebar menu (you might have to open the menu first). After that, enjoy!"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Neural Style Transfer.ipynb",
   "provenance": [],
   "toc_visible": true,
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}