StyleGAN3_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MH1i0fLwbBvk"
      },
      "source": [
        "論文  \n",
        "https://arxiv.org/abs/2201.13433  \n",
        "  \n",
        "GitHub  \n",
        "https://github.com/yuval-alaluf/stylegan3-editing  \n",
        "  \n",
        "<a href=\"https://colab.research.google.com/github/kaz12tech/ai_demos/blob/master/StyleGAN3_demo.ipynb\" target=\"_blank\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hs5lgHy1bBvq"
      },
      "source": [
        "# ランタイムの設定\n",
        "「ランタイム」→「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」をGPUに変更"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IiaCriA8bBvr"
      },
      "source": [
        "# 実行方法\n",
        "「ランタイム」→「すべてのセルを実行」を選択"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kMxw_RANbBvr"
      },
      "source": [
        "# GPU確認"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4UgMBgucbBvs",
        "outputId": "510104fc-2df4-4fa6-ee57-6aad0afe01c9"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wOmsOtenbBvt"
      },
      "source": [
        "# 環境セットアップ"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cw-C-iMzbBvu"
      },
      "source": [
        "## ライブラリのインポート"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2t-0NeEYbBvv"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from pathlib import Path"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SxZf2yrrbBvv"
      },
      "source": [
        "## GitHubからコードを取得"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z3sEvKlpbBvw",
        "outputId": "dcaaa948-4e78-4c9d-ed6d-43ed187f5d65"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "\n",
        "CODE_DIR = 'stylegan3-editing'\n",
        "!git clone https://github.com/yuval-alaluf/stylegan3-editing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7EESUczUbBvx"
      },
      "source": [
        "## ライブラリのインストール"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "P8VuMjGXbBvx",
        "outputId": "e598656e-bdd5-4dfc-f381-954aa4aca988"
      },
      "outputs": [],
      "source": [
        "!wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip\n",
        "!sudo unzip ninja-linux.zip -d /usr/local/bin/\n",
        "!sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force\n",
        "\n",
        "!pip install pyrallis\n",
        "!pip install git+https://github.com/openai/CLIP.git"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-o1gy8UebXhp"
      },
      "source": [
        "## ライブラリのインポート"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DEJNmH-vbBvy",
        "outputId": "fe8fa932-2eb7-4f50-cbf8-bcabdd0c7b4f"
      },
      "outputs": [],
      "source": [
        "%cd /content/{CODE_DIR}\n",
        "\n",
        "import time\n",
        "import sys\n",
        "import pprint\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "import dataclasses\n",
        "import torch\n",
        "import torchvision.transforms as transforms\n",
        "\n",
        "sys.path.append(\".\")\n",
        "sys.path.append(\"..\")\n",
        "\n",
        "from editing.interfacegan.face_editor import FaceEditor\n",
        "from editing.styleclip_global_directions import edit as styleclip_edit\n",
        "from models.stylegan3.model import GeneratorType\n",
        "from notebooks.notebook_utils import Downloader, ENCODER_PATHS, INTERFACEGAN_PATHS, STYLECLIP_PATHS\n",
        "from notebooks.notebook_utils import run_alignment, crop_image, compute_transforms\n",
        "from utils.common import tensor2im\n",
        "from utils.inference_utils import run_on_batch, load_encoder, get_average_image\n",
        "\n",
        "%load_ext autoreload\n",
        "%autoreload 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kFnhM7Qibndh"
      },
      "source": [
        "# 学習済みモデルのセットアップ"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mrfLYhGObqbc"
      },
      "outputs": [],
      "source": [
        "## Downloaderの設定"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4LSppCHmbaKy"
      },
      "outputs": [],
      "source": [
        "download_with_pydrive = False #@param {type:\"boolean\"}\n",
        "downloader = Downloader(code_dir=CODE_DIR,\n",
        "                        use_pydrive=download_with_pydrive,\n",
        "                        subdir=\"pretrained_models\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TbU3rUHIbvLB"
      },
      "outputs": [],
      "source": [
        "#@markdown 学習済みモデルの選択\n",
        "experiment_type = 'restyle_pSp_ffhq' #@param ['restyle_e4e_ffhq', 'restyle_pSp_ffhq']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KlagWhlPcqgG"
      },
      "source": [
        "## 推論時パラメータ設定"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JkIex9zOcpem"
      },
      "outputs": [],
      "source": [
        "EXPERIMENT_DATA_ARGS = {\n",
        "    \"restyle_pSp_ffhq\": {\n",
        "        \"model_path\": \"./pretrained_models/restyle_pSp_ffhq.pt\",\n",
        "        \"image_path\": \"./notebooks/images/face_image.jpg\",\n",
        "        \"transform\": transforms.Compose([\n",
        "            transforms.Resize((256, 256)),\n",
        "            transforms.ToTensor(),\n",
        "            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])\n",
        "    },\n",
        "    \"restyle_e4e_ffhq\": {\n",
        "        \"model_path\": \"./pretrained_models/restyle_e4e_ffhq.pt\",\n",
        "        \"image_path\": \"./notebooks/images/face_image.jpg\",\n",
        "        \"transform\": transforms.Compose([\n",
        "            transforms.Resize((256, 256)),\n",
        "            transforms.ToTensor(),\n",
        "            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])\n",
        "    }\n",
        "}\n",
        "\n",
        "EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS[experiment_type]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jENX4t6sc-nZ"
      },
      "source": [
        "## 学習済みモデルのダウンロード\n",
        "/content/stylegan3-editing/pretrained_models/restyle_pSp_ffhq.pt  をダウンロード"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oNkUU7XcdBD0",
        "outputId": "fce3f15c-1326-4065-f0cc-318f632f1480"
      },
      "outputs": [],
      "source": [
        "if not os.path.exists(EXPERIMENT_ARGS['model_path']) or os.path.getsize(EXPERIMENT_ARGS['model_path']) < 1000000:\n",
        "    print(f'Downloading ReStyle encoder model: {experiment_type}...')\n",
        "    try:\n",
        "      downloader.download_file(file_id=ENCODER_PATHS[experiment_type]['id'],\n",
        "                              file_name=ENCODER_PATHS[experiment_type]['name'])\n",
        "    except Exception as e:\n",
        "      raise ValueError(f\"Unable to download model correctly! {e}\")\n",
        "    # if google drive receives too many requests, we'll reach the quota limit and be unable to download the model\n",
        "    if os.path.getsize(EXPERIMENT_ARGS['model_path']) < 1000000:\n",
        "        raise ValueError(\"Pretrained model was unable to be downloaded correctly!\")\n",
        "    else:\n",
        "        print('Done.')\n",
        "else:\n",
        "    print(f'Model for {experiment_type} already exists!')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wd3eS5kTdS7Q"
      },
      "source": [
        "## 学習済みモデルのロード"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-UfIasuIdWjG",
        "outputId": "05d3f55e-af5f-45d2-e733-bea833c682a0"
      },
      "outputs": [],
      "source": [
        "model_path = EXPERIMENT_ARGS['model_path']\n",
        "net, opts = load_encoder(checkpoint_path=model_path)\n",
        "pprint.pprint(dataclasses.asdict(opts))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XimE7l1VdxK5"
      },
      "source": [
        "# テスト画像のアップロード\n",
        "縦・横同じサイズの画像をアップロードしてください。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 368,
          "resources": {
            "http://localhost:8080/nbextensions/google.colab/files.js": {
              "data": "Ly8gQ29weXJpZ2h0IDIwMTcgR29vZ2xlIExMQwovLwovLyBMaWNlbnNlZCB1bmRlciB0aGUgQXBhY2hlIExpY2Vuc2UsIFZlcnNpb24gMi4wICh0aGUgIkxpY2Vuc2UiKTsKLy8geW91IG1heSBub3QgdXNlIHRoaXMgZmlsZSBleGNlcHQgaW4gY29tcGxpYW5jZSB3aXRoIHRoZSBMaWNlbnNlLgovLyBZb3UgbWF5IG9idGFpbiBhIGNvcHkgb2YgdGhlIExpY2Vuc2UgYXQKLy8KLy8gICAgICBodHRwOi8vd3d3LmFwYWNoZS5vcmcvbGljZW5zZXMvTElDRU5TRS0yLjAKLy8KLy8gVW5sZXNzIHJlcXVpcmVkIGJ5IGFwcGxpY2FibGUgbGF3IG9yIGFncmVlZCB0byBpbiB3cml0aW5nLCBzb2Z0d2FyZQovLyBkaXN0cmlidXRlZCB1bmRlciB0aGUgTGljZW5zZSBpcyBkaXN0cmlidXRlZCBvbiBhbiAiQVMgSVMiIEJBU0lTLAovLyBXSVRIT1VUIFdBUlJBTlRJRVMgT1IgQ09ORElUSU9OUyBPRiBBTlkgS0lORCwgZWl0aGVyIGV4cHJlc3Mgb3IgaW1wbGllZC4KLy8gU2VlIHRoZSBMaWNlbnNlIGZvciB0aGUgc3BlY2lmaWMgbGFuZ3VhZ2UgZ292ZXJuaW5nIHBlcm1pc3Npb25zIGFuZAovLyBsaW1pdGF0aW9ucyB1bmRlciB0aGUgTGljZW5zZS4KCi8qKgogKiBAZmlsZW92ZXJ2aWV3IEhlbHBlcnMgZm9yIGdvb2dsZS5jb2xhYiBQeXRob24gbW9kdWxlLgogKi8KKGZ1bmN0aW9uKHNjb3BlKSB7CmZ1bmN0aW9uIHNwYW4odGV4dCwgc3R5bGVBdHRyaWJ1dGVzID0ge30pIHsKICBjb25zdCBlbGVtZW50ID0gZG9jdW1lbnQuY3JlYXRlRWxlbWVudCgnc3BhbicpOwogIGVsZW1lbnQudGV4dENvbnRlbnQgPSB0ZXh0OwogIGZvciAoY29uc3Qga2V5IG9mIE9iamVjdC5rZXlzKHN0eWxlQXR0cmlidXRlcykpIHsKICAgIGVsZW1lbnQuc3R5bGVba2V5XSA9IHN0eWxlQXR0cmlidXRlc1trZXldOwogIH0KICByZXR1cm4gZWxlbWVudDsKfQoKLy8gTWF4IG51bWJlciBvZiBieXRlcyB3aGljaCB3aWxsIGJlIHVwbG9hZGVkIGF0IGEgdGltZS4KY29uc3QgTUFYX1BBWUxPQURfU0laRSA9IDEwMCAqIDEwMjQ7CgpmdW5jdGlvbiBfdXBsb2FkRmlsZXMoaW5wdXRJZCwgb3V0cHV0SWQpIHsKICBjb25zdCBzdGVwcyA9IHVwbG9hZEZpbGVzU3RlcChpbnB1dElkLCBvdXRwdXRJZCk7CiAgY29uc3Qgb3V0cHV0RWxlbWVudCA9IGRvY3VtZW50LmdldEVsZW1lbnRCeUlkKG91dHB1dElkKTsKICAvLyBDYWNoZSBzdGVwcyBvbiB0aGUgb3V0cHV0RWxlbWVudCB0byBtYWtlIGl0IGF2YWlsYWJsZSBmb3IgdGhlIG5leHQgY2FsbAogIC8vIHRvIHVwbG9hZEZpbGVzQ29udGludWUgZnJvbSBQeXRob24uCiAgb3V0cHV0RWxlbWVudC5zdGVwcyA9IHN0ZXBzOwoKICByZXR1cm4gX3VwbG9hZEZpbGVzQ29udGludWUob3V0cHV0SWQpOwp9CgovLyBUaGlzIGlzIHJvdWdobHkgYW4gYXN5bmMgZ2VuZXJhdG9yIChub3Qgc3VwcG9ydGVkIGluIHRoZSBicm93c2VyIHlldCksCi8vIHdoZXJlIHRoZXJlIGFyZSBtdWx0aXBsZSBhc3luY2hyb25vdXMgc3RlcHMgYW5kIHRoZSBQeXRob24gc2lkZSBpcyBnb2luZwovLyB0byBwb2xsIGZvciBjb21wbGV0aW9uIG9mIGVhY2ggc3RlcC4KLy8gVGhpcyB1c2VzIGEgUHJvbWlzZSB0byBibG9jayB0aGUgcHl0aG9uIHNpZGUgb24gY29tcGxldGlvbiBvZiBlYWNoIHN0ZXAsCi8vIHRoZW4gcGFzc2VzIHRoZSByZXN1bHQgb2YgdGhlIHByZXZpb3VzIHN0ZXAgYXMgdGhlIGlucHV0IHRvIHRoZSBuZXh0IHN0ZXAuCmZ1bmN0aW9uIF91cGxvYWRGaWxlc0NvbnRpbnVlKG91dHB1dElkKSB7CiAgY29uc3Qgb3V0cHV0RWxlbWVudCA9IGRvY3VtZW50LmdldEVsZW1lbnRCeUlkKG91dHB1dElkKTsKICBjb25zdCBzdGVwcyA9IG91dHB1dEVsZW1lbnQuc3RlcHM7CgogIGNvbnN0IG5leHQgPSBzdGVwcy5uZXh0KG91dHB1dEVsZW1lbnQubGFzdFByb21pc2VWYWx1ZSk7CiAgcmV0dXJuIFByb21pc2UucmVzb2x2ZShuZXh0LnZhbHVlLnByb21pc2UpLnRoZW4oKHZhbHVlKSA9PiB7CiAgICAvLyBDYWNoZSB0aGUgbGFzdCBwcm9taXNlIHZhbHVlIHRvIG1ha2UgaXQgYXZhaWxhYmxlIHRvIHRoZSBuZXh0CiAgICAvLyBzdGVwIG9mIHRoZSBnZW5lcmF0b3IuCiAgICBvdXRwdXRFbGVtZW50Lmxhc3RQcm9taXNlVmFsdWUgPSB2YWx1ZTsKICAgIHJldHVybiBuZXh0LnZhbHVlLnJlc3BvbnNlOwogIH0pOwp9CgovKioKICogR2VuZXJhdG9yIGZ1bmN0aW9uIHdoaWNoIGlzIGNhbGxlZCBiZXR3ZWVuIGVhY2ggYXN5bmMgc3RlcCBvZiB0aGUgdXBsb2FkCiAqIHByb2Nlc3MuCiAqIEBwYXJhbSB7c3RyaW5nfSBpbnB1dElkIEVsZW1lbnQgSUQgb2YgdGhlIGlucHV0IGZpbGUgcGlja2VyIGVsZW1lbnQuCiAqIEBwYXJhbSB7c3RyaW5nfSBvdXRwdXRJZCBFbGVtZW50IElEIG9mIHRoZSBvdXRwdXQgZGlzcGxheS4KICogQHJldHVybiB7IUl0ZXJhYmxlPCFPYmplY3Q+fSBJdGVyYWJsZSBvZiBuZXh0IHN0ZXBzLgogKi8KZnVuY3Rpb24qIHVwbG9hZEZpbGVzU3RlcChpbnB1dElkLCBvdXRwdXRJZCkgewogIGNvbnN0IGlucHV0RWxlbWVudCA9IGRvY3VtZW50LmdldEVsZW1lbnRCeUlkKGlucHV0SWQpOwogIGlucHV0RWxlbWVudC5kaXNhYmxlZCA9IGZhbHNlOwoKICBjb25zdCBvdXRwdXRFbGVtZW50ID0gZG9jdW1lbnQuZ2V0RWxlbWVudEJ5SWQob3V0cHV0SWQpOwogIG91dHB1dEVsZW1lbnQuaW5uZXJIVE1MID0gJyc7CgogIGNvbnN0IHBpY2tlZFByb21pc2UgPSBuZXcgUHJvbWlzZSgocmVzb2x2ZSkgPT4gewogICAgaW5wdXRFbGVtZW50LmFkZEV2ZW50TGlzdGVuZXIoJ2NoYW5nZScsIChlKSA9PiB7CiAgICAgIHJlc29sdmUoZS50YXJnZXQuZmlsZXMpOwogICAgfSk7CiAgfSk7CgogIGNvbnN0IGNhbmNlbCA9IGRvY3VtZW50LmNyZWF0ZUVsZW1lbnQoJ2J1dHRvbicpOwogIGlucHV0RWxlbWVudC5wYXJlbnRFbGVtZW50LmFwcGVuZENoaWxkKGNhbmNlbCk7CiAgY2FuY2VsLnRleHRDb250ZW50ID0gJ0NhbmNlbCB1cGxvYWQnOwogIGNvbnN0IGNhbmNlbFByb21pc2UgPSBuZXcgUHJvbWlzZSgocmVzb2x2ZSkgPT4gewogICAgY2FuY2VsLm9uY2xpY2sgPSAoKSA9PiB7CiAgICAgIHJlc29sdmUobnVsbCk7CiAgICB9OwogIH0pOwoKICAvLyBXYWl0IGZvciB0aGUgdXNlciB0byBwaWNrIHRoZSBmaWxlcy4KICBjb25zdCBmaWxlcyA9IHlpZWxkIHsKICAgIHByb21pc2U6IFByb21pc2UucmFjZShbcGlja2VkUHJvbWlzZSwgY2FuY2VsUHJvbWlzZV0pLAogICAgcmVzcG9uc2U6IHsKICAgICAgYWN0aW9uOiAnc3RhcnRpbmcnLAogICAgfQogIH07CgogIGNhbmNlbC5yZW1vdmUoKTsKCiAgLy8gRGlzYWJsZSB0aGUgaW5wdXQgZWxlbWVudCBzaW5jZSBmdXJ0aGVyIHBpY2tzIGFyZSBub3QgYWxsb3dlZC4KICBpbnB1dEVsZW1lbnQuZGlzYWJsZWQgPSB0cnVlOwoKICBpZiAoIWZpbGVzKSB7CiAgICByZXR1cm4gewogICAgICByZXNwb25zZTogewogICAgICAgIGFjdGlvbjogJ2NvbXBsZXRlJywKICAgICAgfQogICAgfTsKICB9CgogIGZvciAoY29uc3QgZmlsZSBvZiBmaWxlcykgewogICAgY29uc3QgbGkgPSBkb2N1bWVudC5jcmVhdGVFbGVtZW50KCdsaScpOwogICAgbGkuYXBwZW5kKHNwYW4oZmlsZS5uYW1lLCB7Zm9udFdlaWdodDogJ2JvbGQnfSkpOwogICAgbGkuYXBwZW5kKHNwYW4oCiAgICAgICAgYCgke2ZpbGUudHlwZSB8fCAnbi9hJ30pIC0gJHtmaWxlLnNpemV9IGJ5dGVzLCBgICsKICAgICAgICBgbGFzdCBtb2RpZmllZDogJHsKICAgICAgICAgICAgZmlsZS5sYXN0TW9kaWZpZWREYXRlID8gZmlsZS5sYXN0TW9kaWZpZWREYXRlLnRvTG9jYWxlRGF0ZVN0cmluZygpIDoKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgJ24vYSd9IC0gYCkpOwogICAgY29uc3QgcGVyY2VudCA9IHNwYW4oJzAlIGRvbmUnKTsKICAgIGxpLmFwcGVuZENoaWxkKHBlcmNlbnQpOwoKICAgIG91dHB1dEVsZW1lbnQuYXBwZW5kQ2hpbGQobGkpOwoKICAgIGNvbnN0IGZpbGVEYXRhUHJvbWlzZSA9IG5ldyBQcm9taXNlKChyZXNvbHZlKSA9PiB7CiAgICAgIGNvbnN0IHJlYWRlciA9IG5ldyBGaWxlUmVhZGVyKCk7CiAgICAgIHJlYWRlci5vbmxvYWQgPSAoZSkgPT4gewogICAgICAgIHJlc29sdmUoZS50YXJnZXQucmVzdWx0KTsKICAgICAgfTsKICAgICAgcmVhZGVyLnJlYWRBc0FycmF5QnVmZmVyKGZpbGUpOwogICAgfSk7CiAgICAvLyBXYWl0IGZvciB0aGUgZGF0YSB0byBiZSByZWFkeS4KICAgIGxldCBmaWxlRGF0YSA9IHlpZWxkIHsKICAgICAgcHJvbWlzZTogZmlsZURhdGFQcm9taXNlLAogICAgICByZXNwb25zZTogewogICAgICAgIGFjdGlvbjogJ2NvbnRpbnVlJywKICAgICAgfQogICAgfTsKCiAgICAvLyBVc2UgYSBjaHVua2VkIHNlbmRpbmcgdG8gYXZvaWQgbWVzc2FnZSBzaXplIGxpbWl0cy4gU2VlIGIvNjIxMTU2NjAuCiAgICBsZXQgcG9zaXRpb24gPSAwOwogICAgZG8gewogICAgICBjb25zdCBsZW5ndGggPSBNYXRoLm1pbihmaWxlRGF0YS5ieXRlTGVuZ3RoIC0gcG9zaXRpb24sIE1BWF9QQVlMT0FEX1NJWkUpOwogICAgICBjb25zdCBjaHVuayA9IG5ldyBVaW50OEFycmF5KGZpbGVEYXRhLCBwb3NpdGlvbiwgbGVuZ3RoKTsKICAgICAgcG9zaXRpb24gKz0gbGVuZ3RoOwoKICAgICAgY29uc3QgYmFzZTY0ID0gYnRvYShTdHJpbmcuZnJvbUNoYXJDb2RlLmFwcGx5KG51bGwsIGNodW5rKSk7CiAgICAgIHlpZWxkIHsKICAgICAgICByZXNwb25zZTogewogICAgICAgICAgYWN0aW9uOiAnYXBwZW5kJywKICAgICAgICAgIGZpbGU6IGZpbGUubmFtZSwKICAgICAgICAgIGRhdGE6IGJhc2U2NCwKICAgICAgICB9LAogICAgICB9OwoKICAgICAgbGV0IHBlcmNlbnREb25lID0gZmlsZURhdGEuYnl0ZUxlbmd0aCA9PT0gMCA/CiAgICAgICAgICAxMDAgOgogICAgICAgICAgTWF0aC5yb3VuZCgocG9zaXRpb24gLyBmaWxlRGF0YS5ieXRlTGVuZ3RoKSAqIDEwMCk7CiAgICAgIHBlcmNlbnQudGV4dENvbnRlbnQgPSBgJHtwZXJjZW50RG9uZX0lIGRvbmVgOwoKICAgIH0gd2hpbGUgKHBvc2l0aW9uIDwgZmlsZURhdGEuYnl0ZUxlbmd0aCk7CiAgfQoKICAvLyBBbGwgZG9uZS4KICB5aWVsZCB7CiAgICByZXNwb25zZTogewogICAgICBhY3Rpb246ICdjb21wbGV0ZScsCiAgICB9CiAgfTsKfQoKc2NvcGUuZ29vZ2xlID0gc2NvcGUuZ29vZ2xlIHx8IHt9OwpzY29wZS5nb29nbGUuY29sYWIgPSBzY29wZS5nb29nbGUuY29sYWIgfHwge307CnNjb3BlLmdvb2dsZS5jb2xhYi5fZmlsZXMgPSB7CiAgX3VwbG9hZEZpbGVzLAogIF91cGxvYWRGaWxlc0NvbnRpbnVlLAp9Owp9KShzZWxmKTsK",
              "headers": [
                [
                  "content-type",
                  "application/javascript"
                ]
              ],
              "ok": true,
              "status": 200,
              "status_text": ""
            }
          }
        },
        "id": "Xy6JOZ6Wd1BR",
        "outputId": "537af3fc-6f80-4ab3-c7c9-ae788e2890ce"
      },
      "outputs": [],
      "source": [
        "%cd /content/{CODE_DIR}\n",
        "\n",
        "!mkdir upload_images\n",
        "%cd upload_images\n",
        "from google.colab import files\n",
        "uploaded = files.upload()\n",
        "uploaded = list(uploaded.keys())\n",
        "dir_path = \"/content/\" + CODE_DIR + \"/upload_images\"\n",
        "image_path = Path( os.path.join(dir_path, uploaded[0]) )\n",
        "\n",
        "#image_path = Path(EXPERIMENT_DATA_ARGS[experiment_type][\"image_path\"]) # ウィル・スミス\n",
        "\n",
        "original_image = Image.open(image_path).convert(\"RGB\")\n",
        "original_image = original_image.resize((256, 256))\n",
        "original_image"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_UUwBBldhmGj"
      },
      "source": [
        "# 顔画像の下処理"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1jcUsyE9hFER"
      },
      "source": [
        "## 顔部分の整列(alignment)、クロップ(crop)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 533
        },
        "id": "lV6BA7fvgcMx",
        "outputId": "6a9f4862-f4c4-4f05-85e7-abe2faa66d07"
      },
      "outputs": [],
      "source": [
        "%cd /content/{CODE_DIR}\n",
        "\n",
        "import dlib\n",
        "from utils.alignment_utils import align_face\n",
        "\n",
        "\n",
        "if not os.path.exists(\"./shape_predictor_68_face_landmarks.dat\"):\n",
        "  !wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2\n",
        "  !bzip2 -d shape_predictor_68_face_landmarks.dat.bz2\n",
        "\n",
        "predictor = dlib.shape_predictor(\"./shape_predictor_68_face_landmarks.dat\")\n",
        "detector = dlib.get_frontal_face_detector()\n",
        "print(\"Aligning image...\")\n",
        "input_image = align_face(filepath=str(image_path), detector=detector, predictor=predictor)\n",
        "\n",
        "#input_image = run_alignment(image_path)\n",
        "cropped_image = crop_image(image_path)\n",
        "joined = np.concatenate([input_image.resize((256, 256)), cropped_image.resize((256, 256))], axis=1)\n",
        "Image.fromarray(joined)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YuGWEsZVhdv3"
      },
      "source": [
        "## LandmarkベースのTransform"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "R4Wjh2F8hclh",
        "outputId": "07ba2a92-4d62-4091-fcfe-2847ddf0ae89"
      },
      "outputs": [],
      "source": [
        "images_dir = Path(\"./images\")\n",
        "images_dir.mkdir(exist_ok=True, parents=True)\n",
        "cropped_path = images_dir / f\"cropped_{image_path.name}\"\n",
        "aligned_path = images_dir / f\"aligned_{image_path.name}\"\n",
        "cropped_image.save(cropped_path)\n",
        "input_image.save(aligned_path)\n",
        "landmarks_transform = compute_transforms(aligned_path=aligned_path, cropped_path=cropped_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TQcgxHHBqQdk"
      },
      "source": [
        "# Inversion"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZPZgDi8dqUeC"
      },
      "outputs": [],
      "source": [
        "n_iters_per_batch =  10#@param {type:\"integer\"}\n",
        "opts.n_iters_per_batch = n_iters_per_batch\n",
        "opts.resize_outputs = False  # generate outputs at full resolution"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hAPn9NMRqZg1",
        "outputId": "27f56e52-bb9a-47be-8729-e28d8be2f367"
      },
      "outputs": [],
      "source": [
        "img_transforms = EXPERIMENT_ARGS['transform']\n",
        "transformed_image = img_transforms(input_image)\n",
        "\n",
        "avg_image = get_average_image(net)\n",
        "\n",
        "with torch.no_grad():\n",
        "    tic = time.time()\n",
        "    result_batch, result_latents = run_on_batch(inputs=transformed_image.unsqueeze(0).cuda().float(),\n",
        "                                                net=net,\n",
        "                                                opts=opts,\n",
        "                                                avg_image=avg_image,\n",
        "                                                landmarks_transform=torch.from_numpy(landmarks_transform).cuda().float())\n",
        "    toc = time.time()\n",
        "    print('Inference took {:.4f} seconds.'.format(toc - tic))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 473
        },
        "id": "h4uS8WEbqm9w",
        "outputId": "89d7c280-253b-44a4-c8d6-afe7e45af09b"
      },
      "outputs": [],
      "source": [
        "def get_coupled_results(result_batch, cropped_image):\n",
        "    result_tensors = result_batch[0]  # there's one image in our batch\n",
        "    resize_amount = (256, 256) if opts.resize_outputs else (opts.output_size, opts.output_size)\n",
        "    final_rec = tensor2im(result_tensors[-1]).resize(resize_amount)\n",
        "    input_im = cropped_image.resize(resize_amount)\n",
        "    res = np.concatenate([np.array(input_im), np.array(final_rec)], axis=1)\n",
        "    res = Image.fromarray(res)\n",
        "    return res\n",
        "\n",
        "res = get_coupled_results(result_batch, cropped_image)\n",
        "res.resize((1024, 512))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t22ijpE0qvQx"
      },
      "source": [
        "# 編集\n",
        "Inversionで得られた潜在コードを編集"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MHUkoAyhrJOQ"
      },
      "source": [
        "## boundaries, styleclipのダウンロード"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OIjs7kBrrAq-",
        "outputId": "674ff875-fc8b-4925-e921-7bab0dfe50fe"
      },
      "outputs": [],
      "source": [
        "download_with_pydrive = False #@param {type:\"boolean\"}\n",
        "\n",
        "# download files for interfacegan\n",
        "downloader = Downloader(code_dir=CODE_DIR,\n",
        "                        use_pydrive=download_with_pydrive,\n",
        "                        subdir=\"editing/interfacegan/boundaries/ffhq\")\n",
        "print(\"Downloading InterFaceGAN boundaries...\")\n",
        "for editing_file, params in INTERFACEGAN_PATHS.items():\n",
        "    print(f\"Downloading {editing_file} boundary...\")\n",
        "    downloader.download_file(file_id=params['id'],\n",
        "                             file_name=params['name'])\n",
        "\n",
        "# download files for styleclip\n",
        "downloader = Downloader(code_dir=CODE_DIR,\n",
        "                        use_pydrive=download_with_pydrive,\n",
        "                        subdir=\"editing/styleclip_global_directions/sg3-r-ffhq-1024\")\n",
        "print(\"Downloading StyleCLIP auxiliary files...\")\n",
        "for editing_file, params in STYLECLIP_PATHS.items():\n",
        "    print(f\"Downloading {editing_file}...\")\n",
        "    downloader.download_file(file_id=params['id'],\n",
        "                             file_name=params['name'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C_20_NRPrp5G"
      },
      "source": [
        "## 編集パラメータ設定"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LwJ_XQWerVnl"
      },
      "outputs": [],
      "source": [
        "editor = FaceEditor(stylegan_generator=net.decoder, generator_type=GeneratorType.ALIGNED)\n",
        "\n",
        "#@markdown 編集パラメータ設定\n",
        "edit_direction = 'Male' #@param ['age', 'smile', 'pose', 'Male']\n",
        "min_value = -5 #@param {type:\"slider\", min:-10, max:10, step:1}\n",
        "max_value = 5 #@param {type:\"slider\", min:-10, max:10, step:1}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z13ipvYfroX1"
      },
      "source": [
        "## 編集実行"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MPZpVLfMrgG_",
        "outputId": "5a836a9c-d0eb-4325-d4ae-bd8ccbc3f188"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "\n",
        "print(f\"Performing edit for {edit_direction}...\")\n",
        "input_latent = torch.from_numpy(result_latents[0][-1]).unsqueeze(0).cuda()\n",
        "edit_images, edit_latents = editor.edit(latents=input_latent,\n",
        "                                        direction=edit_direction,\n",
        "                                        factor_range=(min_value, max_value),\n",
        "                                        user_transforms=landmarks_transform,\n",
        "                                        apply_user_transformations=True)\n",
        "print(\"Done!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bN08rJNbrvLE"
      },
      "source": [
        "## 編集結果表示"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 63
        },
        "id": "ai9pXe6FrxHo",
        "outputId": "0d065c98-8441-4d2b-8df9-c9606e340ff7"
      },
      "outputs": [],
      "source": [
        "def prepare_edited_result(edit_images):\n",
        "  if type(edit_images[0]) == list:\n",
        "      edit_images = [image[0] for image in edit_images]\n",
        "  res = np.array(edit_images[0].resize((512, 512)))\n",
        "  for image in edit_images[1:]:\n",
        "      res = np.concatenate([res, image.resize((512, 512))], axis=1)\n",
        "  res = Image.fromarray(res).convert(\"RGB\")\n",
        "  return res\n",
        "\n",
        "res = prepare_edited_result(edit_images)\n",
        "res"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MuE15upilJxG"
      },
      "source": [
        "## 動画にまとめる"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CQCrLARglNV8",
        "outputId": "01198fad-5907-4ba3-d603-390e2bf03120"
      },
      "outputs": [],
      "source": [
        "%cd /content/{CODE_DIR}\n",
        "!rm -rf videos\n",
        "!mkdir -p videos/frames\n",
        "\n",
        "def create_video(edit_images):\n",
        "  if type(edit_images[0]) == list:\n",
        "    edit_images = [image[0] for image in edit_images]\n",
        "  res = np.array(edit_images[0].resize((512, 512)))\n",
        "  for i, image in enumerate( edit_images[1:] ):\n",
        "    work_dir = \"/content/\" + CODE_DIR\n",
        "    save_dir = os.path.join(work_dir, \"videos/frames\")\n",
        "    filename = edit_direction + \"_\" + str(i) + \".png\"\n",
        "    filename = os.path.join(save_dir, filename)\n",
        "    resize_img = image.resize((512, 512))\n",
        "    resize_img.save(filename)\n",
        "\n",
        "  framename = edit_direction + \"_%d.png\"\n",
        "  framename = os.path.join(save_dir, framename)\n",
        "  dst_video = os.path.join(work_dir, \"videos/result.mp4\")\n",
        "  !!ffmpeg -i {framename} -c:v libx264 -vf \"fps=25,format=yuv420p\" {dst_video}\n",
        "\n",
        "create_video(edit_images)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f1JLS9MWs0Sy"
      },
      "source": [
        "# StyleCLIP"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "q_fP9Aw6s1of",
        "outputId": "6198c731-f953-4aa2-c73b-1a94fca1ea89"
      },
      "outputs": [],
      "source": [
        "styleclip_args = styleclip_edit.EditConfig()\n",
        "global_direction_calculator = styleclip_edit.load_direction_calculator(stylegan_model=net.decoder, opts=styleclip_args)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QgyB1SfTs65C"
      },
      "source": [
        "## 編集パラメータ設定\n",
        "neutral_textはデフォルトが望ましい"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AMBh9-Qvs5-j"
      },
      "outputs": [],
      "source": [
        "neutral_text = \"a face\" #@param {type:\"raw\"}\n",
        "target_text = \"a smiling face\" #@param {type:\"raw\"}\n",
        "alpha = 4 #@param {type:\"slider\", min:-5, max:5, step:0.5}\n",
        "beta = 0.13 #@param {type:\"slider\", min:-1, max:1, step:0.1}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eqVIGgF-tHAj"
      },
      "source": [
        "## 編集実行"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VYLRiO02tGKq",
        "outputId": "4d1d47e9-a67e-4b22-86ba-50888cf7c5d0"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "\n",
        "opts = styleclip_edit.EditConfig()\n",
        "opts.alpha_min = alpha\n",
        "opts.alpha_max = alpha\n",
        "opts.num_alphas = 1\n",
        "opts.beta_min = beta\n",
        "opts.beta_max = beta\n",
        "opts.num_betas = 1\n",
        "opts.neutral_text = neutral_text\n",
        "opts.target_text = target_text\n",
        "\n",
        "input_latent = result_latents[0][-1]\n",
        "input_transforms = torch.from_numpy(landmarks_transform).cpu().numpy()\n",
        "print(f'Performing edit for: \"{opts.target_text}\"...')\n",
        "edit_res, edit_latent = styleclip_edit.edit_image(latent=input_latent,\n",
        "                                                  landmarks_transform=input_transforms,\n",
        "                                                  stylegan_model=net.decoder,\n",
        "                                                  global_direction_calculator=global_direction_calculator,\n",
        "                                                  opts=opts,\n",
        "                                                  image_name=None,\n",
        "                                                  save=False)\n",
        "print(\"Done!\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 473
        },
        "id": "Vwq61Xa3tN4F",
        "outputId": "06b9f3c4-aae9-44c6-d35a-f7a316d17e15"
      },
      "outputs": [],
      "source": [
        "input_im = tensor2im(transformed_image).resize((512, 512))\n",
        "edited_im = tensor2im(edit_res[0]).resize((512, 512))\n",
        "edit_coupled = np.concatenate([np.array(input_im), np.array(edited_im)], axis=1)\n",
        "edit_coupled = Image.fromarray(edit_coupled)\n",
        "edit_coupled.resize((1024, 512))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "StyleGAN3_demo.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 0
}