Style_Transfer_Demo_InceptionV3_Dynamic_Shape.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Style Transfer Demo - Dynamic Shape.ipynb",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [],
      "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/sayakpaul/Adventures-in-TensorFlow-Lite/blob/master/Style_Transfer_Demo_InceptionV3_Dynamic_Shape.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZMpWwQOT8rzG",
        "colab_type": "text"
      },
      "source": [
        "## Stylize Your Images In 5 Minutes\n",
        "\n",
        "Demo developed by [Sayak Paul](https://twitter.com/RisingSayak).\n",
        "\n",
        "Neural style transfer is one of the most interesting applications of deep learning. We've created a demo to help you recreate your images in the style of famous artists. \n",
        "\n",
        "![](https://storage.googleapis.com/download.tensorflow.org/models/tflite/arbitrary_style_transfer/table.png)\n",
        "\n",
        "The image and code id adapted from [this tutorial](https://www.tensorflow.org/lite/models/style_transfer/overview). \n",
        "\n",
        "## To start the demo click the play button to the left of each cell, one by one.\n",
        "\n",
        "You need to run the following cells in order to get stylized images.\n",
        "\n",
        "The easiest way to do that is to click on the cell and press Shift + Enter or press the play button. Each cell takes between 10-30 seconds to execute."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ciYyUYVZT4IA",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title 1. Setup 🧰\n",
        "#@markdown Just run this cell as is. ***Don't modify the code block.*** The setup should not take more than two minutes. \n",
        "!pip uninstall -q -y tensorflow \n",
        "!pip install -q tf-nightly\n",
        "\n",
        "from IPython.display import clear_output\n",
        "import matplotlib.pyplot as plt\n",
        "import tensorflow as tf\n",
        " \n",
        "!pip install wandb -qq\n",
        "import wandb\n",
        "wandb.login(anonymous='allow')\n",
        "clear_output()\n",
        "content_image = None # This needs to be in global scope\n",
        "\n",
        "print('You are all set!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3qttuI19VKOZ",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title 2. Choose a style image from the options below 🎆\n",
        "\n",
        "STYLE_IMAGE_NAME = 'IMAGE_2' #@param ['IMAGE_1', 'IMAGE_2', 'IMAGE_3', 'IMAGE_4', 'IMAGE_5']\n",
        "\n",
        "corresponding_url = {\n",
        "    'IMAGE_1': 'https://storage.googleapis.com/download.tensorflow.org/example_images/Vassily_Kandinsky%2C_1913_-_Composition_7.jpg',\n",
        "    'IMAGE_2': 'https://storage.googleapis.com/khanhlvg-public.appspot.com/arbitrary-style-transfer/style23.jpg',\n",
        "    'IMAGE_3': 'https://upload.wikimedia.org/wikipedia/commons/thumb/a/a5/Tsunami_by_hokusai_19th_century.jpg/1024px-Tsunami_by_hokusai_19th_century.jpg',\n",
        "    'IMAGE_4': 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c5/Edvard_Munch%2C_1893%2C_The_Scream%2C_oil%2C_tempera_and_pastel_on_cardboard%2C_91_x_73_cm%2C_National_Gallery_of_Norway.jpg/800px-Edvard_Munch%2C_1893%2C_The_Scream%2C_oil%2C_tempera_and_pastel_on_cardboard%2C_91_x_73_cm%2C_National_Gallery_of_Norway.jpg',\n",
        "    'IMAGE_5': 'https://upload.wikimedia.org/wikipedia/commons/thumb/e/ea/Van_Gogh_-_Starry_Night_-_Google_Art_Project.jpg/757px-Van_Gogh_-_Starry_Night_-_Google_Art_Project.jpg'\n",
        "}\n",
        "\n",
        "style_image_path = tf.keras.utils.get_file(STYLE_IMAGE_NAME + \".jpg\", corresponding_url[STYLE_IMAGE_NAME])\n",
        "print(\"Style image downloaded!\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wu4KVB0m5f1U",
        "colab_type": "text"
      },
      "source": [
        "<center><img src = 'https://i.ibb.co/Tmnwnbc/Untitled-Diagram.png'></img></center>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b-sI7KVNb-hz",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title 3. Upload an image you want to stylize 🖼\n",
        "from IPython.display import display\n",
        "import ipywidgets as widgets\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "import io\n",
        "\n",
        "clear_output()\n",
        "\n",
        "def button_click(change):\n",
        "    global content_image\n",
        "    print('\\n')\n",
        "    img = Image.open(io.BytesIO(uploader.data[-1]))\n",
        "    content_image = img\n",
        "    img = img.convert('RGB')\n",
        "    img.thumbnail((256, 256))\n",
        "    img.save('thumbnail.jpg')\n",
        "    image = Image.open('thumbnail.jpg')\n",
        "    display(image)\n",
        "    \n",
        "uploader = widgets.FileUpload()\n",
        "show_button = widgets.Button(description='Show image')\n",
        "show_button.on_click(button_click)\n",
        "\n",
        "widgets.VBox([widgets.Label('Upload a content image (must be an RGB or RGBA image). High-res images might take more time to be processed.'), uploader, show_button])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iOdyUkmZ6Caq",
        "colab_type": "text"
      },
      "source": [
        "You can upload as many images as you would want to but we will only pick the last uploaded one. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QP7VXfwZ6KD2",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title [Optional] You can also use your Webcam! 📸\n",
        "#@markdown Just execute this cell and click anywhere in the streaming feed.\n",
        "\n",
        "#@markdown Courtesy: https://ricardodeazambuja.com/\n",
        "from IPython.display import HTML, Audio\n",
        "from google.colab.output import eval_js\n",
        "from base64 import b64decode\n",
        "import numpy as np\n",
        "import io\n",
        "from PIL import Image\n",
        "\n",
        "VIDEO_HTML = \"\"\"\n",
        "<video autoplay\n",
        " width=%d height=%d style='cursor: pointer;'></video>\n",
        "<script>\n",
        "\n",
        "var video = document.querySelector('video')\n",
        "\n",
        "navigator.mediaDevices.getUserMedia({ video: true })\n",
        "  .then(stream=> video.srcObject = stream)\n",
        "  \n",
        "var data = new Promise(resolve=>{\n",
        "  video.onclick = ()=>{\n",
        "    var canvas = document.createElement('canvas')\n",
        "    var [w,h] = [video.offsetWidth, video.offsetHeight]\n",
        "    canvas.width = w\n",
        "    canvas.height = h\n",
        "    canvas.getContext('2d')\n",
        "          .drawImage(video, 0, 0, w, h)\n",
        "    video.srcObject.getVideoTracks()[0].stop()\n",
        "    video.replaceWith(canvas)\n",
        "    resolve(canvas.toDataURL('image/jpeg', %f))\n",
        "  }\n",
        "})\n",
        "</script>\n",
        "\"\"\"\n",
        "def take_photo(filename='photo.jpg', quality=0.8, size=(800,600)):\n",
        "    global content_image\n",
        "    display(HTML(VIDEO_HTML % (size[0],size[1],quality)))\n",
        "    data = eval_js(\"data\")\n",
        "    binary = b64decode(data.split(',')[1])\n",
        "    f = io.BytesIO(binary)\n",
        "    content_image = Image.open(f)\n",
        "    print('\\nImage captured! 🤳')\n",
        "\n",
        "take_photo()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WyENTFmggM19",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title 4. Load the images 👾\n",
        "#@markdown ***Don't modify this code block.***\n",
        "# Function to load an image from a file, and add a batch dimension.\n",
        "def load_img(path_to_img):\n",
        "  img = tf.io.read_file(path_to_img)\n",
        "  img = tf.io.decode_image(img, channels=3)\n",
        "  img = tf.image.convert_image_dtype(img, tf.float32)\n",
        "  img = img[tf.newaxis, :]\n",
        "\n",
        "  return img\n",
        "\n",
        "# Function to load an image from a file, and add a batch dimension.\n",
        "def load_content_img(image_pixels):\n",
        "    if image_pixels.shape[-1] == 4:\n",
        "        image_pixels = Image.fromarray(image_pixels)\n",
        "        img = image_pixels.convert('RGB')\n",
        "        img = np.array(img)\n",
        "        img = tf.convert_to_tensor(img)\n",
        "        img = tf.image.convert_image_dtype(img, tf.float32)\n",
        "        img = img[tf.newaxis, :]\n",
        "        return img\n",
        "    elif image_pixels.shape[-1] == 3:\n",
        "        img = tf.convert_to_tensor(image_pixels)\n",
        "        img = tf.image.convert_image_dtype(img, tf.float32)\n",
        "        img = img[tf.newaxis, :]\n",
        "        return img\n",
        "    elif image_pixels.shape[-1] == 1:\n",
        "        raise Error('Grayscale images not supported! Please try with RGB or RGBA images.')\n",
        "    print('Exception not thrown')\n",
        "\n",
        "# Function to pre-process by resizing an central cropping it.\n",
        "def preprocess_image(image, target_dim):\n",
        "  # Resize the image so that the shorter dimension becomes 256px.\n",
        "  shape = tf.cast(tf.shape(image)[1:-1], tf.float32)\n",
        "  short_dim = min(shape)\n",
        "  scale = target_dim / short_dim\n",
        "  new_shape = tf.cast(shape * scale, tf.int32)\n",
        "  image = tf.image.resize(image, new_shape)\n",
        "\n",
        "  # Central crop the image.\n",
        "  image = tf.image.resize_with_crop_or_pad(image, target_dim, target_dim)\n",
        "\n",
        "  return image\n",
        "\n",
        "# Convert the content image from Bytes to NumPy array.\n",
        "content_image = np.array(content_image)\n",
        "\n",
        "# Load the input images.\n",
        "content_image = load_content_img(content_image)\n",
        "style_image = load_img(style_image_path)\n",
        "print('Images loaded up!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lUDZZN7Gxaz3",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title 5. Download the style transfer networks from TF Hub 💻\n",
        "#@markdown ***Don't modify this code block.***\n",
        "#@markdown \n",
        "\n",
        "# Download the style bottleneck and transfer networks\n",
        "print('Downloading the model files...')\n",
        "\n",
        "style_predict_path = tf.keras.utils.get_file('style_predict.tflite', 'https://tfhub.dev/sayakpaul/lite-model/arbitrary-image-stylization-inceptionv3-dynamic-shapes/int8/predict/1?lite-format=tflite')\n",
        "style_transform_path = style_transform_path = tf.keras.utils.get_file('style_transform.tflite', 'https://tfhub.dev/sayakpaul/lite-model/arbitrary-image-stylization-inceptionv3-dynamic-shapes/int8/transfer/1?lite-format=tflite')\n",
        "\n",
        "print('Model files downloaded...')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sBNjdwyXlx_B",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title 6. Stylize image 🥁\n",
        "\n",
        "content_blending_ratio = 0 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "#@markdown You're encouraged to play with the different values of `content_blending_ratio`.\n",
        "\n",
        "content_image_size = 350 #@param {type:\"slider\", min:192, max:512, step:1}\n",
        "\n",
        "# Preprocess the input images.\n",
        "preprocessed_content_image = preprocess_image(content_image, content_image_size)\n",
        "preprocessed_style_image = preprocess_image(style_image, 256)\n",
        "\n",
        "print('Preprocessing the style and the content images...')\n",
        "print('Style image shape:', preprocessed_style_image.shape)\n",
        "print('Content image shape:', preprocessed_content_image.shape)\n",
        "\n",
        "def imshow(image, title=None):\n",
        "  if len(image.shape) > 3:\n",
        "    image = tf.squeeze(image, axis=0)\n",
        "\n",
        "  plt.imshow(image)\n",
        "  if title:\n",
        "    plt.title(title)\n",
        "\n",
        "plt.figure(figsize=(10, 10))\n",
        "\n",
        "plt.subplot(1, 3, 1)\n",
        "imshow(preprocessed_content_image, 'Content Image')\n",
        "\n",
        "plt.subplot(1, 3, 2)\n",
        "imshow(preprocessed_style_image, 'Style Image')\n",
        "\n",
        "# Function to run style prediction on preprocessed style image.\n",
        "def run_style_predict(preprocessed_style_image):\n",
        "  # Load the model.\n",
        "  interpreter = tf.lite.Interpreter(model_path=style_predict_path)\n",
        "\n",
        "  # Set model input.\n",
        "  interpreter.allocate_tensors()\n",
        "  input_details = interpreter.get_input_details()\n",
        "  interpreter.set_tensor(input_details[0][\"index\"], preprocessed_style_image)\n",
        "\n",
        "  # Calculate style bottleneck.\n",
        "  interpreter.invoke()\n",
        "  style_bottleneck = interpreter.tensor(\n",
        "      interpreter.get_output_details()[0][\"index\"]\n",
        "      )()\n",
        "\n",
        "  return style_bottleneck\n",
        "\n",
        "# Calculate style bottleneck for the preprocessed style image.\n",
        "print('Calculating style bottleneck...')\n",
        "style_bottleneck = run_style_predict(preprocessed_style_image)\n",
        "print('Style Bottleneck Shape:', style_bottleneck.shape)\n",
        "print('Stylizing image. It should not take more than three minutes...')\n",
        "\n",
        "# Run style transform on preprocessed style image\n",
        "def run_style_transform(style_bottleneck, preprocessed_content_image):\n",
        "  # Load the model.\n",
        "  interpreter = tf.lite.Interpreter(model_path=style_transform_path)\n",
        "\n",
        "  # Set model input.\n",
        "  input_details = interpreter.get_input_details()\n",
        "  for index in range(len(input_details)):\n",
        "    if input_details[index][\"name\"]=='content_image':\n",
        "      index = input_details[index][\"index\"]\n",
        "      interpreter.resize_tensor_input(index, [1, content_image_size, content_image_size, 3])\n",
        "  interpreter.allocate_tensors()\n",
        "\n",
        "  # Set model inputs.\n",
        "  for index in range(len(input_details)):\n",
        "    if input_details[index][\"name\"]=='Conv/BiasAdd':\n",
        "      interpreter.set_tensor(input_details[index][\"index\"], style_bottleneck)\n",
        "    elif input_details[index][\"name\"]=='content_image':\n",
        "      interpreter.set_tensor(input_details[index][\"index\"], preprocessed_content_image)\n",
        "  interpreter.invoke()\n",
        "\n",
        "  # Transform content image.\n",
        "  stylized_image = interpreter.tensor(\n",
        "      interpreter.get_output_details()[0][\"index\"]\n",
        "      )()\n",
        "\n",
        "  return stylized_image\n",
        "\n",
        "# Calculate style bottleneck of the content image.\n",
        "style_bottleneck_content = run_style_predict(\n",
        "    preprocess_image(content_image, 256)\n",
        ")\n",
        "\n",
        "# Blend the style bottleneck of style image and content image\n",
        "style_bottleneck_blended = content_blending_ratio * style_bottleneck_content \\\n",
        "                           + (1 - content_blending_ratio) * style_bottleneck\n",
        "\n",
        "# Stylize the content image using the style bottleneck.\n",
        "stylized_image = run_style_transform(style_bottleneck_blended, preprocessed_content_image)\n",
        "\n",
        "# Visualize the output.\n",
        "plt.subplot(1, 3, 3)\n",
        "imshow(stylized_image, 'Stylized Image')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}