{
  "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 this blog post](https://medium.com/mlreview/making-ai-art-with-style-transfer-using-keras-8bb5fa44b216), and a lot of code is borrowed from there. 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": "stderr",
          "output_type": "stream",
          "text": [
            "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
            "                                 Dload  Upload   Total   Spent    Left  Speed\n",
            "\n",
            "  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0\n",
            "  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0\n",
            "100  210k  100  210k    0     0   208k      0  0:00:01  0:00:01 --:--:--  208k\n",
            "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
            "                                 Dload  Upload   Total   Spent    Left  Speed\n",
            "\n",
            "  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0\n",
            "100  131k  100  131k    0     0   184k      0 --:--:-- --:--:-- --:--:--  185k\n"
          ]
        }
      ],
      "source": [
        "!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": 2,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "SE-T2e0Qt5zm"
      },
      "outputs": [],
      "source": [
        "import cv2\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from scipy.optimize import fmin_l_bfgs_b"
      ]
    },
    {
      "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": 5,
      "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, 106148 characters]",
            "text/plain": [
              "<Figure size 432x288 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.randint(256,size=(img_size,img_size,3)).astype(np.float64)\n",
        "\n",
        "fig,ax = plt.subplots(1,3)\n",
        "ax[0].imshow(img_content)\n",
        "ax[1].imshow(img_style)\n",
        "ax[2].imshow(img_result.astype(np.int))\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "colab_type": "code",
        "id": "cUGN26f5tafW",
        "outputId": "1b06890c-1ad2-41b7-cfb2-106d7fad7fbd"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "from tf.keras import backend as K\n",
        "from tf.keras.applications.vgg16 import preprocess_input\n",
        "from tf.keras.applications import VGG16\n",
        "from tf.keras.preprocessing.image import load_img, img_to_array\n",
        "tf_session = K.get_session()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "G5gI9Y34z9LE"
      },
      "outputs": [],
      "source": [
        "img_content_var = K.variable(preprocess_input(np.expand_dims(img_content, axis=0)), dtype='float32')\n",
        "img_style_var = K.variable(preprocess_input(np.expand_dims(img_style, axis=0)), dtype='float32')\n",
        "img_result_1 = preprocess_input(np.expand_dims(img_result, axis=0))\n",
        "img_result_holder = K.placeholder(shape=(1, img_size, img_size, 3))"
      ]
    },
    {
      "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-16, pre-trained on *ImageNet*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "JGC48V4627vk"
      },
      "outputs": [],
      "source": [
        "cModel = VGG16(include_top=False, weights='imagenet', input_tensor=img_content_var)\n",
        "sModel = VGG16(include_top=False, weights='imagenet', input_tensor=img_style_var)\n",
        "gModel = VGG16(include_top=False, weights='imagenet', input_tensor=img_result_holder)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8D3HTgURhZDM"
      },
      "source": [
        "Let's have a look at the model architecture:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "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": [
            "_________________________________________________________________\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",
            "Total params: 14,714,688\n",
            "Trainable params: 14,714,688\n",
            "Non-trainable params: 0\n",
            "_________________________________________________________________\n"
          ]
        }
      ],
      "source": [
        "gModel.summary()"
      ]
    },
    {
      "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": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "XHWUFOK50ORi"
      },
      "outputs": [],
      "source": [
        "def get_feature_reps(x, layer_names, model):\n",
        "    featMatrices = []\n",
        "    for ln in layer_names:\n",
        "        selectedLayer = model.get_layer(ln)\n",
        "        featRaw = selectedLayer.output\n",
        "        featRawShape = K.shape(featRaw).eval(session=tf_session)\n",
        "        N_l = featRawShape[-1]\n",
        "        M_l = featRawShape[1]*featRawShape[2]\n",
        "        featMatrix = K.reshape(featRaw, (M_l, N_l))\n",
        "        featMatrix = K.transpose(featMatrix)\n",
        "        featMatrices.append(featMatrix)\n",
        "    return featMatrices\n",
        "\n",
        "def get_content_loss(F, P):\n",
        "    cLoss = 0.5*K.sum(K.square(F - P))\n",
        "    return cLoss"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "lUGSMAcxkS1w"
      },
      "source": [
        "Let's see how features at different layers affect the image. To do that, we will try to minimize content loss for just one layer. We will use `fmin_l_bfgs_b` function from SciPy, which take the function to be minimized and its gradient (in our case, `help_loss` function returns both the loss function and the gradient).\n",
        "\n",
        "**Important**: in our case, all computations are performed using GPU-aware *TensorFlow* framework. `helper_loss` function returns computational graph, which can be used to compute the loss for a given image, and it uses `K.gradients` to automatically compute gradients."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "j8N5abBPkfSt"
      },
      "outputs": [],
      "source": [
        "layer='block4_conv2'\n",
        "P = get_feature_reps(x=img_content_var, layer_names=[layer], model=cModel)[0]\n",
        "x = img_result.flatten()\n",
        "def helper_loss(img):\n",
        "    R = get_feature_reps(img,[layer],gModel)[0]\n",
        "    return get_content_loss(R,P)\n",
        "  \n",
        "def help_loss(x):\n",
        "    if x.shape != (1, img_size, img_size, 3):\n",
        "        x = x.reshape((1,img_size, img_size, 3))\n",
        "    keras_fcn = K.function([gModel.input], [helper_loss(gModel.input)])\n",
        "    keras_grad = K.function([gModel.input], K.gradients(helper_loss(gModel.input),[gModel.input]))\n",
        "    return keras_fcn([x])[0].astype('float64'),keras_grad([x])[0].flatten().astype('float64')\n",
        "\n",
        "x, _, _ = fmin_l_bfgs_b(help_loss, x, maxiter=30, disp=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "_Jl8BDzvxpf-"
      },
      "outputs": [],
      "source": [
        "def postprocess_array(x):\n",
        "    # Zero-center by mean pixel\n",
        "    if x.shape != (img_size, img_size, 3):\n",
        "        x = x.reshape((img_size, img_size, 3))\n",
        "    x[..., 0] += 103.939\n",
        "    x[..., 1] += 116.779\n",
        "    x[..., 2] += 123.68\n",
        "    # 'BGR'->'RGB'\n",
        "    x = x[..., ::-1]\n",
        "    x = np.clip(x, 0, 255)\n",
        "    x = x.astype('uint8')\n",
        "    return x\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 350
        },
        "colab_type": "code",
        "id": "muT_VWKbtehI",
        "outputId": "c94d654b-a86a-40d0-ddb0-cb3618dbe7b5"
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 319680 characters]",
            "text/plain": [
              "<matplotlib.figure.Figure at 0x7fb265464b00>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "plt.imshow(postprocess_array(x.copy()))\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 350
        },
        "colab_type": "code",
        "id": "JPMe1PHvujyp",
        "outputId": "b470c3d7-6c04-4dc6-c577-474b005be8dc"
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 320836 characters]",
            "text/plain": [
              "<matplotlib.figure.Figure at 0x7fb2620b26d8>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "layer='block3_conv2'\n",
        "P = get_feature_reps(x=img_content_var, layer_names=[layer], model=cModel)[0]\n",
        "x = img_result.flatten()\n",
        "x, _, _ = fmin_l_bfgs_b(help_loss, x, maxiter=30, disp=True)\n",
        "plt.imshow(postprocess_array(x.copy()))\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 350
        },
        "colab_type": "code",
        "id": "tI5yUDXUxBYL",
        "outputId": "5026c6c1-5f3f-434a-a187-a9ba7ed71a93"
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 318384 characters]",
            "text/plain": [
              "<matplotlib.figure.Figure at 0x7fb2620b2908>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "layer='block5_conv1'\n",
        "P = get_feature_reps(x=img_content_var, layer_names=[layer], model=cModel)[0]\n",
        "x = img_result.flatten()\n",
        "x, _, _ = fmin_l_bfgs_b(help_loss, x, maxiter=30, disp=True)\n",
        "plt.imshow(postprocess_array(x.copy()))\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "R1KKFgZ43YyN"
      },
      "source": [
        "## Style Loss\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Cm5c7s4N4Dj4"
      },
      "source": [
        "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": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "F2viTufg3aOt"
      },
      "outputs": [],
      "source": [
        "def get_Gram_matrix(F):\n",
        "    G = K.dot(F, K.transpose(F))\n",
        "    return G\n",
        "\n",
        "def get_style_loss(ws, Gs, As):\n",
        "    sLoss = K.variable(0.)\n",
        "    for w, G, A in zip(ws, Gs, As):\n",
        "        M_l = K.int_shape(G)[1]\n",
        "        N_l = K.int_shape(G)[0]\n",
        "        G_gram = get_Gram_matrix(G)\n",
        "        A_gram = get_Gram_matrix(A)\n",
        "        sLoss+= w*0.25*K.sum(K.square(G_gram - A_gram))/ (N_l**2 * M_l**2)\n",
        "    return sLoss\n",
        "  \n",
        "def get_total_loss(gImPlaceholder, alpha=1.0, beta=30.0):\n",
        "    F = get_feature_reps(gImPlaceholder, layer_names=[content_layer_name], model=gModel)[0]\n",
        "    Gs = get_feature_reps(gImPlaceholder, layer_names=style_layer_names, model=gModel)\n",
        "    contentLoss = get_content_loss(F, P)\n",
        "    styleLoss = get_style_loss(ws, Gs, As)\n",
        "    totalLoss = alpha*contentLoss + beta*styleLoss\n",
        "    return totalLoss\n",
        " "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "n6TiEh1R3y3x"
      },
      "source": [
        "## Putting it all together\n",
        "\n",
        "Here `calualate_loss` function will calculate total loss:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "EXH3PUgm30nl"
      },
      "outputs": [],
      "source": [
        "def calculate_loss(gImArr):\n",
        "  \"\"\"\n",
        "  Calculate total loss using K.function\n",
        "  \"\"\"\n",
        "  if gImArr.shape != (1, img_size, img_size, 3):\n",
        "        gImArr = gImArr.reshape((1,img_size, img_size, 3))\n",
        "  loss_fcn = K.function([gModel.input], [get_total_loss(gModel.input)])\n",
        "  grad_fcn = K.function([gModel.input], \n",
        "                        K.gradients(get_total_loss(gModel.input), [gModel.input]))\n",
        "  return loss_fcn([gImArr])[0].astype('float64'),grad_fcn([gImArr])[0].flatten().astype('float64')\n",
        "\n",
        "content_layer_name = 'block4_conv2'\n",
        "style_layer_names = ['block1_conv1','block2_conv1','block3_conv1','block4_conv1']\n",
        "\n",
        "P = get_feature_reps(x=img_content_var, layer_names=[content_layer_name], model=cModel)[0]\n",
        "As = get_feature_reps(x=img_style_var, layer_names=style_layer_names, model=sModel)\n",
        "ws = np.ones(len(style_layer_names))/float(len(style_layer_names))\n",
        "\n",
        "img_result = np.random.randint(256,size=(img_size,img_size,3)).astype(np.float64)\n",
        "img_result_1 = preprocess_input(np.expand_dims(img_result, axis=0))\n",
        "\n",
        "iterations = 10\n",
        "x_opt = img_result_1.flatten()"
      ]
    },
    {
      "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": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 350
        },
        "colab_type": "code",
        "id": "9vx4mDtohZDn",
        "outputId": "8c948643-e869-4918-8ddf-e747178bc9e9"
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 320636 characters]",
            "text/plain": [
              "<matplotlib.figure.Figure at 0x7f43494f9630>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "xopt, f_val, info= fmin_l_bfgs_b(calculate_loss, x_opt, maxiter=iterations, disp=True)\n",
        "plt.imshow(postprocess_array(xopt.copy()))\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 350
        },
        "colab_type": "code",
        "id": "qkbBCUTKGots",
        "outputId": "ba6c706a-cb8d-4d9f-a01d-db6e4a261b61"
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 319112 characters]",
            "text/plain": [
              "<matplotlib.figure.Figure at 0x7f28933be2d0>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "iterations = 20\n",
        "xopt, f_val, info= fmin_l_bfgs_b(calculate_loss, xopt, fprime=get_grad,\n",
        "                            maxiter=iterations, disp=True)\n",
        "plt.imshow(postprocess_array(xopt.copy()))\n",
        "plt.show()"
      ]
    },
    {
      "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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "JlUBzILGYFLs"
      },
      "outputs": [],
      "source": [
        "def total_variation_loss(x):\n",
        "    a = K.square(x[:,:img_size-1,:img_size-1,:] - x[:, 1:, :img_size-1,:])\n",
        "    b = K.square(x[:,:img_size-1,:img_size-1,:] - x[:, :img_size-1, 1:,:])\n",
        "    return K.sum(K.pow(a + b, 1.25))\n",
        "  \n",
        "def get_total_loss(gImPlaceholder, alpha=1.0, beta=30.0):\n",
        "    F = get_feature_reps(gImPlaceholder, layer_names=[content_layer_name], model=gModel)[0]\n",
        "    Gs = get_feature_reps(gImPlaceholder, layer_names=style_layer_names, model=gModel)\n",
        "    contentLoss = get_content_loss(F, P)\n",
        "    styleLoss = get_style_loss(ws, Gs, As)\n",
        "    variationLoss = total_variation_loss(gImPlaceholder)\n",
        "    totalLoss = alpha*contentLoss + beta*styleLoss + variationLoss\n",
        "    return totalLoss\n",
        "\n",
        "img_result = np.random.randint(256,size=(img_size,img_size,3)).astype(np.float64)\n",
        "img_result_1 = preprocess_input(np.expand_dims(img_result, axis=0))\n",
        "iterations = 10\n",
        "x_opt = img_result_1.flatten()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 350
        },
        "colab_type": "code",
        "id": "gi4w-jKh-WxO",
        "outputId": "ba8278ff-5760-404d-c9a9-98ef05aa95d3"
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 291516 characters]",
            "text/plain": [
              "<matplotlib.figure.Figure at 0x7fb25130ba58>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "xopt, f_val, info= fmin_l_bfgs_b(calculate_loss, x_opt, maxiter=iterations, disp=True)\n",
        "plt.imshow(postprocess_array(xopt.copy()))\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "DENlUJLdNcek"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "StyleTransfer.ipynb",
      "provenance": [],
      "version": "0.3.2"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.8.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
