Style Transfer for Timelapse, see https://ithoughthecamewithyou.com/post/style-transfer-for-time-lapse-photography for instructions

style-transfer-for-timelapse.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Style Transfer for Timelapse.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyNSFwNMHz6Mtpw3zLReIFU+",
      "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/gist/abfo/c2fe8af9976ae09a324f2cc410a86008/style-transfer-for-timelapse.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HEHPFmEEV7Yk"
      },
      "source": [
        "Apply style transfer to a set of images in Google Drive. Based on the style transfer guide [here](https://www.tensorflow.org/hub/tutorials/tf2_arbitrary_image_stylization) and the drive example [here](https://colab.research.google.com/notebooks/io.ipynb#scrollTo=RWSJpsyKqHjH). This code expects a folder in the root of your drive called stylegan. The folder should contain two subfolders, input (images to style), output and a style image called style.jpg. This should be 256x256 pixels. The sample code works with square images and will take the center of the source image if in a different aspect ratio. For my purposes I'm using 4k frames and then setting the output size to 1920x1920 in order to crop a 1920x1080 HD frame from the processed impages.:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CUGH5d1o5jPh"
      },
      "source": [
        "content_img_size = (1920, 1920)\r\n",
        "style_img_size = (256, 256) "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rcj_QQd75oMH"
      },
      "source": [
        "Set up tensorflow and some support functions:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4Qtbh0rXFcCd"
      },
      "source": [
        "# Copyright 2019 The TensorFlow Hub Authors. All Rights Reserved.\r\n",
        "#\r\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\r\n",
        "# you may not use this file except in compliance with the License.\r\n",
        "# You may obtain a copy of the License at\r\n",
        "#\r\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\r\n",
        "#\r\n",
        "# Unless required by applicable law or agreed to in writing, software\r\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\r\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\r\n",
        "# See the License for the specific language governing permissions and\r\n",
        "# limitations under the License.\r\n",
        "# ==============================================================================\r\n",
        "\r\n",
        "from google.colab import drive\r\n",
        "from google.colab import files\r\n",
        "import functools\r\n",
        "import os\r\n",
        "from matplotlib import gridspec\r\n",
        "import matplotlib.pylab as plt\r\n",
        "import numpy as np\r\n",
        "import tensorflow as tf\r\n",
        "import tensorflow_hub as hub\r\n",
        "import os\r\n",
        "import PIL.Image\r\n",
        "\r\n",
        "print(\"TF Version: \", tf.__version__)\r\n",
        "print(\"TF-Hub version: \", hub.__version__)\r\n",
        "print(\"Eager mode enabled: \", tf.executing_eagerly())\r\n",
        "print(\"GPU available: \", tf.test.is_gpu_available())\r\n",
        "\r\n",
        "def crop_center(image):\r\n",
        "  \"\"\"Returns a cropped square image.\"\"\"\r\n",
        "  shape = image.shape\r\n",
        "  new_shape = min(shape[1], shape[2])\r\n",
        "  offset_y = max(shape[1] - shape[2], 0) // 2\r\n",
        "  offset_x = max(shape[2] - shape[1], 0) // 2\r\n",
        "  image = tf.image.crop_to_bounding_box(\r\n",
        "      image, offset_y, offset_x, new_shape, new_shape)\r\n",
        "  return image\r\n",
        "\r\n",
        "@functools.lru_cache(maxsize=None)\r\n",
        "def load_image(image_url, image_size=(256, 256), preserve_aspect_ratio=True):\r\n",
        "  \"\"\"Loads and preprocesses images.\"\"\"\r\n",
        "  # Cache image file locally.\r\n",
        "  #image_path = tf.keras.utils.get_file(os.path.basename(image_url)[-128:], image_url)\r\n",
        "  image_path = image_url\r\n",
        "  # Load and convert to float32 numpy array, add batch dimension, and normalize to range [0, 1].\r\n",
        "  img = plt.imread(image_path).astype(np.float32)[np.newaxis, ...]\r\n",
        "  if img.max() > 1.0:\r\n",
        "    img = img / 255.\r\n",
        "  if len(img.shape) == 3:\r\n",
        "    img = tf.stack([img, img, img], axis=-1)\r\n",
        "  img = crop_center(img)\r\n",
        "  img = tf.image.resize(img, image_size, preserve_aspect_ratio=True)\r\n",
        "  return img\r\n",
        "\r\n",
        "def show_n(images, titles=('',)):\r\n",
        "  n = len(images)\r\n",
        "  image_sizes = [image.shape[1] for image in images]\r\n",
        "  w = (image_sizes[0] * 6) // 320\r\n",
        "  plt.figure(figsize=(w  * n, w))\r\n",
        "  gs = gridspec.GridSpec(1, n, width_ratios=image_sizes)\r\n",
        "  for i in range(n):\r\n",
        "    plt.subplot(gs[i])\r\n",
        "    plt.imshow(images[i][0], aspect='equal')\r\n",
        "    plt.axis('off')\r\n",
        "    plt.title(titles[i] if len(titles) > i else '')\r\n",
        "  plt.show()\r\n",
        "\r\n",
        "def tensor_to_image(tensor):\r\n",
        "  tensor = tensor*255\r\n",
        "  tensor = np.array(tensor, dtype=np.uint8)\r\n",
        "  if np.ndim(tensor)>3:\r\n",
        "    assert tensor.shape[0] == 1\r\n",
        "    tensor = tensor[0]\r\n",
        "  return PIL.Image.fromarray(tensor)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bZwqFmPM43tQ"
      },
      "source": [
        "Load the hub module, this will take a few minutes"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V5KZbbG-47hk"
      },
      "source": [
        "hub_handle = 'https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2'\r\n",
        "hub_module = hub.load(hub_handle)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NM-z5lKXGBQf"
      },
      "source": [
        "Open drive (you'll need to allow access each session, and load the style image)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T9759bbNIp8l"
      },
      "source": [
        "drive.mount('/content/drive', force_remount=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7QWqAgulWd3Z"
      },
      "source": [
        "style_image = load_image('/content/drive/MyDrive/stylegan/style.jpg', style_img_size)\r\n",
        "style_image = tf.nn.avg_pool(style_image, ksize=[3,3], strides=[1,1], padding='SAME')\r\n",
        "show_n([style_image], ['Style image'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GAyE0dNiD44E"
      },
      "source": [
        "Process all the files... This might take a while and sometimes fails so it won't process if the output file alredy exists. You can run this cell as many times as needed to process all the input files (if you get out of memory errors try restarting the runtime and running the cells above again). This will run much faster if you can grab a GPU instance. If you change anything (like the style image) you'll want to clear the output folder and run again. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GK9ifHT4D_s5"
      },
      "source": [
        "directory = '/content/drive/MyDrive/stylegan/input'\r\n",
        "outdirectory = '/content/drive/MyDrive/stylegan/output'\r\n",
        "for file in os.listdir(os.fsencode(directory)):\r\n",
        "     filename = os.fsdecode(file)\r\n",
        "     outputfile = os.path.join(outdirectory, filename)\r\n",
        "     if (os.path.isfile(outputfile)):\r\n",
        "       continue\r\n",
        "     print(filename)\r\n",
        "     content_image = load_image(os.path.join(directory, filename), content_img_size, preserve_aspect_ratio=True)\r\n",
        "     outputs = hub_module(tf.constant(content_image), tf.constant(style_image))\r\n",
        "     stylized_image = outputs[0]\r\n",
        "     tensor_to_image(stylized_image).save(outputfile)\r\n",
        "\r\n",
        "#make sure files are saved when we're done\r\n",
        "print('Done, disconnecting from drive...')\r\n",
        "drive.flush_and_unmount()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}