{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "491ebf48-3611-4a78-9f0f-bdc505b1f3e0"
   },
   "source": [
    "## Neural Style Transfer with VGG19\n",
    "\n",
    "Based on section 12.3 of [*Deep Learning with Python, 2nd Edition*](https://github.com/fchollet/deep-learning-with-python-notebooks/blob/master/second_edition/chapter12_part03_neural-style-transfer.ipynb) by Fran&ccedil;ois Chollet.\n",
    "\n",
    "<img src=\"https://science.slc.edu/jmarshall/bioai/images/neural_style_xfer.jpg\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "import keras"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download the images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "25b81c5b-f5b8-4f69-8bd9-a7d5f9305bd8",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# content/base images\n",
    "base_image_path = keras.utils.get_file(origin=\"https://science.slc.edu/jmarshall/bioai/images/tuebingen.jpg\")\n",
    "#base_image_path = keras.utils.get_file(origin=\"http://science.slc.edu/jmarshall/bioai/images/sanfran.jpg\")\n",
    "#base_image_path = keras.utils.get_file(origin=\"http://science.slc.edu/jmarshall/bioai/images/elephants.jpg\")\n",
    "#base_image_path = keras.utils.get_file(origin=\"http://science.slc.edu/jmarshall/bioai/images/jellyfish.jpg\")\n",
    "#base_image_path = keras.utils.get_file(origin=\"http://science.slc.edu/jmarshall/bioai/images/flamingos.jpg\")\n",
    "#base_image_path = keras.utils.get_file(origin=\"http://science.slc.edu/jmarshall/bioai/images/tiger.jpg\")\n",
    "\n",
    "# style reference images\n",
    "style_image_path = keras.utils.get_file(origin=\"http://science.slc.edu/jmarshall/bioai/images/starry_night.jpg\")\n",
    "#style_image_path = keras.utils.get_file(origin=\"http://science.slc.edu/jmarshall/bioai/images/the_scream.jpg\")\n",
    "#style_image_path = keras.utils.get_file(origin=\"http://science.slc.edu/jmarshall/bioai/images/eiffel_tower.jpg\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Content/Base image\n",
    "\n",
    "The **content** or **base** image will serve as our input image. The visual properties of the **style reference** image will be applied to the base image to create the **combination** image, which will be generated as output. We will start by resizing the base image to be 400 pixels high."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_image_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_jpeg = keras.utils.load_img(base_image_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "type(base_jpeg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_jpeg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "original_width, original_height = base_jpeg.size\n",
    "aspect_ratio = original_width / original_height\n",
    "\n",
    "# height and width of the new image to be generated\n",
    "img_height = 400\n",
    "img_width = round(img_height * aspect_ratio)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_height, img_width"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_jpeg = keras.utils.load_img(base_image_path, target_size=(img_height, img_width))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_jpeg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Style Reference image\n",
    "\n",
    "We also need to resize the style reference image to match the size of the base image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "style_image_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "style_jpeg = keras.utils.load_img(style_image_path, target_size=(img_height, img_width))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "style_jpeg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### Utility functions for VGG19 images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "def preprocess_image(image_path):\n",
    "    # image_path is a string\n",
    "    img = keras.utils.load_img(image_path, target_size=(img_height, img_width))\n",
    "    img = keras.utils.img_to_array(img)\n",
    "    img = np.expand_dims(img, axis=0)\n",
    "    img = keras.applications.vgg19.preprocess_input(img)\n",
    "    return img\n",
    "\n",
    "def deprocess_image(img):\n",
    "    # img is a numpy array\n",
    "    img = img.reshape((img_height, img_width, 3))\n",
    "    # undo zero-centering transformation done by vgg19.preprocess_input\n",
    "    img[:,:,0] += 103.939  # mean pixel values of\n",
    "    img[:,:,1] += 116.779  # each color channel\n",
    "    img[:,:,2] += 123.68   # over ImageNet dataset\n",
    "    # convert image from 'BGR' to 'RGB' format (another effect of vgg19.preprocess_input)\n",
    "    img = img[:,:,::-1]\n",
    "    img = np.clip(img, 0, 255).astype('uint8')\n",
    "    return img"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### Feature extractor for the pre-trained VGG19 network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "from keras.applications.vgg19 import VGG19\n",
    "vgg19 = VGG19(weights=\"imagenet\", include_top=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vgg19.summary(print_fn=print)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "outputs_dictionary = { layer.name: layer.output for layer in vgg19.layers }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "feature_extractor = keras.Model(inputs=vgg19.inputs, outputs=outputs_dictionary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_extractor(vgg19.input)['block1_conv1']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, we can extract the features of the base image from layer `block1_conv1` like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_image = preprocess_image(base_image_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = feature_extractor([base_image])['block1_conv1']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features.numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "33b2a0f4-10ff-4559-9b20-23538cfc3463"
   },
   "source": [
    "### Components of the loss function\n",
    "\n",
    "The overall **loss** function comprises three components: the **content loss**, the **style loss**, and the **total variation loss**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "77d2783d-1dea-4ae9-bb31-1422763ba8b1",
    "tags": []
   },
   "outputs": [],
   "source": [
    "content_loss_layer_name = \"block5_conv2\"\n",
    "\n",
    "style_loss_layer_names = {\n",
    "    \"block1_conv1\",\n",
    "    \"block2_conv1\",\n",
    "    \"block3_conv1\",\n",
    "    \"block4_conv1\",\n",
    "    \"block5_conv1\",\n",
    "}\n",
    "\n",
    "# weight coefficients for each loss component\n",
    "content_loss_weight = 2.5e-8\n",
    "style_loss_weight = 1e-6\n",
    "total_variation_loss_weight = 1e-6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "#### Content loss\n",
    "\n",
    "To compute the **content loss**, we apply the VGG19 network to the base image and to the generated combination image, and then compare the internal activation patterns (features) created for both images. The closer the activation patterns are to each other, the lower the content loss value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "def content_loss(base_image_features, combination_image_features):\n",
    "    return keras.ops.sum(keras.ops.square(combination_image_features - base_image_features))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "#### Style loss\n",
    "\n",
    "The **style loss** measures the similarity of the feature maps computed from the style-reference image and the generated combination image at different spatial scales within the VGG19 convolutional network.  The correlation of the features at each layer specified by `style_loss_layer_names` are computed based on the *Gram matrix* from linear algebra.  These feature correlations capture the statistics of the patterns at a particular spatial scale, which empirically correspond to the appearance of the textures found at this scale."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "c4f1c687-512a-4fca-b1e0-7ca60da23948",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 2-D matrix\n",
    "a = np.array([[1,2,3],[4,5,6],[7,8,9]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ea308f44-22d3-49a4-a642-54bf028be5af",
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5a29b1ce-d3ed-4e5d-9dcc-aa2b0663ca54",
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(tf.transpose(a).numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "824150b7-4846-49b1-ac8a-d8842c713ad4",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 3-D matrix\n",
    "b = np.array([[[1,2,3],[4,5,6],[7,8,9]],\n",
    "              [[11,22,33],[44,55,66],[77,88,99]],\n",
    "              [[111,222,333],[444,555,666],[777,888,999]]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "759aa572-387b-4b55-93cf-27e56e490c65",
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tf.transpose(b).numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "67cb08e5-0a85-49c3-b1e3-82cb875f745f",
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(tf.transpose(b, perm=(2,0,1)).numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "950e6ac4-2dd6-49d4-9d91-f87d8e746bf5",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def show_gram_matrix(x):\n",
    "    print(\"Input matrix:\\n\", x)\n",
    "    x = tf.transpose(x, perm=(2,0,1))  # x must be a 3-D matrix\n",
    "    print(\"\\nTranspose:\\n\", x.numpy())\n",
    "    features = tf.reshape(x, (tf.shape(x)[0], -1))\n",
    "    print(\"\\nFeatures:\\n\", features.numpy())\n",
    "    gram = tf.matmul(features, tf.transpose(features))\n",
    "    print(\"\\nGram matrix:\\n\", gram.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "f4559639-739c-4fe6-ba78-a33c8ec992a4",
    "tags": []
   },
   "outputs": [],
   "source": [
    "show_gram_matrix(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0f98b23e-b80d-4c17-84b7-1d4bc6292963",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def gram_matrix(x):\n",
    "    x = tf.transpose(x, (2,0,1))  # x must be a 3-D matrix\n",
    "    features = tf.reshape(x, (tf.shape(x)[0], -1))\n",
    "    gram = tf.matmul(features, tf.transpose(features))\n",
    "    return gram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cabc1f61-4ed9-49d0-b63e-025b64aa0f51",
    "tags": []
   },
   "outputs": [],
   "source": [
    "gram_matrix(b).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "c6928a2a-5ce8-421b-ba66-b8478d01d953",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def style_loss(style_image_features, combination_image_features):\n",
    "    S = gram_matrix(style_image_features)\n",
    "    G = gram_matrix(combination_image_features)\n",
    "    color_channels = 3\n",
    "    size = img_height * img_width\n",
    "    loss = keras.ops.sum(keras.ops.square(S - G)) / (4.0 * (color_channels ** 2) * (size ** 2))\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "#### Total variation loss\n",
    "\n",
    "The **total variation loss** encourages spatial continuity within the generated combination image, thus avoiding overly pixelated results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0ca3e24d-b339-48eb-b6b6-ebc17ef047c0",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def total_variation_loss(combination_image):\n",
    "    height, width = combination_image.shape[1], combination_image.shape[2]\n",
    "    a = keras.ops.square(combination_image[:,:height-1,:width-1,:] - combination_image[:,1:,:width-1,:])\n",
    "    b = keras.ops.square(combination_image[:,:height-1,:width-1,:] - combination_image[:,:height-1,1:,:])\n",
    "    return keras.ops.sum(keras.ops.power(a + b, 1.25))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "09d683a7-75fa-4037-902e-f648aba55be0",
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Overall loss to be minimized\n",
    "\n",
    "We can now define the overall loss function to be minimized via gradient descent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3febef8c-2a13-47da-bea5-2978ae604027",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def compute_loss(combination_image, base_image, style_reference_image):\n",
    "    input_tensor = keras.ops.concatenate([base_image, style_reference_image, combination_image], axis=0)\n",
    "    # run all three images through the network at once and extract the dictionary of features\n",
    "    features = feature_extractor(input_tensor)\n",
    "    \n",
    "    # initialize total loss to 0\n",
    "    loss = keras.ops.zeros(shape=())\n",
    "\n",
    "    # add the content loss\n",
    "    layer_features = features[content_loss_layer_name]\n",
    "    base_image_features = layer_features[0, :, :, :]\n",
    "    combination_image_features = layer_features[2, :, :, :]\n",
    "    loss += content_loss_weight * content_loss(base_image_features, combination_image_features)\n",
    "\n",
    "    # add the style loss\n",
    "    num_style_layers = len(style_loss_layer_names)\n",
    "    for layer_name in style_loss_layer_names:\n",
    "        layer_features = features[layer_name]\n",
    "        style_image_features = layer_features[1, :, :, :]\n",
    "        combination_image_features = layer_features[2, :, :, :]\n",
    "        style_loss_value = style_loss(style_image_features, combination_image_features)\n",
    "        loss += style_loss_weight * style_loss_value / num_style_layers\n",
    "\n",
    "    # add the total variation loss\n",
    "    loss += total_variation_loss_weight * total_variation_loss(combination_image)\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### Gradient descent process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "# compiled for speed\n",
    "@tf.function\n",
    "def compute_loss_and_gradients(combination_image, base_image, style_reference_image):\n",
    "    with tf.GradientTape() as tape:\n",
    "        loss = compute_loss(combination_image, base_image, style_reference_image)\n",
    "    grads = tape.gradient(loss, combination_image)\n",
    "    return loss, grads"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5004d952-5986-44bb-8a32-1f8302e1947c"
   },
   "source": [
    "We will use a learning-rate schedule, which will start with a high learning rate and gradually decrease it as the loss is minimized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5fac0e6d-0853-4f3b-9dfc-efb072f782f0",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from keras.optimizers.schedules import ExponentialDecay\n",
    "\n",
    "# decrease learning rate by 4% every 100 steps\n",
    "learning_rate_schedule = ExponentialDecay(\n",
    "    initial_learning_rate=100.0,\n",
    "    decay_steps=100,\n",
    "    decay_rate=0.96\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize the images and run gradient descent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "base_image = preprocess_image(base_image_path)\n",
    "style_reference_image = preprocess_image(style_image_path)\n",
    "combination_image = tf.Variable(preprocess_image(base_image_path))\n",
    "\n",
    "iterations = 250\n",
    "\n",
    "optimizer = keras.optimizers.SGD(learning_rate_schedule)\n",
    "print(\"Generating combination image...\")\n",
    "for i in range(1, iterations+1):\n",
    "    loss, grads = compute_loss_and_gradients(combination_image, base_image, style_reference_image)\n",
    "    optimizer.apply_gradients([(grads, combination_image)])\n",
    "    if i == 1 or i % 10 == 0:\n",
    "        print(f\"iteration {i}: loss = {loss:.2f}\")\n",
    "\n",
    "img = deprocess_image(combination_image.numpy())\n",
    "filename = f\"combination_image_{iterations}.png\"\n",
    "keras.utils.save_img(filename, img)\n",
    "print(f\"Saved final image as {filename}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# display the final image\n",
    "keras.utils.load_img(filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cUCKAR-OrxiJ"
   },
   "outputs": [],
   "source": [
    "# download final image if running in Google Colab\n",
    "try:\n",
    "    from google.colab import files\n",
    "    files.download(filename)\n",
    "except:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "chapter12_part03_neural-style-transfer.i",
   "private_outputs": false,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "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.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
