StyleGAN2_projection_interpolation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "StyleGAN2 projection interpolation",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "290168f9058c4abaa846d807ce2984cc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "VBoxModel",
          "state": {
            "_view_name": "VBoxView",
            "_dom_classes": [],
            "_model_name": "VBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_2332ec1b25cd4ca8a0c671e7884e7225",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_5be1c18f016744d1903247cd85d65dbc",
              "IPY_MODEL_54246974c1ec45a48a6064b1517c046b"
            ]
          }
        },
        "2332ec1b25cd4ca8a0c671e7884e7225": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "5be1c18f016744d1903247cd85d65dbc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_b1327f73da28457a8606204a30bf7d77",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": "Creating animation: 100",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_11bbe6b7111e42bbabd18bc86a2e49d8"
          }
        },
        "54246974c1ec45a48a6064b1517c046b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "IntProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_0e7bb66324934d58bf7c058f35fe629e",
            "_dom_classes": [],
            "description": "",
            "_model_name": "IntProgressModel",
            "bar_style": "success",
            "max": 100,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 100,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_af4b3f75693241eca70a86ed02e2362c"
          }
        },
        "b1327f73da28457a8606204a30bf7d77": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "11bbe6b7111e42bbabd18bc86a2e49d8": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "0e7bb66324934d58bf7c058f35fe629e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "af4b3f75693241eca70a86ed02e2362c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/dvschultz/stylegan2/blob/master/StyleGAN2_projection_interpolation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4_s8h-ilzHQc",
        "colab_type": "text"
      },
      "source": [
        "# StyleGAN2 interpolation\n",
        "\n",
        "This notebook demonstrates how to run NVIDIA's StyleGAN2 on Google Colab.\n",
        "Make sure to specify a GPU runtime.\n",
        "\n",
        "Based on previous work by Mikael Christensen, 2019.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SPF0iQFTSTKh",
        "colab_type": "text"
      },
      "source": [
        "## Set up StyleGAN2"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PzDuIoMcqfBT",
        "colab_type": "code",
        "outputId": "7a217cc4-9bba-4941-bf5a-18a8ea826d10",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 513
        }
      },
      "source": [
        "%tensorflow_version 1.x\n",
        "import tensorflow as tf\n",
        "\n",
        "# Download the code\n",
        "!git clone https://github.com/dvschultz/stylegan2\n",
        "%cd 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",
        "!pip install gdown\n",
        "%mkdir pkl\n",
        "%cd pkl\n",
        "!gdown --id 1gbxwfHNOaGjGsLNTmmSrNA85X2VWHHOq\n",
        "%cd ../"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "TensorFlow 1.x selected.\n",
            "Cloning into 'stylegan2'...\n",
            "remote: Enumerating objects: 3, done.\u001b[K\n",
            "remote: Counting objects: 100% (3/3), done.\u001b[K\n",
            "remote: Compressing objects: 100% (3/3), done.\u001b[K\n",
            "remote: Total 264 (delta 0), reused 0 (delta 0), pack-reused 261\u001b[K\n",
            "Receiving objects: 100% (264/264), 15.26 MiB | 6.93 MiB/s, done.\n",
            "Resolving deltas: 100% (140/140), done.\n",
            "/content/stylegan2\n",
            "CPU says hello.\n",
            "GPU says hello.\n",
            "Tensorflow version: 1.15.2\n",
            "GPU 0: Tesla P100-PCIE-16GB (UUID: GPU-076f5574-a36e-a12d-c1ca-c80960633032)\n",
            "GPU Identified at: /device:GPU:0\n",
            "Requirement already satisfied: gdown in /usr/local/lib/python3.6/dist-packages (3.6.4)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from gdown) (1.12.0)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.6/dist-packages (from gdown) (4.38.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.6/dist-packages (from gdown) (2.21.0)\n",
            "Requirement already satisfied: chardet<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests->gdown) (3.0.4)\n",
            "Requirement already satisfied: urllib3<1.25,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests->gdown) (1.24.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests->gdown) (2020.4.5.1)\n",
            "Requirement already satisfied: idna<2.9,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests->gdown) (2.8)\n",
            "/content/stylegan2/pkl\n",
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=1gbxwfHNOaGjGsLNTmmSrNA85X2VWHHOq\n",
            "To: /content/stylegan2/pkl/vgg16_zhang_perceptual.pkl\n",
            "58.9MB [00:01, 33.3MB/s]\n",
            "/content/stylegan2\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cwVXBFaSuoIU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Download the model of choice\n",
        "import argparse\n",
        "import numpy as np\n",
        "import PIL.Image\n",
        "import dnnlib\n",
        "import dnnlib.tflib as tflib\n",
        "import re\n",
        "import sys\n",
        "from io import BytesIO\n",
        "import IPython.display\n",
        "import numpy as np\n",
        "from math import ceil\n",
        "from PIL import Image, ImageDraw\n",
        "import imageio\n",
        "\n",
        "import pretrained_networks\n",
        "\n",
        "# Choose between these pretrained models - I think 'f' is the best choice:\n",
        "\n",
        "# 1024×1024 faces\n",
        "# stylegan2-ffhq-config-a.pkl\n",
        "# stylegan2-ffhq-config-b.pkl\n",
        "# stylegan2-ffhq-config-c.pkl\n",
        "# stylegan2-ffhq-config-d.pkl\n",
        "# stylegan2-ffhq-config-e.pkl\n",
        "# stylegan2-ffhq-config-f.pkl\n",
        "\n",
        "# 512×384 cars\n",
        "# stylegan2-car-config-a.pkl\n",
        "# stylegan2-car-config-b.pkl\n",
        "# stylegan2-car-config-c.pkl\n",
        "# stylegan2-car-config-d.pkl\n",
        "# stylegan2-car-config-e.pkl\n",
        "# stylegan2-car-config-f.pkl\n",
        "\n",
        "# 256x256 horses\n",
        "# stylegan2-horse-config-a.pkl\n",
        "# stylegan2-horse-config-f.pkl\n",
        "\n",
        "# 256x256 churches\n",
        "# stylegan2-church-config-a.pkl\n",
        "# stylegan2-church-config-f.pkl\n",
        "\n",
        "# 256x256 cats\n",
        "# stylegan2-cat-config-f.pkl\n",
        "# stylegan2-cat-config-a.pkl\n",
        "# network_pkl = \"gdrive:networks/stylegan2-ffhq-config-f.pkl\"\n",
        "\n",
        "# If downloads fails, due to 'Google Drive download quota exceeded' you can try downloading manually from your own Google Drive account\n",
        "!gdown --id 1UlDmJVLLnBD9SnLSMXeiZRO6g-OMQCA_\n",
        "network_pkl = \"/content/stylegan2/stylegan2-ffhq-config-f.pkl\"\n",
        "\n",
        "print('Loading networks from \"%s\"...' % network_pkl)\n",
        "_G, _D, Gs = pretrained_networks.load_networks(network_pkl)\n",
        "noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Zxbhe4uLvF_a",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Useful utility functions...\n",
        "\n",
        "# Generates a list of images, based on a list of latent vectors (Z), and a list (or a single constant) of truncation_psi's.\n",
        "def generate_images_in_w_space(dlatents, truncation_psi):\n",
        "    Gs_kwargs = dnnlib.EasyDict()\n",
        "    Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)\n",
        "    Gs_kwargs.randomize_noise = False\n",
        "    Gs_kwargs.truncation_psi = truncation_psi\n",
        "    dlatent_avg = Gs.get_var('dlatent_avg') # [component]\n",
        "\n",
        "    imgs = []\n",
        "    for row, dlatent in log_progress(enumerate(dlatents), name = \"Generating images\"):\n",
        "        #row_dlatents = (dlatent[np.newaxis] - dlatent_avg) * np.reshape(truncation_psi, [-1, 1, 1]) + dlatent_avg\n",
        "        dl = (dlatent-dlatent_avg)*truncation_psi   + dlatent_avg\n",
        "        row_images = Gs.components.synthesis.run(dlatent,  **Gs_kwargs)\n",
        "        imgs.append(PIL.Image.fromarray(row_images[0], 'RGB'))\n",
        "    return imgs       \n",
        "\n",
        "def generate_images(zs, truncation_psi):\n",
        "    Gs_kwargs = dnnlib.EasyDict()\n",
        "    Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)\n",
        "    Gs_kwargs.randomize_noise = False\n",
        "    if not isinstance(truncation_psi, list):\n",
        "        truncation_psi = [truncation_psi] * len(zs)\n",
        "        \n",
        "    imgs = []\n",
        "    for z_idx, z in log_progress(enumerate(zs), size = len(zs), name = \"Generating images\"):\n",
        "        Gs_kwargs.truncation_psi = truncation_psi[z_idx]\n",
        "        noise_rnd = np.random.RandomState(1) # fix noise\n",
        "        tflib.set_vars({var: noise_rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width]\n",
        "        images = Gs.run(z, None, **Gs_kwargs) # [minibatch, height, width, channel]\n",
        "        imgs.append(PIL.Image.fromarray(images[0], 'RGB'))\n",
        "    return imgs\n",
        "\n",
        "def generate_zs_from_seeds(seeds):\n",
        "    zs = []\n",
        "    for seed_idx, seed in enumerate(seeds):\n",
        "        rnd = np.random.RandomState(seed)\n",
        "        z = rnd.randn(1, *Gs.input_shape[1:]) # [minibatch, component]\n",
        "        zs.append(z)\n",
        "    return zs\n",
        "\n",
        "# Generates a list of images, based on a list of seed for latent vectors (Z), and a list (or a single constant) of truncation_psi's.\n",
        "def generate_images_from_seeds(seeds, truncation_psi):\n",
        "    return generate_images(generate_zs_from_seeds(seeds), truncation_psi)\n",
        "\n",
        "def saveImgs(imgs, location):\n",
        "  for idx, img in log_progress(enumerate(imgs), size = len(imgs), name=\"Saving images\"):\n",
        "    file = location+ str(idx) + \".png\"\n",
        "    img.save(file)\n",
        "\n",
        "def imshow(a, format='png', jpeg_fallback=True):\n",
        "  a = np.asarray(a, dtype=np.uint8)\n",
        "  str_file = BytesIO()\n",
        "  PIL.Image.fromarray(a).save(str_file, format)\n",
        "  im_data = str_file.getvalue()\n",
        "  try:\n",
        "    disp = IPython.display.display(IPython.display.Image(im_data))\n",
        "  except IOError:\n",
        "    if jpeg_fallback and format != 'jpeg':\n",
        "      print ('Warning: image was too large to display in format \"{}\"; '\n",
        "             'trying jpeg instead.').format(format)\n",
        "      return imshow(a, format='jpeg')\n",
        "    else:\n",
        "      raise\n",
        "  return disp\n",
        "\n",
        "def showarray(a, fmt='png'):\n",
        "    a = np.uint8(a)\n",
        "    f = StringIO()\n",
        "    PIL.Image.fromarray(a).save(f, fmt)\n",
        "    IPython.display.display(IPython.display.Image(data=f.getvalue()))\n",
        "\n",
        "        \n",
        "def clamp(x, minimum, maximum):\n",
        "    return max(minimum, min(x, maximum))\n",
        "    \n",
        "def drawLatent(image,latents,x,y,x2,y2, color=(255,0,0,100)):\n",
        "  buffer = PIL.Image.new('RGBA', image.size, (0,0,0,0))\n",
        "   \n",
        "  draw = ImageDraw.Draw(buffer)\n",
        "  cy = (y+y2)/2\n",
        "  draw.rectangle([x,y,x2,y2],fill=(255,255,255,180), outline=(0,0,0,180))\n",
        "  for i in range(len(latents)):\n",
        "    mx = x + (x2-x)*(float(i)/len(latents))\n",
        "    h = (y2-y)*latents[i]*0.1\n",
        "    h = clamp(h,cy-y2,y2-cy)\n",
        "    draw.line((mx,cy,mx,cy+h),fill=color)\n",
        "  return PIL.Image.alpha_composite(image,buffer)\n",
        "             \n",
        "  \n",
        "def createImageGrid(images, scale=0.25, rows=1):\n",
        "   w,h = images[0].size\n",
        "   w = int(w*scale)\n",
        "   h = int(h*scale)\n",
        "   height = rows*h\n",
        "   cols = ceil(len(images) / rows)\n",
        "   width = cols*w\n",
        "   canvas = PIL.Image.new('RGBA', (width,height), 'white')\n",
        "   for i,img in enumerate(images):\n",
        "     img = img.resize((w,h), PIL.Image.ANTIALIAS)\n",
        "     canvas.paste(img, (w*(i % cols), h*(i // cols))) \n",
        "   return canvas\n",
        "\n",
        "def convertZtoW(latent, truncation_psi=0.7, truncation_cutoff=9):\n",
        "  dlatent = Gs.components.mapping.run(latent, None) # [seed, layer, component]\n",
        "  dlatent_avg = Gs.get_var('dlatent_avg') # [component]\n",
        "  for i in range(truncation_cutoff):\n",
        "    dlatent[0][i] = (dlatent[0][i]-dlatent_avg)*truncation_psi + dlatent_avg\n",
        "    \n",
        "  return dlatent\n",
        "\n",
        "def interpolate(zs, steps):\n",
        "   out = []\n",
        "   for i in range(len(zs)-1):\n",
        "    for index in range(steps):\n",
        "     fraction = index/float(steps) \n",
        "     out.append(zs[i+1]*fraction + zs[i]*(1-fraction))\n",
        "   return out\n",
        "\n",
        "# Taken from https://github.com/alexanderkuk/log-progress\n",
        "def log_progress(sequence, every=1, size=None, name='Items'):\n",
        "    from ipywidgets import IntProgress, HTML, VBox\n",
        "    from IPython.display import display\n",
        "\n",
        "    is_iterator = False\n",
        "    if size is None:\n",
        "        try:\n",
        "            size = len(sequence)\n",
        "        except TypeError:\n",
        "            is_iterator = True\n",
        "    if size is not None:\n",
        "        if every is None:\n",
        "            if size <= 200:\n",
        "                every = 1\n",
        "            else:\n",
        "                every = int(size / 200)     # every 0.5%\n",
        "    else:\n",
        "        assert every is not None, 'sequence is iterator, set every'\n",
        "\n",
        "    if is_iterator:\n",
        "        progress = IntProgress(min=0, max=1, value=1)\n",
        "        progress.bar_style = 'info'\n",
        "    else:\n",
        "        progress = IntProgress(min=0, max=size, value=0)\n",
        "    label = HTML()\n",
        "    box = VBox(children=[label, progress])\n",
        "    display(box)\n",
        "\n",
        "    index = 0\n",
        "    try:\n",
        "        for index, record in enumerate(sequence, 1):\n",
        "            if index == 1 or index % every == 0:\n",
        "                if is_iterator:\n",
        "                    label.value = '{name}: {index} / ?'.format(\n",
        "                        name=name,\n",
        "                        index=index\n",
        "                    )\n",
        "                else:\n",
        "                    progress.value = index\n",
        "                    label.value = u'{name}: {index} / {size}'.format(\n",
        "                        name=name,\n",
        "                        index=index,\n",
        "                        size=size\n",
        "                    )\n",
        "            yield record\n",
        "    except:\n",
        "        progress.bar_style = 'danger'\n",
        "        raise\n",
        "    else:\n",
        "        progress.bar_style = 'success'\n",
        "        progress.value = index\n",
        "        label.value = \"{name}: {index}\".format(\n",
        "            name=name,\n",
        "            index=str(index or '?')\n",
        "        )\n",
        "\n",
        "\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rYdsgv4i6YPl",
        "colab_type": "text"
      },
      "source": [
        "# Projecting images onto the generatable manifold\n",
        "\n",
        "StyleGAN2 comes with a projector that finds the closest generatable image based on any input image. This allows you to get a feeling for the diversity of the portrait manifold.\n",
        "\n",
        "Let’s set up the folder we’ll need for this project."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "urzy8lw76j_r",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!mkdir projection\n",
        "!mkdir projection/imgs\n",
        "!mkdir projection/out\n",
        "\n",
        "# Now upload a single image to 'stylegan2/projection/imgs' (use the Files side panel). Image should be color PNG, with a size of 1024x1024."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "utoafQ6oSlAr",
        "colab_type": "text"
      },
      "source": [
        "Now upload a single image to `stylegan2/projection/imgs` (use the Files side panel). Image should be color PNG, with a size of 1024x1024.\n",
        "\n",
        "Things you might want to edit in here:\n",
        "\n",
        "- `file_name`: I recommend saving this to be something readable in case you want to do interpolations with it. Change it every time you switch projection samples.\n",
        "- `proj.num_steps`: this is the number of iterations you want the model to run thru. Higher numbers will increase accuracy but also take longer to produce."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Bi2mA9XaQxWb",
        "colab_type": "code",
        "outputId": "759997c3-ce41-49f0-fbbe-0e93797edc12",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "# Convert uploaded images to TFRecords\n",
        "import dataset_tool\n",
        "dataset_tool.create_from_images(\"./projection/records/\", \"./projection/imgs/\", True)\n",
        "\n",
        "# Run the projector\n",
        "import run_projector\n",
        "import projector\n",
        "import training.dataset\n",
        "import training.misc\n",
        "import os\n",
        "\n",
        "!rm projection/out/*.*\n",
        "\n",
        "def project_image_run(proj, targets, png_prefix, num_snapshots):\n",
        "    snapshot_steps = set(proj.num_steps - np.linspace(0, proj.num_steps, num_snapshots, endpoint=False, dtype=int))\n",
        "    training.misc.save_image_grid(targets, png_prefix + 'target.png', drange=[-1,1])\n",
        "    proj.start(targets)\n",
        "    while proj.get_cur_step() < proj.num_steps:\n",
        "        print('\\r%d / %d ... ' % (proj.get_cur_step(), proj.num_steps), end='', flush=True)\n",
        "        proj.step()\n",
        "        if proj.get_cur_step() in snapshot_steps:\n",
        "            training.misc.save_image_grid(proj.get_images(), png_prefix + 'step%04d.png' % proj.get_cur_step(), drange=[-1,1])\n",
        "    print('\\r%-30s\\r' % '', end='', flush=True)\n",
        "    type(proj.get_noises())\n",
        "    #if you want to do interpolations, name the file below something memorable\n",
        "    file_name = 'my_face'\n",
        "    np.save(('./projection/'+file_name+'.npy'), proj.get_dlatents())\n",
        "\n",
        "def project_real_images(dataset_name, data_dir, num_images, num_snapshots):\n",
        "    proj = projector.Projector()\n",
        "    proj.set_network(Gs)\n",
        "    #num_steps = how many iterations; larger = ~more accurate but longer run times \n",
        "    proj.num_steps=4000 \n",
        "\n",
        "    print('Loading images from \"%s\"...' % dataset_name)\n",
        "    dataset_obj = training.dataset.load_dataset(data_dir=data_dir, tfrecord_dir=dataset_name, max_label_size=0, verbose=True, repeat=False, shuffle_mb=0)\n",
        "    assert dataset_obj.shape == Gs.output_shape[1:]\n",
        "\n",
        "    for image_idx in range(num_images):\n",
        "        print('Projecting image %d/%d ...' % (image_idx, num_images))\n",
        "        images, _labels = dataset_obj.get_minibatch_np(1)\n",
        "        images = training.misc.adjust_dynamic_range(images, [0, 255], [-1, 1])\n",
        "        project_image_run(proj, targets=images, png_prefix=dnnlib.make_run_dir_path('projection/out/image%04d-' % image_idx), num_snapshots=num_snapshots)\n",
        "\n",
        "project_real_images(\"records\",\"./projection\",1,100)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\r  0%|          | 0/1 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Loading images from \"./projection/imgs/\"\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 1/1 [00:00<00:00,  3.26it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Projector: Finding W midpoint and stddev using 100 samples...\n",
            "Projector: std = 9.49727\n",
            "Projector: Setting up noise inputs...\n",
            "Projector: G_synthesis/noise0 <tf.Variable 'G_synthesis_4/noise0:0' shape=(1, 1, 4, 4) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise1 <tf.Variable 'G_synthesis_4/noise1:0' shape=(1, 1, 8, 8) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise2 <tf.Variable 'G_synthesis_4/noise2:0' shape=(1, 1, 8, 8) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise3 <tf.Variable 'G_synthesis_4/noise3:0' shape=(1, 1, 16, 16) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise4 <tf.Variable 'G_synthesis_4/noise4:0' shape=(1, 1, 16, 16) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise5 <tf.Variable 'G_synthesis_4/noise5:0' shape=(1, 1, 32, 32) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise6 <tf.Variable 'G_synthesis_4/noise6:0' shape=(1, 1, 32, 32) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise7 <tf.Variable 'G_synthesis_4/noise7:0' shape=(1, 1, 64, 64) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise8 <tf.Variable 'G_synthesis_4/noise8:0' shape=(1, 1, 64, 64) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise9 <tf.Variable 'G_synthesis_4/noise9:0' shape=(1, 1, 128, 128) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise10 <tf.Variable 'G_synthesis_4/noise10:0' shape=(1, 1, 128, 128) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise11 <tf.Variable 'G_synthesis_4/noise11:0' shape=(1, 1, 256, 256) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise12 <tf.Variable 'G_synthesis_4/noise12:0' shape=(1, 1, 256, 256) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise13 <tf.Variable 'G_synthesis_4/noise13:0' shape=(1, 1, 512, 512) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise14 <tf.Variable 'G_synthesis_4/noise14:0' shape=(1, 1, 512, 512) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise15 <tf.Variable 'G_synthesis_4/noise15:0' shape=(1, 1, 1024, 1024) dtype=float32_ref>\n",
            "Projector: G_synthesis/noise16 <tf.Variable 'G_synthesis_4/noise16:0' shape=(1, 1, 1024, 1024) dtype=float32_ref>\n",
            "Projector: Building image output graph...\n",
            "Projector: Building loss graph...\n",
            "Projector: Building noise regularization graph...\n",
            "Projector: Setting up optimizer...\n",
            "Loading images from \"records\"...\n",
            "Streaming data using training.dataset.TFRecordDataset...\n",
            "Dataset shape = [3, 1024, 1024]\n",
            "Dynamic range = [0, 255]\n",
            "Label size    = 0\n",
            "Projecting image 0/1 ...\n",
            "Projector: Preparing target images...\n",
            "Projector: Initializing optimization state...\n",
            "0 / 4000 ... Projector: Running...\n",
            "9 / 4000 ... Projector: 10      0.804985    15051.2     \n",
            "19 / 4000 ... Projector: 20      0.783072    7306.47     \n",
            "29 / 4000 ... Projector: 30      0.74531     1221.62     \n",
            "39 / 4000 ... Projector: 40      0.74135     387.972     \n",
            "49 / 4000 ... Projector: 50      0.724023    506.663     \n",
            "59 / 4000 ... Projector: 60      0.719339    43.9194     \n",
            "69 / 4000 ... Projector: 70      0.69408     76.7358     \n",
            "79 / 4000 ... Projector: 80      0.696007    18.8457     \n",
            "89 / 4000 ... Projector: 90      0.663357    10.0275     \n",
            "99 / 4000 ... Projector: 100     0.672493    15.2415     \n",
            "109 / 4000 ... Projector: 110     0.665343    8.62163     \n",
            "119 / 4000 ... Projector: 120     0.654462    14.3933     \n",
            "129 / 4000 ... Projector: 130     0.651055    15.3232     \n",
            "139 / 4000 ... Projector: 140     0.642931    13.1459     \n",
            "149 / 4000 ... Projector: 150     0.64878     19.3242     \n",
            "159 / 4000 ... Projector: 160     0.647528    35.3043     \n",
            "169 / 4000 ... Projector: 170     0.632534    46.8002     \n",
            "179 / 4000 ... Projector: 180     0.642496    25.2749     \n",
            "189 / 4000 ... Projector: 190     0.629047    32.6448     \n",
            "199 / 4000 ... Projector: 200     0.639432    19.7704     \n",
            "209 / 4000 ... Projector: 210     0.618729    31.5354     \n",
            "219 / 4000 ... Projector: 220     0.633526    20.5684     \n",
            "229 / 4000 ... Projector: 230     0.623258    11.1906     \n",
            "239 / 4000 ... Projector: 240     0.617763    6.67891     \n",
            "249 / 4000 ... Projector: 250     0.621849    7.61257     \n",
            "252 / 4000 ... "
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NTQTHCtDVrjb",
        "colab_type": "text"
      },
      "source": [
        "### Create video from projection search"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-6aZDk54VpE2",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 81,
          "referenced_widgets": [
            "290168f9058c4abaa846d807ce2984cc",
            "2332ec1b25cd4ca8a0c671e7884e7225",
            "5be1c18f016744d1903247cd85d65dbc",
            "54246974c1ec45a48a6064b1517c046b",
            "b1327f73da28457a8606204a30bf7d77",
            "11bbe6b7111e42bbabd18bc86a2e49d8",
            "0e7bb66324934d58bf7c058f35fe629e",
            "af4b3f75693241eca70a86ed02e2362c"
          ]
        },
        "outputId": "04bd2e78-0bfb-4770-b630-807d547b7db2"
      },
      "source": [
        "# Create video \n",
        "import glob\n",
        "\n",
        "imgs = sorted(glob.glob(\"projection/out/*step*.png\"))\n",
        "\n",
        "target_imgs = sorted(glob.glob(\"projection/out/*target*.png\"))\n",
        "assert len(target_imgs) == 1, \"More than one target found?\"\n",
        "target_img = imageio.imread(target_imgs[0])\n",
        "\n",
        "movieName = \"projection/movie.mp4\"\n",
        "with imageio.get_writer(movieName, mode='I') as writer:\n",
        "    for filename in log_progress(imgs, name = \"Creating animation\"):\n",
        "        image = imageio.imread(filename)\n",
        "\n",
        "        # Concatenate images with original target image\n",
        "        w,h = image.shape[0:2]\n",
        "        canvas = PIL.Image.new('RGBA', (w*2,h), 'white')\n",
        "        canvas.paste(Image.fromarray(target_img), (0, 0))\n",
        "        canvas.paste(Image.fromarray(image), (w, 0))\n",
        "\n",
        "        writer.append_data(np.array(canvas))"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "290168f9058c4abaa846d807ce2984cc",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "VBox(children=(HTML(value=''), IntProgress(value=0)))"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61XdlqEBXN_A",
        "colab_type": "text"
      },
      "source": [
        "### Make an interpolation\n",
        "\n",
        "If you generate a couple different projections, you can then interpolate between the different z values.\n",
        "\n",
        "(This is a little janky if someone wants to write a cleaner function please submit a PR!)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_aZvophLZQOw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "latent1 = np.load('/content/stylegan2/projection/latent-oprah.npy')\n",
        "latent2 = np.load('/content/stylegan2/projection/latent-ja.npy')\n",
        "latent3 = np.load('/content/stylegan2/projection/latent-bruce.npy')\n",
        "latent4 = np.load('/content/stylegan2/projection/latent-britney.npy')\n",
        "#noise_vars = np.load('/content/stylegan2/projection/noise.npy',allow_pickle=True)\n",
        "\n",
        "#imgs = generate_images_in_w_space([latent2],0.7)\n",
        "imgs = generate_images_in_w_space(interpolate([latent1,latent2,latent3,latent4,latent1],144),0.7)\n",
        "\n",
        "%mkdir interpolations\n",
        "saveImgs(imgs,'./interpolations/')\n",
        "!zip -r interpolations.zip ./interpolations"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lPjl5_1vu-Ag",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XGVarLre63dL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Now you can download the video (find it in the Files side panel under 'stylegan2/projection')\n",
        "\n",
        "# To cleanup\n",
        "!rm projection/out/*.*\n",
        "!rm projection/records/*.*\n",
        "#!rm projection/imgs/*.*"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9mGDRao-wrDx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!rm interpolations/*.png"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2RPLWN_M_kLv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}