EvgenyKashin_Animal_conditional_generation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "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.6.8"
    },
    "colab": {
      "name": "EvgenyKashin-Animal_conditional_generation.ipynb",
      "provenance": [],
      "private_outputs": true,
      "toc_visible": true,
      "include_colab_link": true
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Norod/my-colab-experiments/blob/master/EvgenyKashin_Animal_conditional_generation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G2zqgaHLgWyu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%cd /content/\n",
        "!mkdir pretrained\n",
        "\n",
        "!wget 'https://github.com/EvgenyKashin/stylegan2/releases/download/v1.0.0/network-snapshot-005532.pkl' -O ./pretrained/network.pkl\n",
        "\n",
        "!ls -latr ./pretrained\n",
        "\n",
        "!mkdir animations\n",
        "\n",
        "!apt-get install imagemagick\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o2bPF6ZIgmOx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%cd /content/\n",
        "%tensorflow_version 1.x\n",
        "import tensorflow as tf\n",
        "\n",
        "# Download the code\n",
        "!git clone https://github.com/EvgenyKashin/stylegan2\n",
        "%cd /content/stylegan2\n",
        "!nvcc test_nvcc.cu -o test_nvcc -run\n",
        "\n",
        "print('Tensorflow version: {}'.format(tf.__version__) )\n",
        "!nvidia-smi -L\n",
        "print('GPU Identified at: {}'.format(tf.test.gpu_device_name()))\n",
        "\n",
        "%cd /content/stylegan2/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "reWnxuHwfTTR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import ipywidgets as widgets\n",
        "import pretrained_networks\n",
        "import PIL.Image\n",
        "import numpy as np\n",
        "import dnnlib\n",
        "import dnnlib.tflib as tflib\n",
        "\n",
        "network_pkl = '/content/pretrained/network.pkl'\n",
        "_G, _D, Gs = pretrained_networks.load_networks(network_pkl)\n",
        "\n",
        "Gs_syn_kwargs = dnnlib.EasyDict()\n",
        "batch_size = 1\n",
        "Gs_syn_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)\n",
        "Gs_syn_kwargs.randomize_noise = True\n",
        "Gs_syn_kwargs.minibatch_size = batch_size"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6SjTpaVgfTTe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def display_sample_conditional(cat, dog, wild, seed, truncation, return_img=False):\n",
        "    batch_size = 1\n",
        "    l1 = np.zeros((1,3))\n",
        "    l1[0][0] = cat\n",
        "    l1[0][1] = dog\n",
        "    l1[0][2] = wild\n",
        "\n",
        "    all_seeds = [seed] * batch_size\n",
        "    all_z = np.stack([np.random.RandomState(seed).randn(*Gs.input_shape[1:]) for seed in all_seeds])\n",
        "    all_w = Gs.components.mapping.run(all_z, np.tile(l1, (batch_size, 1))) # [minibatch, layer, component]\n",
        "    if truncation != 1:\n",
        "        w_avg = Gs.get_var('dlatent_avg')\n",
        "        all_w = w_avg + (all_w - w_avg) * truncation # [minibatch, layer, component]\n",
        "    all_images = Gs.components.synthesis.run(all_w, **Gs_syn_kwargs)\n",
        "    if return_img:\n",
        "        return PIL.Image.fromarray(np.median(all_images, axis=0).astype(np.uint8))\n",
        "    else:\n",
        "        display(PIL.Image.fromarray(np.median(all_images, axis=0).astype(np.uint8)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lY0RSdyAfTTn",
        "colab_type": "text"
      },
      "source": [
        "## Conditional generation of animals"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z6RUViIAfTTo",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "animal = widgets.Dropdown(\n",
        "    options=[('Cat', 0), ('Dog', 1), ('Wild', 2)],\n",
        "    value=0,\n",
        "    description='Animal: '\n",
        ")\n",
        "\n",
        "seed = widgets.IntSlider(min=0, max=100000, step=1, value=0, description='Seed: ')\n",
        "truncation = widgets.FloatSlider(min=0, max=1, step=0.1, value=1, description='Truncation: ')\n",
        "\n",
        "top_box = widgets.HBox([animal])\n",
        "bot_box = widgets.HBox([seed, truncation])\n",
        "ui = widgets.VBox([top_box, bot_box])\n",
        "\n",
        "def display_animal(animal, seed, truncation):\n",
        "    cat = (animal == 0)\n",
        "    dog = (animal == 1)\n",
        "    wild = (animal == 2)\n",
        "    display_sample_conditional(cat, dog, wild, seed, truncation)\n",
        "\n",
        "out = widgets.interactive_output(display_animal, {'animal': animal, 'seed': seed, 'truncation': truncation})\n",
        "\n",
        "display(ui, out)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3gPeZ0PhfTTv",
        "colab_type": "text"
      },
      "source": [
        "## Mixed generation of animal"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JdygfLMRfTTw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "cat = widgets.FloatSlider(min=0, max=1, step=0.05, value=1, description='Cat: ')\n",
        "dog = widgets.FloatSlider(min=0, max=1, step=0.05, value=0, description='Dog: ')\n",
        "wild = widgets.FloatSlider(min=0, max=1, step=0.05, value=0, description='Wild: ')\n",
        "\n",
        "top_box = widgets.HBox([cat, dog, wild])\n",
        "bot_box = widgets.HBox([seed, truncation])\n",
        "ui = widgets.VBox([top_box, bot_box])\n",
        "\n",
        "\n",
        "out = widgets.interactive_output(display_sample_conditional,\n",
        "                                 {'cat': cat, 'dog': dog, 'wild': wild,\n",
        "                                  'seed': seed, 'truncation': truncation})\n",
        "\n",
        "display(ui, out)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F9VQIWMPfTT9",
        "colab_type": "text"
      },
      "source": [
        "## Transition between labels"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tioYhgU2fTT-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "direction = widgets.Dropdown(\n",
        "    options=['cat2wild', 'cat2dog', 'dog2wild'],\n",
        "    value='cat2wild',\n",
        "    description='Animal: '\n",
        ")\n",
        "value = widgets.FloatSlider(min=0, max=1, step=0.05, value=1, description='Value: ')\n",
        "\n",
        "\n",
        "top_box = widgets.HBox([direction, value])\n",
        "bot_box = widgets.HBox([seed, truncation])\n",
        "ui = widgets.VBox([top_box, bot_box])\n",
        "\n",
        "def display_transition(direction, value, truncation, seed, return_img=False):\n",
        "    if direction == 'cat2wild':\n",
        "        wild = value\n",
        "        cat = 1 - value\n",
        "        dog = 0\n",
        "    elif direction == 'cat2dog':\n",
        "        dog = value\n",
        "        cat = 1 - value\n",
        "        wild = 0\n",
        "    elif direction == 'dog2wild':\n",
        "        wild = value\n",
        "        dog = 1 - value\n",
        "        cat = 0\n",
        "    else:\n",
        "        raise ValueError('Wrong direction value')\n",
        "    \n",
        "    if return_img:\n",
        "        return display_sample_conditional(cat, dog, wild, seed, truncation, return_img)\n",
        "    else:\n",
        "        display_sample_conditional(cat, dog, wild, seed, truncation, return_img)\n",
        "\n",
        "out = widgets.interactive_output(display_transition, {'direction': direction, 'value': value,\n",
        "                                                      'seed': seed, 'truncation': truncation})\n",
        "\n",
        "display(ui, out)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I7XJumGpfTUF",
        "colab_type": "text"
      },
      "source": [
        "## Save images for animation with imagemagick and ffmpeg"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "41g0hLM_fTUG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "output_img_dir = '/content/animations/' + str(seed.value) + '-' + str(truncation.value)\n",
        "!mkdir \"$output_img_dir\"\n",
        "\n",
        "imgs = [display_transition(direction.value, i, truncation.value, seed.value, return_img=True) for\n",
        "        i in np.linspace(0, 1, 31)]\n",
        "\n",
        "for i, im in enumerate(imgs):\n",
        "    im.save(f\"{output_img_dir}/{i:03}.jpg\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HRNVhEXZfTUK",
        "colab_type": "text"
      },
      "source": [
        "Generate gif animation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Xd4kzoXKfTUL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "output_gif_file = '/content/animations/' + str(seed.value) + '-' + str(truncation.value) + '-animation.gif'\n",
        "!convert -delay 10 -layers optimize \"$output_img_dir/*.jpg\" \"$output_gif_file\""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "FlFOcesZo7ag"
      },
      "source": [
        "Generate video animation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4_agwmh4nsrU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "output_vid_file = '/content/animations/' + str(seed.value) + '-' + str(truncation.value) + '-video.mp4'\n",
        "!ffmpeg -f image2  -framerate 8 -i \"$output_img_dir/%03d.jpg\" -b:v 8192k -r 30 -y -c:v libx264 \"$output_vid_file\""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P3pwy-1dp1_O",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(\"Your gif is available at \" + output_gif_file)\n",
        "print(\"Your video is available at \" + output_vid_file)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}