CycleGAN.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JKSULKknpuvt"
      },
      "source": [
        "# Practice : CycleGAN\n",
        "\n",
        "[CycleGAN official homepage (from authors)] : https://junyanz.github.io/CycleGAN/ <br>\n",
        "[CycleGAN original paper] : https://arxiv.org/abs/1703.10593"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uJ62r3TgWNlU"
      },
      "source": [
        "### Unpaired Image-to-Image Translation\n",
        "\n",
        "> In many cases, there're no paired images between two distributions. So it's quite hard to directly apply paired image-to-image translation algorithm in this case, such as pix2pix. <br>\n",
        "The CycleGAN leverages two GAN architectures(2 generators, 2 discriminators) and cycle consistency loss to deal with this case"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dWAHOeREPAPe"
      },
      "source": [
        "Install tensorflow 2.8.3 <br>\n",
        "(just to avoid the bugs which makes the implementation of data augmentation extremely slow)\n",
        "\n",
        "make sure to install below version of tensorflow, or there might be an error at restoration step"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nuKl4j5WOMCw",
        "outputId": "11b9be4e-51ab-4bac-fc2f-eee251e5e932"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: tensorflow==2.8.3 in /usr/local/lib/python3.9/dist-packages (2.8.3)\n",
            "Requirement already satisfied: tensorflow-estimator<2.9,>=2.8 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (2.8.0)\n",
            "Requirement already satisfied: keras<2.9,>=2.8.0rc0 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (2.8.0)\n",
            "Requirement already satisfied: grpcio<2.0,>=1.24.3 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (1.53.0)\n",
            "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (1.16.0)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (67.6.1)\n",
            "Requirement already satisfied: protobuf<3.20,>=3.9.2 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (3.19.6)\n",
            "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (3.3.0)\n",
            "Requirement already satisfied: numpy>=1.20 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (1.22.4)\n",
            "Requirement already satisfied: keras-preprocessing>=1.1.1 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (1.1.2)\n",
            "Requirement already satisfied: h5py>=2.9.0 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (3.8.0)\n",
            "Requirement already satisfied: astunparse>=1.6.0 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (1.6.3)\n",
            "Requirement already satisfied: libclang>=9.0.1 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (16.0.0)\n",
            "Requirement already satisfied: typing-extensions>=3.6.6 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (4.5.0)\n",
            "Requirement already satisfied: wrapt>=1.11.0 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (1.14.1)\n",
            "Requirement already satisfied: google-pasta>=0.1.1 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (0.2.0)\n",
            "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (2.2.0)\n",
            "Requirement already satisfied: gast>=0.2.1 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (0.4.0)\n",
            "Requirement already satisfied: tensorboard<2.9,>=2.8 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (2.8.0)\n",
            "Requirement already satisfied: absl-py>=0.4.0 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (1.4.0)\n",
            "Requirement already satisfied: tensorflow-io-gcs-filesystem>=0.23.1 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (0.32.0)\n",
            "Requirement already satisfied: flatbuffers>=1.12 in /usr/local/lib/python3.9/dist-packages (from tensorflow==2.8.3) (23.3.3)\n",
            "Requirement already satisfied: wheel<1.0,>=0.23.0 in /usr/local/lib/python3.9/dist-packages (from astunparse>=1.6.0->tensorflow==2.8.3) (0.40.0)\n",
            "Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.9/dist-packages (from tensorboard<2.9,>=2.8->tensorflow==2.8.3) (2.17.1)\n",
            "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.9/dist-packages (from tensorboard<2.9,>=2.8->tensorflow==2.8.3) (2.2.3)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.9/dist-packages (from tensorboard<2.9,>=2.8->tensorflow==2.8.3) (0.4.6)\n",
            "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.9/dist-packages (from tensorboard<2.9,>=2.8->tensorflow==2.8.3) (2.27.1)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.9/dist-packages (from tensorboard<2.9,>=2.8->tensorflow==2.8.3) (3.4.3)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.9/dist-packages (from tensorboard<2.9,>=2.8->tensorflow==2.8.3) (1.8.1)\n",
            "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.9/dist-packages (from tensorboard<2.9,>=2.8->tensorflow==2.8.3) (0.6.1)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.9/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow==2.8.3) (0.2.8)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.9/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow==2.8.3) (4.9)\n",
            "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.9/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow==2.8.3) (5.3.0)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.9/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.9,>=2.8->tensorflow==2.8.3) (1.3.1)\n",
            "Requirement already satisfied: importlib-metadata>=4.4 in /usr/local/lib/python3.9/dist-packages (from markdown>=2.6.8->tensorboard<2.9,>=2.8->tensorflow==2.8.3) (6.1.0)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.9/dist-packages (from requests<3,>=2.21.0->tensorboard<2.9,>=2.8->tensorflow==2.8.3) (1.26.15)\n",
            "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.9/dist-packages (from requests<3,>=2.21.0->tensorboard<2.9,>=2.8->tensorflow==2.8.3) (2.0.12)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.9/dist-packages (from requests<3,>=2.21.0->tensorboard<2.9,>=2.8->tensorflow==2.8.3) (3.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.9/dist-packages (from requests<3,>=2.21.0->tensorboard<2.9,>=2.8->tensorflow==2.8.3) (2022.12.7)\n",
            "Requirement already satisfied: MarkupSafe>=2.1.1 in /usr/local/lib/python3.9/dist-packages (from werkzeug>=0.11.15->tensorboard<2.9,>=2.8->tensorflow==2.8.3) (2.1.2)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.9/dist-packages (from importlib-metadata>=4.4->markdown>=2.6.8->tensorboard<2.9,>=2.8->tensorflow==2.8.3) (3.15.0)\n",
            "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.9/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow==2.8.3) (0.4.8)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.9/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.9,>=2.8->tensorflow==2.8.3) (3.2.2)\n"
          ]
        }
      ],
      "source": [
        "!pip install tensorflow==2.8.3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WEXOkAlhjrl7",
        "outputId": "e7040b96-1b86-4f2b-8a15-e8bf1b33a848"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: tensorflow-addons in /usr/local/lib/python3.9/dist-packages (0.20.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.9/dist-packages (from tensorflow-addons) (23.0)\n",
            "Requirement already satisfied: typeguard<3.0.0,>=2.7 in /usr/local/lib/python3.9/dist-packages (from tensorflow-addons) (2.13.3)\n"
          ]
        }
      ],
      "source": [
        "!pip install tensorflow-addons"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lEwDAoJgCa-m",
        "outputId": "0d404d1a-029d-4df2-fa16-73d0e97a7adf"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Python 3.9.16\n"
          ]
        }
      ],
      "source": [
        "!python --version"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "r06ThFnZM32y",
        "outputId": "b7cd9485-c71e-4c08-c32b-9f85308613db"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mon Apr 10 15:55:25 2023       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   55C    P8    10W /  70W |      0MiB / 15360MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QSV3wtL0sqpi"
      },
      "source": [
        "### Import useful libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bhvaU1zcps6P",
        "outputId": "3a7b350c-aef1-4a92-962f-687d8d524b22"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.9/dist-packages/tensorflow_addons/utils/tfa_eol_msg.py:23: UserWarning: \n",
            "\n",
            "TensorFlow Addons (TFA) has ended development and introduction of new features.\n",
            "TFA has entered a minimal maintenance and release mode until a planned end of life in May 2024.\n",
            "Please modify downstream libraries to take dependencies from other repositories in our TensorFlow community (e.g. Keras, Keras-CV, and Keras-NLP). \n",
            "\n",
            "For more information see: https://github.com/tensorflow/addons/issues/2807 \n",
            "\n",
            "  warnings.warn(\n",
            "/usr/local/lib/python3.9/dist-packages/tensorflow_addons/utils/ensure_tf_install.py:53: UserWarning: Tensorflow Addons supports using Python ops for all Tensorflow versions above or equal to 2.10.0 and strictly below 2.13.0 (nightly versions are not supported). \n",
            " The versions of TensorFlow you are currently using is 2.8.3 and is not supported. \n",
            "Some things might work, some things might not.\n",
            "If you were to encounter a bug, do not file an issue.\n",
            "If you want to make sure you're using a tested and supported configuration, either change the TensorFlow version or the TensorFlow Addons's version. \n",
            "You can find the compatibility matrix in TensorFlow Addon's readme:\n",
            "https://github.com/tensorflow/addons\n",
            "  warnings.warn(\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "import glob\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_addons as tfa\n",
        "from tqdm.auto import tqdm\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wKFiKyT-NIXY",
        "outputId": "5e81ea80-8f08-4e41-831e-21730e5f2c8d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tf_version : 2.8.3\n"
          ]
        }
      ],
      "source": [
        "print(f'tf_version : {tf.__version__}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2yS53AXKtlu0"
      },
      "source": [
        "### Connect to google drive (where the images are in)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PgSONgtatlDi",
        "outputId": "4d98c7c0-c07c-4cd9-f965-e02890fb664c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Drive already mounted at /tmp/drive; to attempt to forcibly remount, call drive.mount(\"/tmp/drive\", force_remount=True).\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/tmp/drive')\n",
        "\n",
        "# base directory\n",
        "base_path = \"/tmp/drive/MyDrive/practice/CycleGAN/\"\n",
        "# change the current working directory to base_path\n",
        "os.chdir(base_path)\n",
        "# images directory\n",
        "image_path = os.path.join(base_path, \"images/\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Yux9xxIC_WLH",
        "outputId": "9a7bf27c-609a-4b52-fa10-f256c618fc0b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "image_directory_A : /tmp/drive/MyDrive/practice/CycleGAN/images/henesys\n",
            "image_directory_B : /tmp/drive/MyDrive/practice/CycleGAN/images/ellinia\n",
            "number of images of category A : 173\n",
            "number of images of category B : 135\n"
          ]
        }
      ],
      "source": [
        "# define category name (string)\n",
        "name_A = 'henesys'\n",
        "name_B = 'ellinia'\n",
        "\n",
        "# image directory for each distribution\n",
        "image_directory_A = os.path.join(image_path, name_A)\n",
        "image_directory_B = os.path.join(image_path, name_B)\n",
        "\n",
        "# just to make sure\n",
        "images_A = glob.glob(image_directory_A + '/*.*')\n",
        "images_B = glob.glob(image_directory_B + '/*.*')\n",
        "\n",
        "num_of_examples_A = len(images_A)\n",
        "num_of_examples_B = len(images_B)\n",
        "\n",
        "# print it\n",
        "print(f'image_directory_A : {image_directory_A}')\n",
        "print(f'image_directory_B : {image_directory_B}')\n",
        "print(f'number of images of category A : {num_of_examples_A}')\n",
        "print(f'number of images of category B : {num_of_examples_B}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lm5_PQ-kKHT1"
      },
      "source": [
        "### Make our datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zfXlVVz3Y_Ca"
      },
      "source": [
        "Note: `Dataset.cache` stores the data from the first epoch and replays it in order. So, using the `cache` method disables any shuffles earlier in the pipeline. Below, `Dataset.shuffle` is added back in after `Dataset.cache`.\n",
        "\n",
        "from https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/tutorials/load_data/csv.ipynb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "82VPyuV0BjkI"
      },
      "outputs": [],
      "source": [
        "BATCH_SIZE = 1\n",
        "AUTOTUNE = tf.data.AUTOTUNE\n",
        "\n",
        "def map_training_images(file):\n",
        "    '''\n",
        "    convert the images to floats and preprocess them\n",
        "    '''\n",
        "\n",
        "    img = tf.io.decode_png(tf.io.read_file(file), channels=3) # decode it as RGB images (3 channels), not RGBA\n",
        "    img = tf.cast(img, tf.float32)\n",
        "    img = img / 127.5 - 1                                     # image tensor should lie in [-1, 1]\n",
        "    img = tf.clip_by_value(img, -1, 1)\n",
        "    \n",
        "    return img\n",
        "\n",
        "\n",
        "def generate_dataset(directory, num_of_examples):\n",
        "    '''\n",
        "    function that return tf.data.Dataset instance containing images in given directory\n",
        "    '''\n",
        "\n",
        "    dataset = tf.data.Dataset.list_files(os.path.join(directory, \"*.*\"))\n",
        "    dataset = dataset.map(map_training_images).cache().shuffle(num_of_examples).batch(BATCH_SIZE).prefetch(AUTOTUNE)\n",
        "\n",
        "    return dataset\n",
        "\n",
        "\n",
        "training_dataset_A = generate_dataset(image_directory_A, num_of_examples_A)\n",
        "training_dataset_B = generate_dataset(image_directory_B, num_of_examples_B)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bhsuUJod3Dty"
      },
      "source": [
        "### Visualize it for test"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zYcWZ_5Q5Ef4"
      },
      "outputs": [],
      "source": [
        "def plot_image(image_tensor, category_name):\n",
        "    '''\n",
        "    function that plot the given image tensor\n",
        "    '''\n",
        "    plt.figure(figsize=(10,10))\n",
        "    plt.grid(False)\n",
        "    plt.title(category_name)\n",
        "\n",
        "    image_tensor = np.squeeze(image_tensor)\n",
        "    # matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers)\n",
        "    # convert it from [-1, 1] to [0, 1]\n",
        "    image_tensor = (image_tensor + 1) / 2.0\n",
        "    \n",
        "    plt.imshow(image_tensor)\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "# # comment below if it's unnecessary\n",
        "# for image_batch in training_dataset_A.take(2):\n",
        "#     print(f'image_batch_shape : {image_batch.shape}')\n",
        "#     plot_image(image_batch, name_A)\n",
        "\n",
        "# for image_batch in training_dataset_B.take(2):\n",
        "#     print(f'image_batch_shape : {image_batch.shape}')\n",
        "#     plot_image(image_batch, name_B)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B1tpVqgzDSwB"
      },
      "source": [
        "## Define our model\n",
        "\n",
        "(followed the architecture of the original paper)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Xi9vGYzDaJ4"
      },
      "source": [
        "As the original authors had noted, since the CycleGAN has 4 models (2 for generators, 2 for discriminators), it is quite memory-intensive <br>\n",
        "<br>\n",
        "So, instead of using the original image, we need to crop the portion of it <br>\n",
        "1. crop the image into size (IMAGE_SIZE, IMAGE_SIZE, 3)\n",
        "2. randomly flip it horizontally"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jZF8obvmDz87"
      },
      "outputs": [],
      "source": [
        "# cropped image\n",
        "IMAGE_SIZE = 480    # it should be multiple of 4, since we will use resnet based generator with 4x downsampling & upsampling\n",
        "                    # to compute the cycle consistency loss later, the size of input & output should match\n",
        "\n",
        "# data augmentation\n",
        "data_augmentation = tf.keras.Sequential([\n",
        "    tf.keras.layers.RandomCrop(IMAGE_SIZE, IMAGE_SIZE),\n",
        "    tf.keras.layers.RandomFlip(mode='horizontal'),\n",
        "])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Y4oBRvsVufw"
      },
      "source": [
        "augmented image looks like this"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rZeQcX4-3BGm"
      },
      "outputs": [],
      "source": [
        "# # comment below if it's unnecessary\n",
        "# for image_batch in training_dataset_A.take(2):\n",
        "#     augmented_batch = data_augmentation(image_batch)\n",
        "#     print(f'augmented_batch_shape : {augmented_batch.shape}')\n",
        "#     plot_image(augmented_batch, name_A)\n",
        "\n",
        "# for image_batch in training_dataset_B.take(2):\n",
        "#     augmented_batch = data_augmentation(image_batch)\n",
        "#     print(f'augmented_batch_shape : {augmented_batch.shape}')\n",
        "#     plot_image(augmented_batch, name_B)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EqsT27axWGgl"
      },
      "source": [
        "### Generator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7clU9SeaWa41"
      },
      "source": [
        "First define our residual block"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2WGb_SAk3Yt3"
      },
      "outputs": [],
      "source": [
        "class Residual_Block(tf.keras.Model):\n",
        "    '''\n",
        "    Residual Block class:\n",
        "        consists of Conv2d - InstanceNorm - Relu - Conv2d - InstanceNorm - Add(Residual Connection)\n",
        "        reflection padding was used to reduce artifacts\n",
        "    '''\n",
        "\n",
        "    def __init__(self, input_channels):\n",
        "        super(Residual_Block, self).__init__()\n",
        "\n",
        "        self.conv_1 = tf.keras.layers.Conv2D(filters=input_channels, kernel_size=3, padding='valid', use_bias=False, \n",
        "                                             kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.02))\n",
        "        self.instance_norm_1 = tfa.layers.InstanceNormalization()\n",
        "\n",
        "        self.conv_2 = tf.keras.layers.Conv2D(filters=input_channels, kernel_size=3, padding='valid', use_bias=False, \n",
        "                                             kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.02))\n",
        "        self.instance_norm_2 = tfa.layers.InstanceNormalization()\n",
        "\n",
        "        self.activation = tf.keras.layers.ReLU()\n",
        "\n",
        "    def reflection_pad(self, input, pad_size):\n",
        "        return tf.pad(input, [[0, 0], [pad_size, pad_size], [pad_size, pad_size], [0, 0]], mode='REFLECT')\n",
        "\n",
        "    def call(self, inputs):\n",
        "        x = self.reflection_pad(inputs, 1)\n",
        "        x = self.conv_1(x)\n",
        "        x = self.instance_norm_1(x)\n",
        "        x = self.activation(x)\n",
        "\n",
        "        x = self.reflection_pad(x, 1)\n",
        "        x = self.conv_2(x)\n",
        "        x = self.instance_norm_2(x)\n",
        "\n",
        "        return x + inputs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ezkNf7lVO58B"
      },
      "source": [
        "Next, define the CycleGAN generator composed of contracting block & 9 residual blocks & expanding block"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M2Cs_sgTaoPO"
      },
      "outputs": [],
      "source": [
        "class CycleGAN_Generator(tf.keras.Model):\n",
        "    '''\n",
        "    CycleGAN Generator class\n",
        "    contracting block + 9 residual blocks + expanding block\n",
        "    '''\n",
        "\n",
        "    def __init__(self, input_channels, output_channels, hidden_channels=64, name=\"\"):\n",
        "        super(CycleGAN_Generator, self).__init__()\n",
        "\n",
        "        if name:\n",
        "            self._name = name\n",
        "\n",
        "        # followed the notation of the paper\n",
        "        # for c7s1-64\n",
        "        self.c7s1_64_conv = tf.keras.layers.Conv2D(filters=hidden_channels, kernel_size=7, padding='valid', use_bias=False, \n",
        "                                                   kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.02))\n",
        "        self.c7s1_64_instance_norm = tfa.layers.InstanceNormalization()\n",
        "\n",
        "        # for d128\n",
        "        self.d128_conv = tf.keras.layers.Conv2D(filters=2*hidden_channels, kernel_size=3, strides=2, padding='valid', use_bias=False, \n",
        "                                                kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.02))\n",
        "        self.d128_instance_norm = tfa.layers.InstanceNormalization()\n",
        "\n",
        "        # for d256\n",
        "        self.d256_conv = tf.keras.layers.Conv2D(filters=4*hidden_channels, kernel_size=3, strides=2, padding='valid', use_bias=False, \n",
        "                                                kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.02))\n",
        "        self.d256_instance_norm = tfa.layers.InstanceNormalization()\n",
        "\n",
        "        # for residual blocks\n",
        "        R_channels = 4*hidden_channels\n",
        "        self.R256_1 = Residual_Block(R_channels)\n",
        "        self.R256_2 = Residual_Block(R_channels)\n",
        "        self.R256_3 = Residual_Block(R_channels)\n",
        "        self.R256_4 = Residual_Block(R_channels)\n",
        "        self.R256_5 = Residual_Block(R_channels)\n",
        "        self.R256_6 = Residual_Block(R_channels)\n",
        "        self.R256_7 = Residual_Block(R_channels)\n",
        "        self.R256_8 = Residual_Block(R_channels)\n",
        "        self.R256_9 = Residual_Block(R_channels)\n",
        "\n",
        "        # for u128\n",
        "        self.u128_conv_transpose = tf.keras.layers.Conv2DTranspose(filters=2*hidden_channels, kernel_size=3, strides=2, padding='same', use_bias=False, \n",
        "                                                                   kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.02))\n",
        "        self.u128_instance_norm = tfa.layers.InstanceNormalization()\n",
        "\n",
        "        # for u256\n",
        "        self.u256_conv_transpose = tf.keras.layers.Conv2DTranspose(filters=hidden_channels, kernel_size=3, strides=2, padding='same', use_bias=False, \n",
        "                                                                   kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.02))\n",
        "        self.u256_instance_norm = tfa.layers.InstanceNormalization()\n",
        "\n",
        "        # for c7s1-3\n",
        "        self.c7s1_3_conv = tf.keras.layers.Conv2D(filters=output_channels, kernel_size=7, padding='valid', \n",
        "                                                  kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.02))\n",
        "\n",
        "        # activation\n",
        "        self.relu = tf.keras.layers.ReLU()\n",
        "\n",
        "\n",
        "    def reflection_pad(self, input, pad_size):\n",
        "        return tf.pad(input, [[0, 0], [pad_size, pad_size], [pad_size, pad_size], [0, 0]], mode='REFLECT')\n",
        "\n",
        "\n",
        "    def call(self, inputs):\n",
        "        # c7s1-64\n",
        "        x = self.reflection_pad(inputs, 3)\n",
        "        x = self.c7s1_64_conv(x)\n",
        "        x = self.c7s1_64_instance_norm(x)\n",
        "        x = self.relu(x)\n",
        "\n",
        "        # d128\n",
        "        x = self.reflection_pad(x, 1)\n",
        "        x = self.d128_conv(x)\n",
        "        x = self.d128_instance_norm(x)\n",
        "        x = self.relu(x)\n",
        "\n",
        "        # d256\n",
        "        x = self.reflection_pad(x, 1)\n",
        "        x = self.d256_conv(x)\n",
        "        x = self.d256_instance_norm(x)\n",
        "        x = self.relu(x)\n",
        "\n",
        "        # R256_1~9\n",
        "        x = self.R256_1(x)\n",
        "        x = self.R256_2(x)\n",
        "        x = self.R256_3(x)\n",
        "        x = self.R256_4(x)\n",
        "        x = self.R256_5(x)\n",
        "        x = self.R256_6(x)\n",
        "        x = self.R256_7(x)\n",
        "        x = self.R256_8(x)\n",
        "        x = self.R256_9(x)\n",
        "\n",
        "        # u128\n",
        "        x = self.u128_conv_transpose(x)\n",
        "        x = self.u128_instance_norm(x)\n",
        "        x = self.relu(x)\n",
        "\n",
        "        # u256\n",
        "        x = self.u256_conv_transpose(x)\n",
        "        x = self.u256_instance_norm(x)\n",
        "        x = self.relu(x)\n",
        "\n",
        "        # c7s1-3\n",
        "        x = self.reflection_pad(x, 3)\n",
        "        x = self.c7s1_3_conv(x)\n",
        "\n",
        "        return tf.keras.activations.tanh(x)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N8Egy2PwtAbS"
      },
      "outputs": [],
      "source": [
        "# temp_generator = CycleGAN_Generator(3, 3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QD1s7c6BtAX_"
      },
      "outputs": [],
      "source": [
        "# for image_batch in training_dataset_A.take(1):\n",
        "#     augmented_batch = data_augmentation(image_batch)\n",
        "#     temp_generated = temp_generator(augmented_batch)\n",
        "#     print(temp_generated.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c_bQ3ZgbtAU6"
      },
      "outputs": [],
      "source": [
        "# temp_generator.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u0VG23bI-HNi"
      },
      "source": [
        "### Discriminator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mYxPGykT_HMJ"
      },
      "source": [
        "PatchGAN distriminator with receptive field : 70 <br>\n",
        "- C64(w/o instance_norm) \n",
        "- C128 \n",
        "- C256 \n",
        "- C512 \n",
        "- (Conv to output 1 channel prediction map)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "000qrKSr-G4x"
      },
      "outputs": [],
      "source": [
        "class CycleGAN_discriminator(tf.keras.Model):\n",
        "    '''\n",
        "    PatchGAN discriminator with receptive field : 70\n",
        "    '''\n",
        "\n",
        "    def __init__(self, input_channels, hidden_channels=64, name=\"\"):\n",
        "        super(CycleGAN_discriminator, self).__init__()\n",
        "\n",
        "        if name:\n",
        "            self._name = name\n",
        "\n",
        "        # for C64\n",
        "        self.C64_conv = tf.keras.layers.Conv2D(filters=hidden_channels, kernel_size=4, strides=2, padding='valid', use_bias=True, \n",
        "                                               kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.02))\n",
        "        \n",
        "        # for C128\n",
        "        self.C128_conv = tf.keras.layers.Conv2D(filters=2*hidden_channels, kernel_size=4, strides=2, padding='valid', use_bias=False, \n",
        "                                                kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.02))\n",
        "        self.C128_instance_norm = tfa.layers.InstanceNormalization()\n",
        "\n",
        "        # for C256\n",
        "        self.C256_conv = tf.keras.layers.Conv2D(filters=4*hidden_channels, kernel_size=4, strides=2, padding='valid', use_bias=False, \n",
        "                                                kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.02))\n",
        "        self.C256_instance_norm = tfa.layers.InstanceNormalization()\n",
        "\n",
        "        # for C512\n",
        "        self.C512_conv = tf.keras.layers.Conv2D(filters=8*hidden_channels, kernel_size=4, strides=1, padding='valid', use_bias=False, \n",
        "                                                kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.02))\n",
        "        self.C512_instance_norm = tfa.layers.InstanceNormalization()\n",
        "\n",
        "        # for output\n",
        "        self.output_conv = tf.keras.layers.Conv2D(filters=1, kernel_size=4, strides=1, padding='valid', use_bias=True, \n",
        "                                                  kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.02))\n",
        "\n",
        "        # activation\n",
        "        self.LeakyReLU = tf.keras.layers.LeakyReLU(alpha=0.2)\n",
        "\n",
        "\n",
        "    def reflection_pad(self, input, pad_size):\n",
        "        return tf.pad(input, [[0, 0], [pad_size, pad_size], [pad_size, pad_size], [0, 0]], mode='REFLECT')\n",
        "\n",
        "    \n",
        "    def call(self, inputs):\n",
        "        # C64\n",
        "        x = self.reflection_pad(inputs, 1)\n",
        "        x = self.C64_conv(x)\n",
        "        # authors didn't used instance norm in the very first C64 block in discriminator\n",
        "        x = self.LeakyReLU(x)\n",
        "\n",
        "        # C128\n",
        "        x = self.reflection_pad(x, 1)\n",
        "        x = self.C128_conv(x)\n",
        "        x = self.C128_instance_norm(x)\n",
        "        x = self.LeakyReLU(x)\n",
        "\n",
        "        # C256\n",
        "        x = self.reflection_pad(x, 1)\n",
        "        x = self.C256_conv(x)\n",
        "        x = self.C256_instance_norm(x)\n",
        "        x = self.LeakyReLU(x)\n",
        "\n",
        "        # C512\n",
        "        x = self.reflection_pad(x, 1)\n",
        "        x = self.C512_conv(x)\n",
        "        x = self.C512_instance_norm(x)\n",
        "        x = self.LeakyReLU(x)\n",
        "\n",
        "        # output\n",
        "        x = self.reflection_pad(x, 1)\n",
        "        x = self.output_conv(x)\n",
        "        return x\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5iDJSvKp90qZ"
      },
      "outputs": [],
      "source": [
        "# temp_discriminator = CycleGAN_discriminator(3)\n",
        "# temp_result = temp_discriminator(temp_generated)\n",
        "# print(temp_result.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BSWvv8ZJ90nh"
      },
      "outputs": [],
      "source": [
        "# temp_discriminator.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "27fJ2xPPIK04"
      },
      "source": [
        "## Define our loss"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ML3K9YqUXY2"
      },
      "source": [
        "### Discriminator Loss\n",
        "\n",
        "- Adversarial Loss (least square loss from LSGAN) <br>\n",
        "https://arxiv.org/abs/1611.04076"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kUoC9Xk_90ki"
      },
      "outputs": [],
      "source": [
        "def discriminator_loss(real_D_out, fake_D_out):\n",
        "    '''\n",
        "    CycleGAN discriminator loss (LSGAN loss)\n",
        "\n",
        "    <params>\n",
        "        real_D_out : discriminator's output given real images from certain distribution\n",
        "        fake_D_out : discriminator's output given fake images generated from the other distribution\n",
        "                     (by putting the image of other distribution to corresponding generator)\n",
        "    '''\n",
        "\n",
        "    # followed the authors, divide it by 2, which slows down the the rate at which D learns\n",
        "    return 0.5 * (tf.math.reduce_mean(tf.math.squared_difference(real_D_out, tf.ones_like(real_D_out))) + \n",
        "                  tf.math.reduce_mean(tf.math.squared_difference(fake_D_out, tf.zeros_like(fake_D_out))))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ejVgEst7joH9"
      },
      "source": [
        "### Generator Loss\n",
        "\n",
        "- Adversarial Loss (least square loss from LSGAN) <br>\n",
        "https://arxiv.org/abs/1611.04076\n",
        "- Cycle Consistency Loss\n",
        "- (optional) Identity Loss"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BS3pdfkrmiTG"
      },
      "source": [
        "Adversarial Loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Yy-3mLKntAQy"
      },
      "outputs": [],
      "source": [
        "def generator_adversarial_loss(fake_D_out):\n",
        "    '''\n",
        "    adversarial loss (LSGAN loss) of CycleGAN generator loss\n",
        "\n",
        "    <params>\n",
        "        fake_D_out : discriminator's output given fake images generated from the other distribution\n",
        "    '''\n",
        "\n",
        "    return tf.math.reduce_mean(tf.math.squared_difference(fake_D_out, tf.ones_like(fake_D_out)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2F1ZEVsxmgv-"
      },
      "source": [
        "Cycle Consistency Loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x4slTf4Trx7U"
      },
      "outputs": [],
      "source": [
        "# weight for cycle consistency loss\n",
        "LAMBDA = 10"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bFoKbRPGtAJT"
      },
      "outputs": [],
      "source": [
        "def generator_cycle_consistency_loss(real_images, cycled_images):\n",
        "    '''\n",
        "    cycle consistency loss of CycleGAN generator loss\n",
        "\n",
        "    <params>\n",
        "        real_images : real images from certain distribution\n",
        "        cycled_images : images generated by putting the above real_images\n",
        "                        to the two generators (in appropriate order)\n",
        "    '''\n",
        "\n",
        "    return tf.math.reduce_mean(tf.math.abs(real_images - cycled_images))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0IA68sEco3We"
      },
      "source": [
        "Omit the Identity Loss"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5UWwd_tjUlnF"
      },
      "source": [
        "## Initialize our Models & Optimizers & Checkpoints"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nB1ddcesiNBI"
      },
      "outputs": [],
      "source": [
        "####### initialize our models #######\n",
        "# generators\n",
        "generator_A_to_B = CycleGAN_Generator(input_channels=3, output_channels=3, name=f'{name_A}2{name_B}_generator')\n",
        "generator_B_to_A = CycleGAN_Generator(input_channels=3, output_channels=3, name=f'{name_B}2{name_A}_generator')\n",
        "\n",
        "# discriminators\n",
        "discriminator_A = CycleGAN_discriminator(input_channels=3, name=f'{name_A}_discriminator')\n",
        "discriminator_B = CycleGAN_discriminator(input_channels=3, name=f'{name_B}_discriminator')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "phX9DPKGXKbW"
      },
      "outputs": [],
      "source": [
        "####### corresponding optimizers #######\n",
        "learning_rate = 2e-4  # 0.0002\n",
        "\n",
        "generator_A_to_B_optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.5, beta_2=0.999)\n",
        "generator_B_to_A_optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.5, beta_2=0.999)\n",
        "\n",
        "discriminator_A_optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.5, beta_2=0.999)\n",
        "discriminator_B_optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.5, beta_2=0.999)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YwKsgICkq-ku"
      },
      "source": [
        "Checkpoints"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nm3QowZ2XKZG",
        "outputId": "8c22c259-778d-4f1d-a6f5-438e3ebdc7bf"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Latest checkpoint restored!!\n"
          ]
        }
      ],
      "source": [
        "# current model name (A2B)\n",
        "cur_model_checkpoint_directory = f\"checkpoints/{name_A}2{name_B}/\"\n",
        "\n",
        "# checkpoint_path\n",
        "checkpoint_path = os.path.join(base_path, cur_model_checkpoint_directory)\n",
        "\n",
        "# checkpoint\n",
        "ckpt = tf.train.Checkpoint(generator_A_to_B=generator_A_to_B,\n",
        "                           generator_B_to_A=generator_B_to_A,\n",
        "                           discriminator_A=discriminator_A,\n",
        "                           discriminator_B=discriminator_B,\n",
        "                           generator_A_to_B_optimizer=generator_A_to_B_optimizer,\n",
        "                           generator_B_to_A_optimizer=generator_B_to_A_optimizer,\n",
        "                           discriminator_A_optimizer=discriminator_A_optimizer,\n",
        "                           discriminator_B_optimizer=discriminator_B_optimizer)\n",
        "\n",
        "# checkpoint manager\n",
        "ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)\n",
        "\n",
        "# if checkpoint exists, restore the latest checkpoint\n",
        "if ckpt_manager.latest_checkpoint:\n",
        "    ckpt.restore(ckpt_manager.latest_checkpoint)\n",
        "    print('Latest checkpoint restored!!')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y3__LeKq5QGw"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WEx1hEVLXKWn"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def train_step(real_A, real_B, LAMBDA):\n",
        "    '''\n",
        "    function that proceeds 1 step of training process\n",
        "\n",
        "    <params>\n",
        "        real_A : batch of real images in A\n",
        "        real_B : batch of real images in B\n",
        "    '''\n",
        "\n",
        "    with tf.GradientTape(persistent=True) as tape:\n",
        "        # fake_A & fake_B\n",
        "        fake_A = generator_B_to_A(real_B)\n",
        "        fake_B = generator_A_to_B(real_A)\n",
        "\n",
        "        # cycled_A & cycled_B\n",
        "        cycled_A = generator_B_to_A(fake_B)\n",
        "        cycled_B = generator_A_to_B(fake_A)\n",
        "\n",
        "        # discriminator's outputs for real & fake images\n",
        "        disc_real_A = discriminator_A(real_A)\n",
        "        disc_real_B = discriminator_B(real_B)\n",
        "\n",
        "        disc_fake_A = discriminator_A(fake_A)\n",
        "        disc_fake_B = discriminator_B(fake_B)\n",
        "\n",
        "        # calculate the loss\n",
        "        # note that the cycled_A & cycled_B used both generator_A_to_B & B_to_A -> precompute it\n",
        "        total_cycle_consistency_loss = generator_cycle_consistency_loss(real_A, cycled_A) + \\\n",
        "                                       generator_cycle_consistency_loss(real_B, cycled_B)\n",
        "\n",
        "        # generator loss\n",
        "        generator_A_to_B_loss = generator_adversarial_loss(disc_fake_B) + \\\n",
        "                                LAMBDA * total_cycle_consistency_loss\n",
        "        generator_B_to_A_loss = generator_adversarial_loss(disc_fake_A) + \\\n",
        "                                LAMBDA * total_cycle_consistency_loss\n",
        "        \n",
        "        # discriminator loss\n",
        "        discriminator_A_loss = discriminator_loss(disc_real_A, disc_fake_A)\n",
        "        discriminator_B_loss = discriminator_loss(disc_real_B, disc_fake_B)\n",
        "\n",
        "    # compute the gradients by backpropagation\n",
        "    # we set the persistent parameter to True above in the tf.GradientTape\n",
        "    # since we're going to calculate gradients more than 1 times\n",
        "    # (otherwise, after one call, the tape will expire)\n",
        "    generator_A_to_B_gradients = tape.gradient(generator_A_to_B_loss, \n",
        "                                               generator_A_to_B.trainable_variables)\n",
        "    generator_B_to_A_gradients = tape.gradient(generator_B_to_A_loss,\n",
        "                                               generator_B_to_A.trainable_variables)\n",
        "    discriminator_A_gradients = tape.gradient(discriminator_A_loss,\n",
        "                                              discriminator_A.trainable_variables)\n",
        "    discriminator_B_gradients = tape.gradient(discriminator_B_loss,\n",
        "                                              discriminator_B.trainable_variables)\n",
        "\n",
        "    # update the weights in the models using optimizer\n",
        "    generator_A_to_B_optimizer.apply_gradients(zip(generator_A_to_B_gradients,\n",
        "                                                   generator_A_to_B.trainable_variables))\n",
        "    generator_B_to_A_optimizer.apply_gradients(zip(generator_B_to_A_gradients,\n",
        "                                                   generator_B_to_A.trainable_variables))\n",
        "    discriminator_A_optimizer.apply_gradients(zip(discriminator_A_gradients,\n",
        "                                                  discriminator_A.trainable_variables))\n",
        "    discriminator_B_optimizer.apply_gradients(zip(discriminator_B_gradients,\n",
        "                                                  discriminator_B.trainable_variables))\n",
        "\n",
        "    # free it\n",
        "    del tape\n",
        "\n",
        "    return generator_A_to_B_loss, generator_B_to_A_loss, discriminator_A_loss, discriminator_B_loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SQ3BJABZXKUF"
      },
      "outputs": [],
      "source": [
        "def plot_images_during_training(real_A, real_B, name_A, name_B):\n",
        "    '''\n",
        "    function that plot the images during training\n",
        "    just to see how it progress\n",
        "    '''\n",
        "    fake_A = generator_B_to_A(real_B)\n",
        "    fake_B = generator_A_to_B(real_A)\n",
        "\n",
        "    plt.figure(figsize=(10, 10))\n",
        "\n",
        "    display_list = [np.squeeze(real_A), np.squeeze(fake_B), np.squeeze(real_B), np.squeeze(fake_A)]\n",
        "    title = [f'real {name_A}', f'fake {name_B}', f'real {name_B}', f'fake {name_A}']\n",
        "\n",
        "    for i in range(4):\n",
        "        plt.subplot(2, 2, i+1)\n",
        "        plt.title(title[i])\n",
        "        # from [-1, 1] to [0, 1]\n",
        "        plt.imshow((display_list[i] + 1) / 2.0)\n",
        "        plt.axis('off')\n",
        "\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5PZbM22o3p8n"
      },
      "source": [
        "Now we could finally train our models"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5lFiDLdIXKQ1"
      },
      "outputs": [],
      "source": [
        "# if we load one & decide to continue training\n",
        "# set the epoch offset\n",
        "epoch_offset = 2000\n",
        "num_of_epochs = 0\n",
        "\n",
        "display_frequency = 1     # in epochs\n",
        "save_frequency = 25       # in epochs\n",
        "\n",
        "# generator_A_to_B_loss, generator_B_to_A_loss, discriminator_A_loss, discriminator_B_loss\n",
        "sum_of_generator_A_to_B_loss = 0\n",
        "sum_of_generator_B_to_A_loss = 0\n",
        "sum_of_discriminator_A_loss = 0\n",
        "sum_of_discriminator_B_loss = 0\n",
        "\n",
        "# since we have the datasets with different number of examples,\n",
        "# the minimum of their size will become the number of steps per epoch\n",
        "steps_per_epoch = min(num_of_examples_A, num_of_examples_B)\n",
        "\n",
        "# training process\n",
        "for epoch in range(epoch_offset, epoch_offset + num_of_epochs):\n",
        "\n",
        "    print(f'================= Epoch {epoch + 1} begins =================')\n",
        "    print()\n",
        "\n",
        "    # for epoch 1000 to 2000, apply linear decay to 0 for learning rate\n",
        "    if epoch >= 1000:\n",
        "        # compute the learning rate of cur epoch\n",
        "        cur_learning_rate = learning_rate - (learning_rate/1000)*(epoch - 1000)\n",
        "        \n",
        "        # update the learning rate of 4 optimizers\n",
        "        tf.keras.backend.set_value(generator_A_to_B_optimizer.learning_rate, cur_learning_rate)\n",
        "        tf.keras.backend.set_value(generator_B_to_A_optimizer.learning_rate, cur_learning_rate)\n",
        "        tf.keras.backend.set_value(discriminator_A_optimizer.learning_rate, cur_learning_rate)\n",
        "        tf.keras.backend.set_value(discriminator_B_optimizer.learning_rate, cur_learning_rate)\n",
        "\n",
        "        print(f'Linearly decay the learning rate => current learning rate : {cur_learning_rate}')\n",
        "    \n",
        "    for real_A_raw, real_B_raw in tqdm(tf.data.Dataset.zip((training_dataset_A, training_dataset_B))):\n",
        "        # augment it\n",
        "        real_A = data_augmentation(real_A_raw)\n",
        "        real_B = data_augmentation(real_B_raw)\n",
        "\n",
        "        # train_step\n",
        "        generator_A_to_B_loss, generator_B_to_A_loss, discriminator_A_loss, discriminator_B_loss = train_step(real_A, real_B, LAMBDA)\n",
        "\n",
        "        # cumulate it\n",
        "        # later we need to divide it by steps_per_epoch\n",
        "        sum_of_generator_A_to_B_loss += generator_A_to_B_loss\n",
        "        sum_of_generator_B_to_A_loss += generator_B_to_A_loss\n",
        "        sum_of_discriminator_A_loss += discriminator_A_loss\n",
        "        sum_of_discriminator_B_loss += discriminator_B_loss\n",
        "\n",
        "    # print the losses\n",
        "    print(f'mean_generator_A_to_B_loss : {sum_of_generator_A_to_B_loss / steps_per_epoch}')\n",
        "    print(f'mean_generator_B_to_A_loss : {sum_of_generator_B_to_A_loss / steps_per_epoch}')\n",
        "    print(f'mean_discriminator_A_loss : {sum_of_discriminator_A_loss / steps_per_epoch}')\n",
        "    print(f'mean_discriminator_B_loss : {sum_of_discriminator_B_loss / steps_per_epoch}')\n",
        "\n",
        "    # updata them for the next epoch\n",
        "    sum_of_generator_A_to_B_loss = 0\n",
        "    sum_of_generator_B_to_A_loss = 0\n",
        "    sum_of_discriminator_A_loss = 0\n",
        "    sum_of_discriminator_B_loss = 0\n",
        "\n",
        "    # display for each display_frequency epochs\n",
        "    if (epoch + 1) % display_frequency == 0:\n",
        "        plot_images_during_training(real_A, real_B, name_A, name_B)\n",
        "\n",
        "    # save the model for each save_frequency epochs\n",
        "    if (epoch + 1) % save_frequency == 0:\n",
        "        ckpt_save_path = ckpt_manager.save()\n",
        "        print('Saving checkpoint for epoch {} at {}'.format(epoch+1, ckpt_save_path))\n",
        "    \n",
        "    print()\n",
        "    print(f'================= Epoch {epoch + 1} ends ===================')\n",
        "\n",
        "    "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ite8IxGxef10"
      },
      "source": [
        "## Application"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0TKt84vXeMHU"
      },
      "source": [
        "- Since the network is fully convolutional, when training is done, we could apply it to arbitrary size image (even the original one)<br>\n",
        "- If width or height is not multiple of 4, then the size of the resulting image might not be the same as that of the original"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0mFAQRRFfuor"
      },
      "source": [
        "### To original image"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sbBxi5SJgBWC"
      },
      "source": [
        "- Simply by skipping the augmentation part"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xc6xssrhXKMn"
      },
      "outputs": [],
      "source": [
        "from PIL import Image\n",
        "\n",
        "# # the number of the original images to print\n",
        "# num_of_original_images_to_print = 10\n",
        "\n",
        "# counter = 1\n",
        "\n",
        "# # take the original images\n",
        "# for real_A, real_B in tqdm(tf.data.Dataset.zip((training_dataset_A, training_dataset_B)).take(num_of_original_images_to_print)):\n",
        "\n",
        "#     # crop the images to make width & height multiple of 4 (otherwise, at the concatenation step, error will occur)\n",
        "#     real_A = real_A[:, :(real_A.shape[1] // 4) * 4, :(real_A.shape[2] // 4) * 4, :]\n",
        "#     real_B = real_B[:, :(real_B.shape[1] // 4) * 4, :(real_B.shape[2] // 4) * 4, :]\n",
        "\n",
        "#     # code block in the function plot_images_during_training\n",
        "#     fake_A = generator_B_to_A(real_B)\n",
        "#     fake_B = generator_A_to_B(real_A)\n",
        "\n",
        "#     plt.figure(figsize=(18, 18))\n",
        "\n",
        "#     squeezed_real_A = np.squeeze(real_A)\n",
        "#     squeezed_fake_B = np.squeeze(fake_B)\n",
        "#     squeezed_real_B = np.squeeze(real_B)\n",
        "#     squeezed_fake_A = np.squeeze(fake_A)\n",
        "\n",
        "#     display_list = [squeezed_real_A, squeezed_fake_B, squeezed_real_B, squeezed_fake_A]\n",
        "#     title = [f'real {name_A}', f'fake {name_B}', f'real {name_B}', f'fake {name_A}']\n",
        "\n",
        "#     for i in range(4):\n",
        "#         plt.subplot(2, 2, i+1)\n",
        "#         plt.title(title[i])\n",
        "#         # from [-1, 1] to [0, 1]\n",
        "#         plt.imshow((display_list[i] + 1) / 2.0)\n",
        "#         plt.axis('off')\n",
        "\n",
        "#     plt.show()\n",
        "\n",
        "#     # concatenate them & download it\n",
        "#     concatenated1 = np.concatenate((squeezed_real_A, squeezed_fake_B), axis=np.argmin(squeezed_real_A.shape[:2]))\n",
        "#     concatenated2 = np.concatenate((squeezed_real_B, squeezed_fake_A), axis=np.argmin(squeezed_real_B.shape[:2]))\n",
        "\n",
        "#     # from [-1, 1] to [0, 255]\n",
        "#     concatenated1 = ((concatenated1 + 1) * 127.5).astype(np.uint8)\n",
        "#     concatenated2 = ((concatenated2 + 1) * 127.5).astype(np.uint8)\n",
        "\n",
        "#     im1 = Image.fromarray(concatenated1, mode=\"RGB\")\n",
        "#     im2 = Image.fromarray(concatenated2, mode=\"RGB\")\n",
        "\n",
        "#     # since our current working directory base_path (in the google drive)\n",
        "#     # we could see the result there\n",
        "#     im1.save(f\"{name_A}2{name_B} result{counter}.jpeg\")\n",
        "#     im2.save(f\"{name_B}2{name_A} result{counter}.jpeg\")\n",
        "\n",
        "#     counter += 1"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### To test image\n"
      ],
      "metadata": {
        "id": "iAi450U9wgIu"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "- test it for the image which hasn't been used in training\n",
        "- use the same scale (in pixels) as in the game (not necessary, but recommended)\n",
        "- to make the test distribution as similar as possible to training one"
      ],
      "metadata": {
        "id": "DIG0NZgHLU00"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# test images directory\n",
        "test_image_path = os.path.join(base_path, \"test_images/\")\n",
        "\n",
        "# test image directory for each distribution\n",
        "test_image_directory_A = os.path.join(test_image_path, name_A)\n",
        "test_image_directory_B = os.path.join(test_image_path, name_B)\n",
        "\n",
        "# just to make sure\n",
        "test_images_A = glob.glob(test_image_directory_A + '/*.*')\n",
        "test_images_B = glob.glob(test_image_directory_B + '/*.*')\n",
        "\n",
        "num_of_test_examples_A = len(test_images_A)\n",
        "num_of_test_examples_B = len(test_images_B)\n",
        "\n",
        "# print it\n",
        "print(f'test_image_directory_A : {test_image_directory_A}')\n",
        "print(f'test_image_directory_B : {test_image_directory_B}')\n",
        "print(f'number of test images of category A : {num_of_test_examples_A}')\n",
        "print(f'number of test images of category B : {num_of_test_examples_B}')\n",
        "\n",
        "test_dataset_A = generate_dataset(test_image_directory_A, num_of_test_examples_A)\n",
        "test_dataset_B = generate_dataset(test_image_directory_B, num_of_test_examples_B)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ds2H8_pOWUCI",
        "outputId": "b9634db2-67f1-4d77-aac8-dde7795f2d1a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "test_image_directory_A : /tmp/drive/MyDrive/practice/CycleGAN/test_images/henesys\n",
            "test_image_directory_B : /tmp/drive/MyDrive/practice/CycleGAN/test_images/ellinia\n",
            "number of test images of category A : 6\n",
            "number of test images of category B : 5\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# # A to B\n",
        "\n",
        "# counter = 1\n",
        "\n",
        "# for real_A in tqdm(test_dataset_A):\n",
        "#     real_A = real_A[:, :(real_A.shape[1] // 4) * 4, :(real_A.shape[2] // 4) * 4, :]\n",
        "\n",
        "#     # generate the fake image\n",
        "#     fake_B = generator_A_to_B(real_A)\n",
        "\n",
        "#     plt.figure(figsize=(18, 18))\n",
        "\n",
        "#     # squeeze the batch dimension\n",
        "#     squeezed_real_A = np.squeeze(real_A)\n",
        "#     squeezed_fake_B = np.squeeze(fake_B)\n",
        "\n",
        "#     display_list = [squeezed_real_A, squeezed_fake_B]\n",
        "#     title = [f'real {name_A}', f'fake {name_B}']\n",
        "\n",
        "#     for i in range(2):\n",
        "#         plt.subplot(1, 2, i+1)\n",
        "#         plt.title(title[i])\n",
        "#         # from [-1, 1] to [0, 1]\n",
        "#         plt.imshow((display_list[i] + 1) / 2.0)\n",
        "#         plt.axis('off')\n",
        "\n",
        "#     plt.show()\n",
        "\n",
        "#     # concatenate them & download it\n",
        "#     concatenated = np.concatenate((squeezed_real_A, squeezed_fake_B), axis=np.argmin(squeezed_real_A.shape[:2]))\n",
        "\n",
        "#     # from [-1, 1] to [0, 255]\n",
        "#     concatenated = ((concatenated + 1) * 127.5).astype(np.uint8)\n",
        "\n",
        "#     im = Image.fromarray(concatenated, mode=\"RGB\")\n",
        "#     image_name = f\"{name_A}2{name_B} result{counter}.jpeg\"\n",
        "#     im.save(image_name)\n",
        "\n",
        "#     counter += 1"
      ],
      "metadata": {
        "id": "yoZBHAWawgZe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# # B to A\n",
        "\n",
        "# counter = 1\n",
        "\n",
        "# for real_B in tqdm(test_dataset_B):\n",
        "#     real_B = real_B[:, :(real_B.shape[1] // 4) * 4, :(real_B.shape[2] // 4) * 4, :]\n",
        "\n",
        "#     # generate the fake image\n",
        "#     fake_A = generator_B_to_A(real_B)\n",
        "\n",
        "#     plt.figure(figsize=(18, 18))\n",
        "\n",
        "#     # squeeze the batch dimension\n",
        "#     squeezed_real_B = np.squeeze(real_B)\n",
        "#     squeezed_fake_A = np.squeeze(fake_A)\n",
        "\n",
        "#     display_list = [squeezed_real_B, squeezed_fake_A]\n",
        "#     title = [f'real {name_B}', f'fake {name_A}']\n",
        "\n",
        "#     for i in range(2):\n",
        "#         plt.subplot(1, 2, i+1)\n",
        "#         plt.title(title[i])\n",
        "#         # from [-1, 1] to [0, 1]\n",
        "#         plt.imshow((display_list[i] + 1) / 2.0)\n",
        "#         plt.axis('off')\n",
        "\n",
        "#     plt.show()\n",
        "\n",
        "#     # concatenate them & download it\n",
        "#     concatenated = np.concatenate((squeezed_real_B, squeezed_fake_A), axis=np.argmin(squeezed_real_B.shape[:2]))\n",
        "\n",
        "#     # from [-1, 1] to [0, 255]\n",
        "#     concatenated = ((concatenated + 1) * 127.5).astype(np.uint8)\n",
        "\n",
        "#     im = Image.fromarray(concatenated, mode=\"RGB\")\n",
        "#     image_name = f\"{name_B}2{name_A} result{counter}.jpeg\"\n",
        "#     im.save(image_name)\n",
        "\n",
        "#     counter += 1"
      ],
      "metadata": {
        "id": "oTK52rJdMs9g"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mr02jRfDf34s"
      },
      "source": [
        "### To video (frame by frame)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8kRuAKzzi2mr"
      },
      "source": [
        "https://www.tensorflow.org/tutorials/load_data/video\n",
        "<br>\n",
        "- during processing, you will likely to encounter out of memory error\n",
        "- then consider reducing the number of frames\n",
        "- or free the unnecessary variables\n",
        "- or restart runtime & proceed again\n",
        "- the generated gif file would likely to have large file size => try compressing it"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# !pip install scikit-video"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "q8--XrA_vFeq",
        "outputId": "ac114bdb-1834-4976-a534-efed32899aab"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: scikit-video in /usr/local/lib/python3.9/dist-packages (1.1.11)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.9/dist-packages (from scikit-video) (1.10.1)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.9/dist-packages (from scikit-video) (1.22.4)\n",
            "Requirement already satisfied: pillow in /usr/local/lib/python3.9/dist-packages (from scikit-video) (8.4.0)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b-QWZFPQXKG_",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "caecc600-84f6-4d33-e3a8-fb640aa4b5ec"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "video_path : /tmp/drive/MyDrive/practice/CycleGAN/videos/where the forest sings.mp4\n",
            "\n",
            "type of video_frames : <class 'numpy.ndarray'>\n",
            "shape of video_frames : (52, 776, 1376, 3)\n"
          ]
        }
      ],
      "source": [
        "# import skvideo.io  \n",
        "\n",
        "# def preprocess_frames(frames):\n",
        "#     '''\n",
        "#     function that preprocess the video frames\n",
        "\n",
        "#     <params>\n",
        "#         frames : video frames (shape : [number_of_frames, height, width, channels])\n",
        "#     '''\n",
        "\n",
        "#     # frames might be considered as image batch\n",
        "#     frames = frames.astype(np.float32)\n",
        "#     frames = frames / 127.5 - 1\n",
        "#     frames = np.clip(frames, -1, 1)\n",
        "\n",
        "#     return frames\n",
        "\n",
        "# # set the video_name & the generator_model that we want to apply\n",
        "# video_name = \"where the forest sings.mp4\"\n",
        "# generator_model = generator_B_to_A\n",
        "\n",
        "# video_directory = os.path.join(base_path, 'videos/')\n",
        "# video_path = os.path.join(video_directory, video_name)\n",
        "# print(f'video_path : {video_path}')\n",
        "\n",
        "# # extract frames\n",
        "# video_frames = skvideo.io.vread(video_path)[::8, :, :, :]\n",
        "# video_frames = preprocess_frames(video_frames)\n",
        "# print()\n",
        "# print(f'type of video_frames : {type(video_frames)}')\n",
        "# print(f'shape of video_frames : {video_frames.shape}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4s_tUhO_XJ-1"
      },
      "outputs": [],
      "source": [
        "# # generate the fake videos\n",
        "# generated_frames = generator_model.predict(video_frames, batch_size=8)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# # concatenate them vertically\n",
        "# concatenated_frames = tf.concat([video_frames, generated_frames], axis=1)\n",
        "# print(f'concatenated shape : {concatenated_frames.shape}')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-lIG3jDTwW6S",
        "outputId": "0d8adeeb-b875-4b01-db0a-3ee6acf08d48"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "concatenated shape : (52, 1552, 1376, 3)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# !pip install -q git+https://github.com/tensorflow/docs"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jR5b4Ndi1VxR",
        "outputId": "50836deb-3d3c-4d92-c186-52282e7254e2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# # https://www.tensorflow.org/tutorials/load_data/video#create_frames_from_each_video_file\n",
        "# import imageio\n",
        "# from tensorflow_docs.vis import embed\n",
        "\n",
        "# def to_gif(images):\n",
        "#     converted_images = np.clip((images + 1) * 127.5, 0, 255).astype(np.uint8)\n",
        "#     imageio.mimsave('./animation.gif', converted_images, fps=5)\n",
        "#     return embed.embed_file('./animation.gif')\n",
        "\n",
        "# to_gif(concatenated_frames)"
      ],
      "metadata": {
        "id": "FUTvcUc4wW3B"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MX-0b3nal7ZX"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "machine_shape": "hm",
      "toc_visible": true,
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}