style_transfer_2d.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "style-transfer-2d.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "JndnmDMp66FL",
        "FsFc1mE51tCd"
      ],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/fzantalis/colab_collection/blob/master/style_transfer_2d.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JndnmDMp66FL",
        "colab_type": "text"
      },
      "source": [
        "##### Copyright 2018 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hMqWDc_m6rUC",
        "colab_type": "code",
        "cellView": "both",
        "colab": {}
      },
      "source": [
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5nno2iRZXJbX",
        "colab_type": "text"
      },
      "source": [
        "# Style Transfer\n",
        "\n",
        "\n",
        "![header image](https://storage.googleapis.com/tensorflow-lucid/static/img/notebook-styletransfer-header.jpg =512x) \n",
        "\n",
        "This notebook uses [**Lucid**](https://github.com/tensorflow/lucid) to perform style transfer between two images, and show how different parameterizations affect that process.\n",
        "\n",
        "This notebook doesn't introduce the abstractions behind lucid; you may wish to also read the [Lucid tutorial](https://colab.research.google.com/github/tensorflow/lucid/blob/master/notebooks/tutorial.ipynb).\n",
        "\n",
        "**Note**: The easiest way to use this tutorial is [as a colab notebook](https://colab.research.google.com/drive/1XuxLjIZj9MV-lRCpXHBhLo5A41Zs0f8E), which allows you to dive in with no setup. We recommend you enable a free GPU by going:\n",
        "\n",
        "> **Runtime**   →   **Change runtime type**   →   **Hardware Accelerator: GPU**\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FsFc1mE51tCd",
        "colab_type": "text"
      },
      "source": [
        "## Install, Import, Load Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tavMPe3KQ8Cs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install  lucid\n",
        "!pip install tensorflow-gpu==1.15"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RBr8QbboRAdU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Imports\n",
        "\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "\n",
        "from lucid.modelzoo import vision_models\n",
        "from lucid.misc.io import show, load, save\n",
        "from lucid.misc.tfutil import create_session\n",
        "\n",
        "import lucid.optvis.objectives as objectives\n",
        "import lucid.optvis.param as param\n",
        "import lucid.optvis.render as render"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CbtI802fnv_A",
        "colab_type": "text"
      },
      "source": [
        "We will be writing a custom objective for style transfer, so let's grab this handy decorator that allows us to take a simple function and mark it as an objective in lucid:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WemFM-0dOOxL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from lucid.optvis.objectives import wrap_objective"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WBmcdJion8oJ",
        "colab_type": "text"
      },
      "source": [
        "And we'll use our default vision model, GoogLeNet:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yNALaA0QRJVT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = vision_models.InceptionV1()\n",
        "model.load_graphdef()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J-gqG5sqQgHG",
        "colab_type": "text"
      },
      "source": [
        "You can use this to see if you're running on a machine with a correctly configured GPU:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TpFer0KYz-lE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "tf.test.is_gpu_available()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FB4nsWJFb9VQ",
        "colab_type": "text"
      },
      "source": [
        "## Load input images"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G5LDE5mraW9u",
        "colab_type": "text"
      },
      "source": [
        "First, let's take a look at the two images we want to transfer styles between. This code should feel familiar from the Lucid tutorial:\n",
        "\n",
        "Later in the notebook there'll be an opportunity to upload your own images, too."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QdppjZM1aa88",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import files \n",
        "uploaded = files.upload() \n",
        "for ci in uploaded.keys():\n",
        "  print('User uploaded file \"{name}\" with length {length} bytes'.format(\n",
        "      name=ci, length=len(uploaded[ci])))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VgL2Cl0xakwS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import files \n",
        "uploaded = files.upload() \n",
        "for si in uploaded.keys():\n",
        "  print('User uploaded file \"{name}\" with length {length} bytes'.format(\n",
        "      name=si, length=len(uploaded[si])))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ji42rN8sN2sb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "content_image = load(ci)\n",
        "style_image = load(si)[..., :3] # removes transparency channel\n",
        "\n",
        "print(content_image.shape, style_image.shape)\n",
        "\n",
        "show(content_image)\n",
        "show(style_image)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O1cIiSgvcCl1",
        "colab_type": "text"
      },
      "source": [
        "## Introducing style transfer\n",
        "\n",
        "In lucid, style transfer (rendering an image in the visual style of a different image) can be expressed compactly.\n",
        "\n",
        "We feed two images, a **style image** and a **content image** into our model, which produces activations at all of its layers. \n",
        "From that we extract two subsets of activations, the style activations (orange) and the content activations (green), at different depths in the network. \n",
        "\n",
        "We then optimize an input to produce activations similar to the combined set of activations we extracted.\n",
        "\n",
        "![styletransfer-diagram](https://storage.googleapis.com/tensorflow-lucid/static/img/notebook-styletransfer-diagram.png =563x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EQwPDc4ddHIW",
        "colab_type": "text"
      },
      "source": [
        "Which layers work best for that is an aesthetic choice. Recommended reading: [Feature Visualization](https://distill.pub/2017/feature-visualization), in which the hero graphic shows what kind of visual concepts get captured at which layers in the model we are using."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tV7n64xVb-4k",
        "colab_type": "text"
      },
      "source": [
        "## YOU MAY NOW PRESS (CTRL + F10) TO AUTOMATICALLY EXECUTE ALL THE REMAINING STEPS"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_7dwzDwgOZVt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "style_layers = [\n",
        "  'conv2d2',\n",
        "  'mixed3a',\n",
        "  'mixed4a',\n",
        "  'mixed4b',\n",
        "  'mixed4c',\n",
        "]\n",
        "\n",
        "content_layers = [\n",
        "  'mixed3b',\n",
        "]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Oigu8czPfa6P",
        "colab_type": "text"
      },
      "source": [
        "To feed two images + one that is being optimized into our network we will use the batch dimension of a cutsom parameterization. It consists simply of the paramaterization of the optimized image stacked with the two images.\n",
        "\n",
        "(We also use cropping to allow style images of a larger size than content images.)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9Fsw1wK_nuMJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def style_transfer_param(content_image, style_image, decorrelate=True, fft=True):\n",
        "  style_transfer_input = param.image(*content_image.shape[:2], decorrelate=decorrelate, fft=fft)[0]\n",
        "  content_input = content_image\n",
        "  style_input = tf.random_crop(style_image, content_image.shape)\n",
        "  return tf.stack([style_transfer_input, content_input, style_input])\n",
        "\n",
        "# these constants help remember which image is at which batch dimension\n",
        "TRANSFER_INDEX = 0\n",
        "CONTENT_INDEX = 1\n",
        "STYLE_INDEX = 2"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qeJU0wIDpm0r",
        "colab_type": "text"
      },
      "source": [
        "Let's sanity check that shapes fit and the three inputs get stacked on their batch dimension correctly:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B2eE6dGQpYRV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "style_transfer_param(content_image, style_image).shape"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fpVnFlaxqlnz",
        "colab_type": "text"
      },
      "source": [
        "Here we use `@wrap_objective` to create a custom `activation_difference` objective in lucid.\n",
        "It takes as arguments the layer names specifying the subset of activations, and an image.\n",
        "It will extract the specified subset of activations from both the supplied image and the one we are optimizing, and compute their difference.\n",
        "\n",
        "Optionally one can pass a loss function defining how these differences should be penalized. By default, we penalize mean L1 difference.\n",
        "Also, one can pass a function transforming the subset of activations. By default, we don't transform the activations."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uKCaeYwf7-jI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def mean_L1(a, b):\n",
        "  return tf.reduce_mean(tf.abs(a-b))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Lq6k5NRsDzBy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "@wrap_objective\n",
        "def activation_difference(layer_names, activation_loss_f=mean_L1, transform_f=None, difference_to=CONTENT_INDEX):\n",
        "  def inner(T):\n",
        "    # first we collect the (constant) activations of image we're computing the difference to\n",
        "    image_activations = [T(layer_name)[difference_to] for layer_name in layer_names]\n",
        "    if transform_f is not None:\n",
        "      image_activations = [transform_f(act) for act in image_activations]\n",
        "\n",
        "    # we also set get the activations of the optimized image which will change during optimization\n",
        "    optimization_activations = [T(layer)[TRANSFER_INDEX] for layer in layer_names]\n",
        "    if transform_f is not None:\n",
        "      optimization_activations = [transform_f(act) for act in optimization_activations]\n",
        "\n",
        "    # we use the supplied loss function to compute the actual losses\n",
        "    losses = [activation_loss_f(a, b) for a, b in zip(image_activations, optimization_activations)]\n",
        "    return tf.add_n(losses)\n",
        "\n",
        "  return inner"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9jlOTiujr7ff",
        "colab_type": "text"
      },
      "source": [
        "For the style loss, we do not care about the spatial information in the activations. Using a gram matrix to transform the activations before penalizing their difference works well for this:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RSXAKI1D5eVq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def gram_matrix(array, normalize_magnitue=True):\n",
        "  channels = tf.shape(array)[-1]\n",
        "  array_flat = tf.reshape(array, [-1, channels])\n",
        "  gram_matrix = tf.matmul(array_flat, array_flat, transpose_a=True)\n",
        "  if normalize_magnitue:\n",
        "    length = tf.shape(array_flat)[0]\n",
        "    gram_matrix /= tf.cast(length, tf.float32)\n",
        "  return gram_matrix"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jTo-lDjmdWFN",
        "colab_type": "text"
      },
      "source": [
        "Now we have all the ingredients to set up an optimization in lucid: we will create two `activation_difference` objectives, a normal one for the **content loss**, and one transforming the activations using `gram_matrix` for the **style loss**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w8hZHgPy8wBh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "param_f = lambda: style_transfer_param(content_image, style_image)\n",
        "\n",
        "content_obj = 100 * activation_difference(content_layers, difference_to=CONTENT_INDEX)\n",
        "content_obj.description = \"Content Loss\"\n",
        "\n",
        "style_obj = activation_difference(style_layers, transform_f=gram_matrix, difference_to=STYLE_INDEX)\n",
        "style_obj.description = \"Style Loss\"\n",
        "\n",
        "objective = - content_obj - style_obj\n",
        "\n",
        "vis = render.render_vis(model, objective, param_f=param_f, thresholds=[512], verbose=False, print_objectives=[content_obj, style_obj])[-1]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W9WIoX0FCaSW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "show(vis)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1dx5C9mTs_FT",
        "colab_type": "text"
      },
      "source": [
        "That looks like style transfer to me!"
      ]
    }
  ]
}