image-style-transfer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bdb45c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from PIL import Image\n",
    "import time\n",
    "import functools\n",
    "import tensorflow as tf\n",
    "import tensorflow.keras as keras\n",
    "from tensorflow.keras.models import Model\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "mpl.rcParams['figure.figsize'] = (10,10)\n",
    "mpl.rcParams['axes.grid'] = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9138cfbd",
   "metadata": {},
   "source": [
    "# Abstract\n",
    "This notebook is based on a sample code hosted on Google's SeedBank project. However, the initiative doesn't seem to be active anymore, and the code samples disappeared.\n",
    "\n",
    "I added a few code mofications to improve the readability of some functions. The comments are mine were there to assist my understanding as I progressed through the original paper.\n",
    "\n",
    "The original sample code was based on \"A Neural Algorithm for Artistic Style\" https://arxiv.org/pdf/1508.06576.pdf. Additionally I include here another paper that dwellved more on the mathematics of the problem: \"Desmytifying Neural Style Transfer\" https://arxiv.org/pdf/1701.01036.pdf."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5992af1",
   "metadata": {},
   "source": [
    "# Eager execution\n",
    "This configures TensorFlow to execute eagerly. Eager execution is an alternative mode of execution in TensorFlow that allows operations to be executed immediately and return concrete values, as opposed to being added to a graph to be executed later. \n",
    "\n",
    "Eager execution is enabled by default in TensorFlow 2.0, but since this notebook was ported from a TensorFlow 1.x version I decided to leave it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "940da1c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.compat.v1.enable_eager_execution()\n",
    "print(\"Eager execution: {}\".format(tf.executing_eagerly()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5867f7e",
   "metadata": {},
   "source": [
    "# Image visualization utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cac02800",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_img(path_to_img):\n",
    "  \"\"\"\n",
    "  inputs:\n",
    "  - path:to_image: str containing file path on disk\n",
    "  output:\n",
    "  - np.array with shapoe (batch_size, row, col, channel)\n",
    "  \"\"\"\n",
    "  max_dim = 512\n",
    "  img = Image.open(path_to_img)\n",
    "  img_width = img.size[0]\n",
    "  img_height = img.size[1]\n",
    "  longest_dim = max(img.size)\n",
    "  scale = max_dim/longest_dim\n",
    "    \n",
    "  img = img.resize((round(img_width*scale), round(img_height*scale)), Image.ANTIALIAS)\n",
    "  img = tf.keras.utils.img_to_array(img)\n",
    "  img = np.expand_dims(img, axis=0) # add batch_size dimension\n",
    "\n",
    "  return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b0f0cfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def imshow(img, title=None):\n",
    "  # Remove the batch dimension\n",
    "  out = np.squeeze(img, axis=0)\n",
    "\n",
    "  # Normalize for display \n",
    "  out = out.astype('uint8')\n",
    "\n",
    "  # Display\n",
    "  if title is not None:\n",
    "    plt.title(title)\n",
    "\n",
    "  plt.imshow(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c16b7b7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_and_process_img(path_to_img):\n",
    "  img = load_img(path_to_img)\n",
    "  img = tf.keras.applications.vgg19.preprocess_input(img)\n",
    "  return img\n",
    "\n",
    "def deprocess_img(processed_img):\n",
    "  x = processed_img.copy()\n",
    "  if len(x.shape) == 4:\n",
    "    x = np.squeeze(x, 0)\n",
    "  assert len(x.shape) == 3, (\"Input to deprocess image must be an image of \"\n",
    "                             \"dimension [1, height, width, channel] or [height, width, channel]\")\n",
    "  if len(x.shape) != 3:\n",
    "    raise ValueError(\"Invalid input to deprocessing image\")\n",
    "  \n",
    "  # perform the inverse of the preprocessing step\n",
    "  x[:, :, 0] += 103.939\n",
    "  x[:, :, 1] += 116.779\n",
    "  x[:, :, 2] += 123.68\n",
    "  x = x[:, :, ::-1]\n",
    "\n",
    "  x = np.clip(x, 0, 255).astype('uint8')\n",
    "  return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1986a0e",
   "metadata": {},
   "source": [
    "# Display the content and style images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83960e45",
   "metadata": {},
   "outputs": [],
   "source": [
    "content_path = './data/three_samurai.jpeg'\n",
    "style_path = './data/The_Great_Wave_off_Kanagawa.jpg'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6431ddc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10,10))\n",
    "\n",
    "content = load_img(content_path).astype('uint8')\n",
    "style = load_img(style_path).astype('uint8')\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "imshow(content, 'Content Image')\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "imshow(style, 'Style Image')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ba8a1a9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "4cca8faf",
   "metadata": {},
   "source": [
    "# The model\n",
    "The basic idea from the paper is to load the VGG19 model and access the intermediate layers.\n",
    "These layers will then be used to create a new model that will take an input image and return the outputs from these intermediate layers from the VGG model. \n",
    "\n",
    "These intermidiate layers represent the style and content weights of the pre-trained model. This is inferred from direct experiments mentioned in the paper where they could directly visualise the information each layer contains about the input image by reconstructing the image only from the feature maps in that layer. Therefore, they could observe how the representations changed as they visualized deeper layers in the network. \n",
    "\n",
    "When matching the style representations up to higher layers in the network, local images structures are matched on an increasingly large scale, leading to a smoother and more continuous visual experience. Thus, the visually most appealing images are usually created by matching the style representation up to the highest layers in the network\n",
    "\n",
    "There is a ratio α/β that represents the emphasis on matching the content of the photograph and the style of the artwork"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bc2d11a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Content layer where will pull our feature maps\n",
    "content_layers = ['block5_conv2'] \n",
    "\n",
    "# Style layer we are interested in\n",
    "style_layers = ['block1_conv1',\n",
    "                'block2_conv1',\n",
    "                'block3_conv1', \n",
    "                'block4_conv1', \n",
    "                'block5_conv1'\n",
    "               ]\n",
    "\n",
    "num_content_layers = len(content_layers)\n",
    "num_style_layers = len(style_layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "590105ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model():\n",
    "  \"\"\" Creates our model with access to intermediate layers. \n",
    "  \n",
    "  This function will load the VGG19 model and access the intermediate layers. \n",
    "  These layers will then be used to create a new model that will take input image\n",
    "  and return the outputs from these intermediate layers from the VGG model. \n",
    "  \n",
    "  Returns:\n",
    "    returns a keras model that takes image inputs and outputs the style and \n",
    "      content intermediate layers. \n",
    "  \"\"\"\n",
    "  # Load our model. We load pretrained VGG, trained on imagenet data\n",
    "  vgg = tf.keras.applications.vgg19.VGG19(include_top=False, weights='imagenet')\n",
    "  vgg.trainable = False\n",
    "  # Get output layers corresponding to style and content layers \n",
    "  style_outputs = [vgg.get_layer(name).output for name in style_layers]\n",
    "  content_outputs = [vgg.get_layer(name).output for name in content_layers]\n",
    "  model_outputs = style_outputs + content_outputs\n",
    "  # Build model \n",
    "  return Model(vgg.input, model_outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24f9047c",
   "metadata": {},
   "source": [
    "# Loss and gradient\n",
    "\n",
    "The images are synthesised by finding an image that simultaneously matches the content representation of the photograph and the style representation of the respective piece of art.\n",
    "\n",
    "However, the loss function we minimise during image synthesis contains two terms for content and style respectively, that are well separated (see Methods). We can therefore smoothly regulate the emphasis on either reconstructing the content or the style.\n",
    "\n",
    "## Gram matrix\n",
    "The Gram matrix of an image represents the second-order statistics of the image, or the relationships between the different features in the image. It is calculated by taking the dot product of each feature vector in the image with itself and with every other feature vector. This results in a square matrix, where each row and column represents a different feature vector. The diagonal elements of the matrix represent the variance of each feature, and the off-diagonal elements represent the covariance between different features.\n",
    "\n",
    "The Gram matrix can be used to represent the style of an image. The style of an image is captured by the relationships between the different features in the image, rather than the individual features themselves. For example, a painting with a lot of brushstrokes will have a different Gram matrix than a painting with a lot of smooth areas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71961ded",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1d6fc97",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7401ca02",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gram_matrix(input_tensor):\n",
    "  # We make the image channels first \n",
    "  channels = int(input_tensor.shape[-1])\n",
    "  a = tf.reshape(input_tensor, [-1, channels])\n",
    "  n = tf.shape(a)[0]\n",
    "  gram = tf.matmul(a, a, transpose_a=True)\n",
    "  return gram / tf.cast(n, tf.float32)\n",
    "\n",
    "def get_style_loss(base_style, gram_target):\n",
    "  \"\"\"Expects two images of dimension h, w, c\"\"\"\n",
    "  # height, width, num filters of each layer\n",
    "  # We scale the loss at a given layer by the size of the feature map and the number of filters\n",
    "  height, width, channels = base_style.get_shape().as_list()\n",
    "  gram_style = gram_matrix(base_style)\n",
    "  \n",
    "  return tf.reduce_mean(tf.square(gram_style - gram_target))# / (4. * (channels ** 2) * (width * height) ** 2)\n",
    "\n",
    "def get_content_loss(base_content, target):\n",
    "  return tf.reduce_mean(tf.square(base_content - target))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "979cc62e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_loss(model, loss_weights, init_image, gram_style_features, content_features):\n",
    "  \"\"\"This function will compute the loss total loss.\n",
    "  \n",
    "  Arguments:\n",
    "    model: The model that will give us access to the intermediate layers\n",
    "    loss_weights: The weights of each contribution of each loss function. \n",
    "      (style weight, content weight, and total variation weight)\n",
    "    init_image: Our initial base image. This image is what we are updating with \n",
    "      our optimization process. We apply the gradients wrt the loss we are \n",
    "      calculating to this image.\n",
    "    gram_style_features: Precomputed gram matrices corresponding to the \n",
    "      defined style layers of interest.\n",
    "    content_features: Precomputed outputs from defined content layers of \n",
    "      interest.\n",
    "      \n",
    "  Returns:\n",
    "    returns the total loss, style loss, content loss, and total variational loss\n",
    "  \"\"\"\n",
    "  style_weight, content_weight = loss_weights\n",
    "  \n",
    "  # Feed our init image through our model. This will give us the content and \n",
    "  # style representations at our desired layers. Since we're using eager\n",
    "  # our model is callable just like any other function!\n",
    "  model_outputs = model(init_image)\n",
    "  \n",
    "  style_output_features = model_outputs[:num_style_layers]\n",
    "  content_output_features = model_outputs[num_style_layers:]\n",
    "  \n",
    "  style_score = 0\n",
    "  content_score = 0\n",
    "\n",
    "  # Accumulate style losses from all layers\n",
    "  # Here, we equally weight each contribution of each loss layer\n",
    "  weight_per_style_layer = 1.0 / float(num_style_layers)\n",
    "  for target_style, comb_style in zip(gram_style_features, style_output_features):\n",
    "    style_score += weight_per_style_layer * get_style_loss(comb_style[0], target_style)\n",
    "    \n",
    "  # Accumulate content losses from all layers \n",
    "  weight_per_content_layer = 1.0 / float(num_content_layers)\n",
    "  for target_content, comb_content in zip(content_features, content_output_features):\n",
    "    content_score += weight_per_content_layer* get_content_loss(comb_content[0], target_content)\n",
    "  \n",
    "  style_score *= style_weight\n",
    "  content_score *= content_weight\n",
    "\n",
    "  # Get total loss\n",
    "  loss = style_score + content_score \n",
    "  return loss, style_score, content_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d346c892",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_grads(cfg):\n",
    "  with tf.GradientTape() as tape: \n",
    "    all_loss = compute_loss(**cfg)\n",
    "  # Compute gradients wrt input image\n",
    "  total_loss = all_loss[0]\n",
    "  return tape.gradient(total_loss, cfg['init_image']), all_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3027ca71",
   "metadata": {},
   "source": [
    "# Style transfer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1a38a20",
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython.display\n",
    "\n",
    "def get_feature_representations(model, content_path, style_path):\n",
    "  \"\"\"Helper function to compute our content and style feature representations.\n",
    "\n",
    "  This function will simply load and preprocess both the content and style \n",
    "  images from their path. Then it will feed them through the network to obtain\n",
    "  the outputs of the intermediate layers. \n",
    "  \n",
    "  Arguments:\n",
    "    model: The model that we are using.\n",
    "    content_path: The path to the content image.\n",
    "    style_path: The path to the style image\n",
    "    \n",
    "  Returns:\n",
    "    returns the style features and the content features. \n",
    "  \"\"\"\n",
    "  # Load our images in \n",
    "  content_image = load_and_process_img(content_path)\n",
    "  style_image = load_and_process_img(style_path)\n",
    "  \n",
    "  # batch compute content and style features\n",
    "  style_outputs = model(style_image)\n",
    "  content_outputs = model(content_image)\n",
    "  \n",
    "  # Get the style and content feature representations from our model  \n",
    "  style_features = [style_layer[0] for style_layer in style_outputs[:num_style_layers]]\n",
    "  content_features = [content_layer[0] for content_layer in content_outputs[num_style_layers:]]\n",
    "  return style_features, content_features\n",
    "\n",
    "def run_style_transfer(content_path, \n",
    "                       style_path,\n",
    "                       num_iterations=1000,\n",
    "                       content_weight=1e3, \n",
    "                       style_weight=1e-5): \n",
    "  # We don't need to (or want to) train any layers of our model, so we set their\n",
    "  # trainable to false. \n",
    "  model = get_model() \n",
    "  for layer in model.layers:\n",
    "    layer.trainable = False\n",
    "  \n",
    "  # Get the style and content feature representations (from our specified intermediate layers) \n",
    "  style_features, content_features = get_feature_representations(model, content_path, style_path)\n",
    "  gram_style_features = [gram_matrix(style_feature) for style_feature in style_features]\n",
    "  \n",
    "  # Set initial image\n",
    "  init_image = load_and_process_img(content_path)\n",
    "  init_image = tf.Variable(init_image, dtype=tf.float32)\n",
    "  # Create our optimizer\n",
    "  opt = tf.optimizers.Adam(learning_rate=5, epsilon=1e-1)\n",
    "\n",
    "  # For displaying intermediate images \n",
    "  iter_count = 1\n",
    "  \n",
    "  # Store our best result\n",
    "  best_loss, best_img = float('inf'), None\n",
    "  \n",
    "  # Create a nice config \n",
    "  loss_weights = (style_weight, content_weight)\n",
    "  cfg = {\n",
    "      'model': model,\n",
    "      'loss_weights': loss_weights,\n",
    "      'init_image': init_image,\n",
    "      'gram_style_features': gram_style_features,\n",
    "      'content_features': content_features\n",
    "  }\n",
    "    \n",
    "  # For displaying\n",
    "  num_rows = 2\n",
    "  num_cols = 5\n",
    "  display_interval = num_iterations/(num_rows*num_cols)\n",
    "  start_time = time.time()\n",
    "  global_start = time.time()\n",
    "  \n",
    "  norm_means = np.array([103.939, 116.779, 123.68])\n",
    "  min_vals = -norm_means\n",
    "  max_vals = 255 - norm_means   \n",
    "  \n",
    "  imgs = []\n",
    "  for i in range(num_iterations):\n",
    "    grads, all_loss = compute_grads(cfg)\n",
    "    loss, style_score, content_score = all_loss\n",
    "    opt.apply_gradients([(grads, init_image)])\n",
    "    clipped = tf.clip_by_value(init_image, min_vals, max_vals)\n",
    "    init_image.assign(clipped)\n",
    "    end_time = time.time() \n",
    "    \n",
    "    if loss < best_loss:\n",
    "      # Update best loss and best image from total loss. \n",
    "      best_loss = loss\n",
    "      best_img = deprocess_img(init_image.numpy())\n",
    "\n",
    "    if i % display_interval== 0:\n",
    "      start_time = time.time()\n",
    "      \n",
    "      # Use the .numpy() method to get the concrete numpy array\n",
    "      plot_img = init_image.numpy()\n",
    "      plot_img = deprocess_img(plot_img)\n",
    "      imgs.append(plot_img)\n",
    "      IPython.display.clear_output(wait=True)\n",
    "      IPython.display.display_png(Image.fromarray(plot_img))\n",
    "      print('Iteration: {}'.format(i))        \n",
    "      print('Total loss: {:.4e}, ' \n",
    "            'style loss: {:.4e}, '\n",
    "            'content loss: {:.4e}, '\n",
    "            'time: {:.4f}s'.format(loss, style_score, content_score, time.time() - start_time))\n",
    "  print('Total time: {:.4f}s'.format(time.time() - global_start))\n",
    "  IPython.display.clear_output(wait=True)\n",
    "  plt.figure(figsize=(14,4))\n",
    "  for i,img in enumerate(imgs):\n",
    "      plt.subplot(num_rows,num_cols,i+1)\n",
    "      plt.imshow(img)\n",
    "      plt.xticks([])\n",
    "      plt.yticks([])\n",
    "      \n",
    "  return best_img, best_loss "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18488420",
   "metadata": {},
   "outputs": [],
   "source": [
    "best, best_loss = run_style_transfer(content_path, \n",
    "                                     style_path,\n",
    "                                     num_iterations=1000)\n",
    "Image.fromarray(best)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74fa650e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_results(best_img, content_path, style_path, show_large_final=True):\n",
    "  plt.figure(figsize=(10, 5))\n",
    "  content = load_img(content_path) \n",
    "  style = load_img(style_path)\n",
    "\n",
    "  plt.subplot(1, 2, 1)\n",
    "  imshow(content, 'Content Image')\n",
    "\n",
    "  plt.subplot(1, 2, 2)\n",
    "  imshow(style, 'Style Image')\n",
    "\n",
    "  if show_large_final: \n",
    "    plt.figure(figsize=(10, 10))\n",
    "\n",
    "    plt.imshow(best_img)\n",
    "    plt.title('Output Image')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32b2ec42",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f659b7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_results(best, content_path, style_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7283c3b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21d89e6f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ac06090",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "208799b2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Environment (conda_st-image)",
   "language": "python",
   "name": "conda_st-image"
  },
  "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.16"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}