generate_stimulus.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "generate_stimulus.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "2TW55Tpa5mOW",
        "trsEuU8G6lTY"
      ],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.7.3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2TW55Tpa5mOW",
        "colab_type": "text"
      },
      "source": [
        "## Clone the official repo of style-GAN from Nvidia"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SQ45wgry4898",
        "colab_type": "code",
        "outputId": "43682d4f-79c4-4f3f-d558-c6833449efb9",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 84
        }
      },
      "source": [
        "! git clone https://github.com/NVlabs/stylegan"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into 'stylegan'...\n",
            "remote: Enumerating objects: 83, done.\u001b[K\n",
            "remote: Total 83 (delta 0), reused 0 (delta 0), pack-reused 83\u001b[K\n",
            "Unpacking objects: 100% (83/83), done.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KJz-FYa7NLAO",
        "colab_type": "code",
        "outputId": "b575f5ed-d0c6-479b-ec80-1368ef1d128c",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "# Go to the directory with model\n",
        "%cd stylegan/"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/content/stylegan\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "trsEuU8G6lTY",
        "colab_type": "text"
      },
      "source": [
        "## Load necessary library and the pretrained model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uM0wNzDS_5am",
        "colab_type": "text"
      },
      "source": [
        "Note that the original code was in tensorflow 1.x version so you need to make sure you use version 1.x (not 2.x)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I6IwPQ2t8jCV",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "8d9781e7-cac2-4edd-c9db-01ed6c975b54"
      },
      "source": [
        "%tensorflow_version 1.x"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "TensorFlow 1.x selected.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "HNv5DyXG7mvn",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import pickle\n",
        "import numpy as np\n",
        "import PIL.Image\n",
        "import dnnlib\n",
        "import dnnlib.tflib as tflib\n",
        "import timeit\n",
        "\n",
        "# Check gpu\n",
        "!nvidia-smi\n",
        "\n",
        "#----------------------------------------------------------------------------\n",
        "# Helpers for loading and using pre-trained generators.\n",
        "\n",
        "url_ffhq        = 'https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ' # karras2019stylegan-ffhq-1024x1024.pkl\n",
        "url_celebahq    = 'https://drive.google.com/uc?id=1MGqJl28pN4t7SAtSrPdSRJSQJqahkzUf' # karras2019stylegan-celebahq-1024x1024.pkl\n",
        "url_bedrooms    = 'https://drive.google.com/uc?id=1MOSKeGF0FJcivpBI7s63V9YHloUTORiF' # karras2019stylegan-bedrooms-256x256.pkl\n",
        "url_cars        = 'https://drive.google.com/uc?id=1MJ6iCfNtMIRicihwRorsM3b7mmtmK9c3' # karras2019stylegan-cars-512x384.pkl\n",
        "url_cats        = 'https://drive.google.com/uc?id=1MQywl0FNt6lHu8E_EUqnRbviagS7fbiJ' # karras2019stylegan-cats-256x256.pkl\n",
        "\n",
        "synthesis_kwargs = dict(output_transform=dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True), minibatch_size=8)\n",
        "_Gs_cache = dict()\n",
        "def load_Gs(url):\n",
        "    if url not in _Gs_cache:\n",
        "        with dnnlib.util.open_url(url, num_attempts=10, cache_dir=os.path.join(os.getcwd(), \"cache_dir\")) as f:\n",
        "            _G, _D, Gs = pickle.load(f)\n",
        "        _Gs_cache[url] = Gs\n",
        "    return _Gs_cache[url]\n",
        "\n",
        "# Load the pretrained model for face\n",
        "tflib.init_tf()\n",
        "Gs = load_Gs(url_ffhq)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dXXFK5d85kDG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Mount google drive to download the images (if you use Google cloud service)\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive/')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_pbuUxeg_kI1",
        "colab_type": "text"
      },
      "source": [
        "## Create a set of fake images from random seeds"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-YO2kxxw_0gz",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 87
        },
        "outputId": "8ea205cb-2fb8-4b68-a2d9-cc774498b49c"
      },
      "source": [
        "# Start timestamp\n",
        "start = timeit.timeit()\n",
        "\n",
        "# Create a directory to store images\n",
        "root_dir = '/content/drive/My Drive'\n",
        "result_dir = os.path.join(root_dir, \"result_dir\")\n",
        "os.makedirs(result_dir, exist_ok=True)\n",
        "example_file_name = os.path.join(result_dir, 'example-fake-faces.png')\n",
        "\n",
        "# Image resolution\n",
        "w = 1024\n",
        "h = 1024\n",
        "\n",
        "# Create a set of random integer seeds, each corresponds to one image\n",
        "# Note that the image is high-resolution so your RAM may be overwhelmed if you put n_image too high\n",
        "n_image = 100\n",
        "rand_seeds = list(np.random.randint(low=1, high=3000, size=(n_image,)))\n",
        "\n",
        "# Make the latent input z from the seed\n",
        "rand_latent_z = np.stack(np.random.RandomState(seed).randn(Gs.input_shape[1]) for seed in rand_seeds)\n",
        "\n",
        "# Feed the latent z through the mapping networks to get the latent embedding w\n",
        "# Note that if you don't have gpu, set num_gpus to 0\n",
        "rand_latent_w = Gs.components.mapping.run(rand_latent_z, None) # [seed, layer, component]\n",
        "\n",
        "# Finally feed the latent embedding w to the generator to get the final images\n",
        "rand_images = Gs.components.synthesis.run(rand_latent_w, randomize_noise=False, **synthesis_kwargs)\n",
        "\n",
        "# Save the images\n",
        "n_col = 10\n",
        "n_row = n_image // n_col\n",
        "canvas = PIL.Image.new('RGB', (w * n_col, h * n_row), 'white')\n",
        "count = 0\n",
        "for row in range(n_row):\n",
        "    for col in range(n_col):\n",
        "        canvas.paste(PIL.Image.fromarray(rand_images[count], 'RGB'), (col * w, row * h))\n",
        "        count = count + 1\n",
        "\n",
        "canvas.save(example_file_name)\n",
        "\n",
        "# End timestamp\n",
        "end = timeit.timeit()\n",
        "print(end - start)"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/IPython/core/interactiveshell.py:2822: FutureWarning: arrays to stack must be passed as a \"sequence\" type such as list or tuple. Support for non-sequence iterables such as generators is deprecated as of NumPy 1.16 and will raise an error in the future.\n",
            "  if self.run_code(code, result):\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "0.005844636999881914\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LOyDvDn5b55N",
        "colab_type": "text"
      },
      "source": [
        "## Show example mixing images"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "Ik8wwD3R74GI",
        "colab": {}
      },
      "source": [
        "# Style mixing - mix n_mix_layer of mix n_mix_layer of mixing image to reference image.\n",
        "# Save file name\n",
        "example_file_name = os.path.join(result_dir, 'example-style-mixing.png')\n",
        "\n",
        "# Seeds of good images\n",
        "w = 1024\n",
        "h = 1024\n",
        "style_ranges = [range(8, 18), range(2, 4), range(0, 4)] \n",
        "master_seeds = [701,687,615,2268,888,829,1898,1733,1614,845,1967,1555,1157,1012,91,388,201,500]\n",
        "src_seeds = [701,687,615,2268,888,829,1898,1733,1614,1967,1555,1157,1012,91,388,201,500] # mixing images\n",
        "dst_seeds = [master_seeds[9]] * len(style_ranges) # original image\n",
        "\n",
        "# Create latent z, w and finally the images\n",
        "src_latents = np.stack(np.random.RandomState(seed).randn(Gs.input_shape[1]) for seed in src_seeds)\n",
        "dst_latents = np.stack(np.random.RandomState(seed).randn(Gs.input_shape[1]) for seed in dst_seeds)\n",
        "src_dlatents = Gs.components.mapping.run(src_latents, None) # [seed, layer, component]\n",
        "dst_dlatents = Gs.components.mapping.run(dst_latents, None) # [seed, layer, component]\n",
        "src_images = Gs.components.synthesis.run(src_dlatents, randomize_noise=False, **synthesis_kwargs)\n",
        "dst_images = Gs.components.synthesis.run(dst_dlatents, randomize_noise=False, **synthesis_kwargs)\n",
        "\n",
        "# Mix the latent embedding w and create new mix images\n",
        "canvas = PIL.Image.new('RGB', (w * (len(src_seeds) + 1), h * (len(dst_seeds) + 1)), 'white')\n",
        "for col, src_image in enumerate(list(src_images)):\n",
        "    canvas.paste(PIL.Image.fromarray(src_image, 'RGB'), ((col + 1) * w, 0))\n",
        "for row, dst_image in enumerate(list(dst_images)):\n",
        "    canvas.paste(PIL.Image.fromarray(dst_image, 'RGB'), (0, (row + 1) * h))\n",
        "    row_dlatents = np.stack([dst_dlatents[row]] * len(src_seeds))\n",
        "    row_dlatents[:, style_ranges[row]] = src_dlatents[:, style_ranges[row]]\n",
        "    row_images = Gs.components.synthesis.run(row_dlatents, randomize_noise=False, **synthesis_kwargs)\n",
        "    for col, image in enumerate(list(row_images)):\n",
        "        canvas.paste(PIL.Image.fromarray(image, 'RGB'), ((col + 1) * w, (row + 1) * h))\n",
        "canvas.save(example_file_name)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6q8_my28c1MC",
        "colab_type": "text"
      },
      "source": [
        "## Create the set of mixing images and save"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "qxvu-OBro9IT",
        "colab": {}
      },
      "source": [
        "# Style mixing - mix n_mix_layer of mix n_mix_layer of mixing image to reference image.\n",
        "# Start timestamp\n",
        "start = timeit.timeit()\n",
        "\n",
        "# Seeds of good images\n",
        "w = 1024\n",
        "h = 1024\n",
        "style_ranges = [range(8,18), range(2,4)]\n",
        "master_seeds = [701,687,615,2268,888,829,1898,1733,1614,845,1967,1555,1157,1012,91,388,201,500]\n",
        "mix_seeds = master_seeds[:2] \n",
        "original_seeds = master_seeds[:2]\n",
        "\n",
        "# Create latent z, w and finally the images\n",
        "mix_latents = np.stack(np.random.RandomState(seed).randn(Gs.input_shape[1]) for seed in mix_seeds)\n",
        "original_latents = np.stack(np.random.RandomState(seed).randn(Gs.input_shape[1]) for seed in original_seeds)\n",
        "mix_dlatents = Gs.components.mapping.run(mix_latents, None) # [seed, layer, component]\n",
        "original_dlatents = Gs.components.mapping.run(original_latents, None) # [seed, layer, component]\n",
        "mix_images = Gs.components.synthesis.run(mix_dlatents, randomize_noise=False, **synthesis_kwargs)\n",
        "original_images = Gs.components.synthesis.run(original_dlatents, randomize_noise=False, **synthesis_kwargs)\n",
        "\n",
        "# Mix the latent embedding w and create new mix images\n",
        "for original_ind, original_image in enumerate(list(original_images)):\n",
        "    # Save original image\n",
        "    canvas_original = PIL.Image.new(mode = 'RGB', size = (w, h), color = 'white')\n",
        "    canvas_original.paste(PIL.Image.fromarray(original_image, 'RGB'), (0, 0))\n",
        "    file_name = \"{}_original.png\".format(original_seeds[original_ind])\n",
        "    canvas_original.save(os.path.join(result_dir, file_name))\n",
        "\n",
        "    for style_ind, style in enumerate(list(style_ranges)):\n",
        "        # Mix the images\n",
        "        row_dlatents = np.stack([original_dlatents[original_ind]] * len(mix_seeds))\n",
        "        row_dlatents[:, style] = mix_dlatents[:, style]\n",
        "        row_images = Gs.components.synthesis.run(row_dlatents, randomize_noise=False, **synthesis_kwargs)\n",
        "\n",
        "        # Save the mix images\n",
        "        for variant, image in enumerate(list(row_images)):\n",
        "            canvas_mix = PIL.Image.new('RGB', (w, h), 'white')\n",
        "            canvas_mix.paste(PIL.Image.fromarray(image, 'RGB'), (0, 0))\n",
        "            file_name = \"{}_{}_{}.png\".format(original_seeds[original_ind], style, mix_seeds[variant])\n",
        "            canvas_mix.save(os.path.join(result_dir, file_name))\n",
        "\n",
        "end = timeit.timeit()\n",
        "print(end - start)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}