Colorization_Transformer.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Colorization Transformer.ipynb",
      "private_outputs": true,
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/dvschultz/ml-art-colabs/blob/master/Colorization_Transformer.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WsYJ6BLXq6St"
      },
      "source": [
        "# Colorization Transformer\n",
        "\n",
        "Source code accompanying the paper [Colorization Transformer](https://openreview.net/forum?id=5NA1PinlGFu) to be presented at\n",
        "ICLR 2021.\n",
        "Work by Manoj Kumar, Dirk Weissenborn and Nal Kalchbrenner.\n",
        "\n",
        "\n",
        "<img src=\"https://github.com/google-research/google-research/raw/master/coltran/coltran_images.png\" alt=\"Model images\" width=\"1000\"/>\n",
        "\n",
        "### Important: Note on Evaluation:\n",
        "\n",
        "Colorization is a multi-modal problem, where a given object can have multiple plausible colors.\n",
        "On recolorizing a RGB image, a single random sample represents another plausible colorization. It\n",
        "is unlikely to have the same colors as the ground-truth RGB image. So while evaluating the Colorization Transformer,\n",
        "one should use distribution-level metrics such as the FID and NOT per-image metrics such as LPIPS or SSIM.\n",
        "\n",
        "\n",
        "## Paper summary\n",
        "\n",
        "\n",
        "ColTran consists of three components, a `colorizer`, `color upsampler` and `spatial upsampler`.\n",
        "\n",
        "<img src=\"https://github.com/google-research/google-research/raw/master/coltran/coltran.png\" alt=\"Model figure\" width=\"800\"/>\n",
        "\n",
        "* The `colorizer` is an autoregressive, self-attention based architecture comprised of conditional transformer layers. It coarsely colorizes low resolution `64x64` grayscale images pixel-by-pixel.\n",
        "* The `color upsampler` is a parallel, deterministic self-attention based network. It refines a coarse low resolution image\n",
        "into a `64x64` RGB image.\n",
        "* The architecture of the `spatial upsampler` is similar to the `color upsampler`. It superesolves the low resolution RGB image into the final output.\n",
        "* The `colorizer` has an auxiliary parallel color prediction model composed of a single linear layer.\n",
        "\n",
        "We report results after training each of these individual components on `4x4 TPUv2` chips. Please adjust the model size and batch-size while training using fewer resources. Our results on training these components using lower resources are available in the [appendix](https://openreview.net/pdf?id=5NA1PinlGFu).\n",
        "\n",
        "Full configurations used to train models in the paper are available in the directory `configs`. Configs for extremely small models are provided at `test_configs` to test that the model build quickly. Set the flag `--steps_per_summaries=100` to output logs quickly. When sampling, set `config.sample.log_dir` to an\n",
        "appropriate write directory."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kfwpb0Us2czD"
      },
      "source": [
        "!nvidia-smi -L"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uwqu4svZsA6_"
      },
      "source": [
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jHAsrV7WsC_K",
        "cellView": "form"
      },
      "source": [
        "#@title Clone repo and install dependencies\n",
        "%cd /content/\n",
        "!git clone --recurse-submodules https://github.com/google-research/google-research.git\n",
        "%cd google-research/coltran/\n",
        "!pip install -r requirements.txt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O5mL1yo-tDRQ",
        "cellView": "form"
      },
      "source": [
        "#@title Download pre-trained model\n",
        "!wget https://storage.googleapis.com/gresearch/coltran/coltran.zip -O coltran.zip\n",
        "!unzip coltran.zip -d ./logdir"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M-lHwKr5up6-"
      },
      "source": [
        "# Image colorisation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DQ-XmsHkuy4I"
      },
      "source": [
        "#@title Set environment variables\n",
        " \n",
        "import os\n",
        " \n",
        "os.environ[\"LOGDIR\"] = \"/content/google-research/coltran/logdir/coltran\"\n",
        "!mkdir -p /content/input/ /content/output/\n",
        "os.environ['IMG_DIR'] = \"/content/input\" \n",
        "os.environ['STORE_DIR'] = \"/content/output\"\n",
        " \n",
        "mode = \"colorize\" #@param [\"recolorize\", \"colorize\"]\n",
        "os.environ['MODE'] = mode"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "toLwtBQGwCgY"
      },
      "source": [
        "#@title Run the colorizer to get a coarsely colorized image.\n",
        "%cd /content/google-research/\n",
        "!rm -rf $IMG_DIR/.ipynb_checkpoints/\n",
        "!python -m coltran.custom_colorize --config=coltran/configs/colorizer.py --logdir=$LOGDIR/colorizer --img_dir=$IMG_DIR --store_dir=$STORE_DIR --mode=$MODE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "fnao3UxPxpWL"
      },
      "source": [
        "#@title Run the color upsampler to upsample the coarsely colored image.\n",
        "%cd /content/google-research\n",
        "!rm -rf $IMG_DIR/.ipynb_checkpoints/ $STORE_DIR/stage1/.ipynb_checkpoints\n",
        "!python -m coltran.custom_colorize --config=coltran/configs/color_upsampler.py --logdir=$LOGDIR/color_upsampler --img_dir=$IMG_DIR --store_dir=$STORE_DIR --gen_data_dir=$STORE_DIR/stage1 --mode=$MODE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "TyhE_Yzpx958"
      },
      "source": [
        "#@title Run the spatial upsampler to super-resolve into the final output.\n",
        "%cd /content/google-research/\n",
        "!rm -rf $IMG_DIR/.ipynb_checkpoints/ $STORE_DIR/stage2/.ipynb_checkpoints\n",
        "!python -m coltran.custom_colorize --config=coltran/configs/spatial_upsampler.py --logdir=$LOGDIR/spatial_upsampler --img_dir=$IMG_DIR --store_dir=$STORE_DIR --gen_data_dir=$STORE_DIR/stage2 --mode=$MODE"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}