{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GS-p39QFra42"
   },
   "source": [
    "# Style Transfer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "URWNLiNb0jHd"
   },
   "source": [
    "An example below is inspired [by original tutorial on TensorFlow](https://www.tensorflow.org/tutorials/generative/style_transfer), as well as [by this blog post](https://medium.com/mlreview/making-ai-art-with-style-transfer-using-keras-8bb5fa44b216). Another good example of Style Transfer using CNTK framework is [here](https://github.com/Microsoft/CNTK/blob/master/Tutorials/CNTK_205_Artistic_Style_Transfer.ipynb). Here is an original paper on [Artistic Style Transfer](https://arxiv.org/abs/1508.06576).\n",
    "\n",
    "Main ideas behind style transfer are the following:\n",
    "\n",
    "* Starting from white noise, we try to optimize the current image $x$ to minimize some loss function\n",
    "* Loss function consists of three components $\\mathcal{L(x)} = \\alpha\\mathcal{L}_c(x,i) + \\beta\\mathcal{L}_s(x,s)+\\gamma\\mathcal{L}_t(x)$\n",
    "   - $\\mathcal{L}_c$ - content loss - shows how close the current image $x$ is to original image $i$\n",
    "   - $\\mathcal{L}_s$ - style loss - shows how close the current image $x$ is to style image $s$\n",
    "   - $\\mathcal{L}_t$ - total variation loss (we will not consider it in our example) - makes sure that the resulting image is smooth, i.e. it shows the mean squared error of neighbouring pixels of the image $x$\n",
    "   \n",
    "Those loss functions have to be designed in a clever way, so that for example style loss corresponds to styles of the images being similar, and not the actual content. For that, we will compare some deeper feature layers of a CNN which looks at the image."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Vnz39nOYM-Vh"
   },
   "source": [
    "Let's start by loading a couple of images:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 126
    },
    "colab_type": "code",
    "id": "gAmHWlU0rZuz",
    "outputId": "919fcfe4-adb9-478b-b46d-a407383b0cb7"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "100  210k  100  210k    0     0  2670k      0 --:--:-- --:--:-- --:--:-- 2670k\n",
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "100  131k  100  131k    0     0   459k      0 --:--:-- --:--:-- --:--:--  459k\n"
     ]
    }
   ],
   "source": [
    "!mkdir -p images\n",
    "!curl https://cdn.pixabay.com/photo/2016/05/18/00/27/franz-marc-1399594_960_720.jpg > images/style.jpg\n",
    "!curl https://upload.wikimedia.org/wikipedia/commons/thumb/b/bd/Golden_tabby_and_white_kitten_n01.jpg/1280px-Golden_tabby_and_white_kitten_n01.jpg > images/image.jpg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SE-T2e0Qt5zm"
   },
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras.applications.vgg16 import preprocess_input\n",
    "import IPython.display as display"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FrdnrziqOUmV"
   },
   "source": [
    "Let's load those images and resize them to $512\\times512$. Also, we will generate resulting image `img_result` as a random array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 184
    },
    "colab_type": "code",
    "id": "9zJLOhaVtXC3",
    "outputId": "b4069219-fa54-4f50-8d3a-8e5905f2d7b5"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 424636 characters]\n",
      "text/plain": [
       "<Figure size 864x864 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "img_size = 256\n",
    "\n",
    "def load_image(fn):\n",
    "  x = cv2.imread(fn)\n",
    "  return cv2.cvtColor(x, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "img_style = load_image('images/style.jpg')\n",
    "img_content = load_image('images/image.jpg')\n",
    "img_content = img_content[:,200:200+857,:]\n",
    "img_content = cv2.resize(img_content,(img_size,img_size))\n",
    "img_style = img_style[:,200:200+671,:]\n",
    "img_style = cv2.resize(img_style,(img_size,img_size))\n",
    "\n",
    "img_result = np.random.uniform(size=(img_size,img_size,3))\n",
    "\n",
    "matplotlib.rcParams['figure.figsize'] = (12, 12)\n",
    "matplotlib.rcParams['axes.grid'] = False\n",
    "\n",
    "fig,ax = plt.subplots(1,3)\n",
    "ax[0].imshow(img_content)\n",
    "ax[1].imshow(img_style)\n",
    "ax[2].imshow((255*img_result).astype(int))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "V6uAbz-IhZDG"
   },
   "source": [
    "To calculate style loss and content loss, we need to work in the feature space extracted by a CNN. We can use different CNN architectures, but for simplicity in our case we will chose VGG-19, pre-trained on *ImageNet*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JGC48V4627vk"
   },
   "outputs": [],
   "source": [
    "vgg = tf.keras.applications.VGG16(include_top=False, weights='imagenet')\n",
    "vgg.trainable = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8D3HTgURhZDM"
   },
   "source": [
    "Let's have a look at the model architecture:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 835
    },
    "colab_type": "code",
    "id": "P7nSJfH6hZDO",
    "outputId": "f9d48267-9e63-4bd4-8684-f32878c9a092"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"vgg16\"\n",
      "_________________________________________________________________\n",
      " Layer (type)                Output Shape              Param #   \n",
      "=================================================================\n",
      " input_3 (InputLayer)        [(None, None, None, 3)]   0         \n",
      "                                                                 \n",
      " block1_conv1 (Conv2D)       (None, None, None, 64)    1792      \n",
      "                                                                 \n",
      " block1_conv2 (Conv2D)       (None, None, None, 64)    36928     \n",
      "                                                                 \n",
      " block1_pool (MaxPooling2D)  (None, None, None, 64)    0         \n",
      "                                                                 \n",
      " block2_conv1 (Conv2D)       (None, None, None, 128)   73856     \n",
      "                                                                 \n",
      " block2_conv2 (Conv2D)       (None, None, None, 128)   147584    \n",
      "                                                                 \n",
      " block2_pool (MaxPooling2D)  (None, None, None, 128)   0         \n",
      "                                                                 \n",
      " block3_conv1 (Conv2D)       (None, None, None, 256)   295168    \n",
      "                                                                 \n",
      " block3_conv2 (Conv2D)       (None, None, None, 256)   590080    \n",
      "                                                                 \n",
      " block3_conv3 (Conv2D)       (None, None, None, 256)   590080    \n",
      "                                                                 \n",
      " block3_pool (MaxPooling2D)  (None, None, None, 256)   0         \n",
      "                                                                 \n",
      " block4_conv1 (Conv2D)       (None, None, None, 512)   1180160   \n",
      "                                                                 \n",
      " block4_conv2 (Conv2D)       (None, None, None, 512)   2359808   \n",
      "                                                                 \n",
      " block4_conv3 (Conv2D)       (None, None, None, 512)   2359808   \n",
      "                                                                 \n",
      " block4_pool (MaxPooling2D)  (None, None, None, 512)   0         \n",
      "                                                                 \n",
      " block5_conv1 (Conv2D)       (None, None, None, 512)   2359808   \n",
      "                                                                 \n",
      " block5_conv2 (Conv2D)       (None, None, None, 512)   2359808   \n",
      "                                                                 \n",
      " block5_conv3 (Conv2D)       (None, None, None, 512)   2359808   \n",
      "                                                                 \n",
      " block5_pool (MaxPooling2D)  (None, None, None, 512)   0         \n",
      "                                                                 \n",
      "=================================================================\n",
      "Total params: 14,714,688\n",
      "Trainable params: 0\n",
      "Non-trainable params: 14,714,688\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "vgg.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define a function that will allow us to extract intermediate features from VGG network: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def layer_extractor(layers):\n",
    "    outputs = [vgg.get_layer(x).output for x in layers]\n",
    "    model = tf.keras.Model([vgg.input],outputs)\n",
    "    return model "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "K_fA5ue-09XG"
   },
   "source": [
    "## Content Loss\n",
    "\n",
    "**Content loss** will show how close our current image $x$ is to the original image. It looks at the intermediate feature layers in CNN, and computes square error. Content loss on layer $l$ will be defined as\n",
    "$$\n",
    "\\mathcal{L}_c = {1\\over2}\\sum_{i,j} (F_{ij}^{(l)}-P_{ij}^{(l)})^2\n",
    "$$\n",
    "where $F^{(l)}$ and $P^{(l)}$ -- features at layer $l$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XHWUFOK50ORi"
   },
   "outputs": [],
   "source": [
    "content_layers = ['block4_conv2'] \n",
    "content_extractor = layer_extractor(content_layers)\n",
    "\n",
    "content_target = content_extractor(preprocess_input(tf.expand_dims(img_content,axis=0)))\n",
    "\n",
    "def content_loss(img):\n",
    "    z = content_extractor(preprocess_input(tf.expand_dims(255*img,axis=0))) \n",
    "    return 0.5*tf.reduce_sum((z-content_target)**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lUGSMAcxkS1w"
   },
   "source": [
    "Now we will implement the main trick of style transfer - **optimization**. We will start with random image, and then use TensorFlow optimizer to adjust this image to minimize content loss.\n",
    "\n",
    "**Important**: in our case, all computations are performed using GPU-aware *TensorFlow* framework, which allows this code to run much more efficiently on GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1572756 characters]\n",
      "text/plain": [
       "<Figure size 864x864 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "img = tf.Variable(img_result)\n",
    "opt = tf.optimizers.Adam(learning_rate=0.002, beta_1=0.99, epsilon=1e-1)\n",
    "\n",
    "clip = lambda x : tf.clip_by_value(x,clip_value_min=0,clip_value_max=1)\n",
    "\n",
    "def optimize(img,loss_fn):\n",
    "    with tf.GradientTape() as tape:\n",
    "        loss = loss_fn(img)\n",
    "    grad = tape.gradient(loss,img)\n",
    "    opt.apply_gradients([(grad,img)])\n",
    "    #img.assign(tf.clip_by_value(img,clip_value_min=0,clip_value_max=1))\n",
    "\n",
    "def train(img,loss_fn,epochs=10,steps_per_epoch=100):\n",
    "    for _ in range(epochs):\n",
    "        display.clear_output(wait=True)\n",
    "        plt.imshow((255*clip(img)).numpy().astype(int))\n",
    "        plt.show()\n",
    "        for _ in range(steps_per_epoch):\n",
    "            optimize(img,loss_fn=loss_fn)\n",
    "\n",
    "train(img,content_loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Cm5c7s4N4Dj4"
   },
   "source": [
    "> **Exercise**: Try experimenting with different layers in the network and see what happens. You can also try optimizing for several layers together, but you would have to change the code for `content_loss` a bit.\n",
    "\n",
    "## Style Loss\n",
    "\n",
    "Style loss is the main idea behind Style Transfer. We compare not the actual features, but their Gram matrices, which are defined as $$G=A\\times A^T$$\n",
    "\n",
    "Gram matrix is similar to correlation matrix, and it shows how some filters depend on the others. Style Loss is computed as a sum of losses from different layers, which are often considered with weighted coefficients.\n",
    "\n",
    "Total loss function for style transfer is a sum of *content loss* and *style loss*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "F2viTufg3aOt"
   },
   "outputs": [],
   "source": [
    "def gram_matrix(x):\n",
    "  result = tf.linalg.einsum('bijc,bijd->bcd', x, x)\n",
    "  input_shape = tf.shape(x)\n",
    "  num_locations = tf.cast(input_shape[1]*input_shape[2], tf.float32)\n",
    "  return result/(num_locations)\n",
    "\n",
    "style_layers = ['block1_conv1','block2_conv1','block3_conv1','block4_conv1']\n",
    "\n",
    "def style_extractor(img):\n",
    "    return [gram_matrix(x) for x in layer_extractor(style_layers)(img)]\n",
    "\n",
    "style_target = style_extractor(preprocess_input(tf.expand_dims(img_style,axis=0)))\n",
    "\n",
    "def style_loss(img):\n",
    "    z = style_extractor(preprocess_input(tf.expand_dims(255*img,axis=0)))\n",
    "    loss = tf.add_n([tf.reduce_mean((x-target)**2) \n",
    "                           for x,target in zip(z,style_target)])\n",
    "    return loss / len(style_layers)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "n6TiEh1R3y3x"
   },
   "source": [
    "## Putting it all together\n",
    "\n",
    "We will define `total_loss` function that will calculate combined loss, and run the optimization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EXH3PUgm30nl"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1554488 characters]\n",
      "text/plain": [
       "<Figure size 864x864 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def total_loss(img):\n",
    "    return 2*content_loss(img)+style_loss(img)\n",
    "\n",
    "img.assign(img_result)\n",
    "\n",
    "train(img,loss_fn=total_loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "D5sJJZ3t304U"
   },
   "source": [
    "The code below performs the actual optimization of loss. Keep in mind that even with GPU the optimization takes significant amount of time. You can run the cell below several times to improve the result."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "95xk2ji-zGnX"
   },
   "source": [
    "## Add variation loss\n",
    "\n",
    "**Variation loss** allows us to make the image less noisy, by minimizing the amount of difference between neighbouring pixels.\n",
    "\n",
    "We will also start optimization from the original content image, which allows us to keep more content details in the image, without complicating content loss function. We will add some noise, though."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JlUBzILGYFLs"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1369912 characters]\n",
      "text/plain": [
       "<Figure size 864x864 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def variation_loss(img):\n",
    "  img = tf.cast(img,tf.float32)\n",
    "  x_var = img[ :, 1:, :] - img[ :, :-1, :]\n",
    "  y_var = img[ 1:, :, :] - img[ :-1, :, :]\n",
    "  return tf.reduce_sum(tf.abs(x_var)) + tf.reduce_sum(tf.abs(y_var))\n",
    "  \n",
    "def total_loss_var(img):\n",
    "    return content_loss(img)+150*style_loss(img)+30*variation_loss(img)\n",
    "\n",
    "img.assign(clip(np.random.normal(-0.3,0.3,size=img_content.shape)+img_content/255.0))\n",
    "\n",
    "train(img,loss_fn=total_loss_var)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "DENlUJLdNcek"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(256, 256, 3)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cv2.imwrite('result.jpg',(img.numpy()[:,:,::-1]*255))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "StyleTransfer.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "azureml_py38_PT_and_TF",
   "language": "python",
   "name": "conda-env-azureml_py38_PT_and_TF-py"
  },
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
