HairCLIP_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rKu2--YjjiNd"
      },
      "source": [
        "論文  \n",
        "https://arxiv.org/abs/2112.05142<br>\n",
        "GitHub<br>\n",
        "https://github.com/wty-ustc/HairCLIP<br>\n",
        "<br>\n",
        "<a href=\"https://colab.research.google.com/github/kaz12tech/ai_demos/blob/master/HairCLIP_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": "YWhJj2RejiNg"
      },
      "source": [
        "# 環境セットアップ"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2IuH74fQjiNg"
      },
      "source": [
        "## GPU確認"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2PcII5fLdHnJ",
        "outputId": "08da5e9b-8c39-480e-f04b-b504abe1c494"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XjrCXTWkjiNi"
      },
      "source": [
        "## GitHubからコード取得"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "R-Y5auuvnnnT",
        "outputId": "f543251c-6051-4992-b2a5-176ecda7064a"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "\n",
        "!git clone https://github.com/wty-ustc/HairCLIP.git\n",
        "\n",
        "!git clone https://github.com/omertov/encoder4editing.git\n",
        "!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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6pEp-7bNjiNj"
      },
      "source": [
        "## ライブラリのインストール"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LScsvO3UppqT",
        "outputId": "64f5a04c-9f0c-4cd0-83a0-dea7ad00ed62"
      },
      "outputs": [],
      "source": [
        "%cd /content/HairCLIP\n",
        "\n",
        "!pip install ftfy regex tqdm\n",
        "!pip install git+https://github.com/openai/CLIP.git\n",
        "!pip install tensorflow-io\n",
        "!pip install --upgrade --no-cache-dir gdown"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nDsSRIjsjiNj"
      },
      "source": [
        "## ライブラリのインポート"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CqJzYJTmEI_I",
        "outputId": "c14895a0-bbdc-4487-f226-006630980f65"
      },
      "outputs": [],
      "source": [
        "%cd /content/encoder4editing\n",
        "\n",
        "from utils.alignment import align_face\n",
        "from models.psp import pSp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "o2lxptJeo6kr",
        "outputId": "d195a5a9-e052-41d0-f6a4-6d3ec486a30b"
      },
      "outputs": [],
      "source": [
        "%cd /content/HairCLIP\n",
        "\n",
        "import os\n",
        "import gdown\n",
        "from argparse import ArgumentParser\n",
        "\n",
        "import sys\n",
        "sys.path.append(\".\")\n",
        "sys.path.append(\"..\")\n",
        "import tempfile\n",
        "from argparse import Namespace\n",
        "\n",
        "import dlib\n",
        "import torch\n",
        "import torchvision\n",
        "from torchvision import transforms\n",
        "from torch.utils.data import DataLoader\n",
        "from tqdm import tqdm\n",
        "\n",
        "import imageio\n",
        "from IPython.display import HTML\n",
        "from base64 import b64encode\n",
        "import glob\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "import random\n",
        "\n",
        "from criteria.parse_related_loss import average_lab_color_loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "w_9K81IODcZ3",
        "outputId": "e4644434-493d-4b56-c6e6-34732d52818e"
      },
      "outputs": [],
      "source": [
        "%cd /content/HairCLIP/mapper\n",
        "\n",
        "from mapper.datasets.latents_dataset_inference import LatentsDatasetInference\n",
        "from mapper.hairclip_mapper import HairCLIPMapper"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uoREcCRMjiNk"
      },
      "source": [
        "# 学習済みモデルのダウンロード\n",
        "Access denied with the following error:<br>\n",
        "が発生する場合、何回か実行"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rNPXkcuCpD3-",
        "outputId": "051fb6ed-7282-4f2d-9380-85acbd496892"
      },
      "outputs": [],
      "source": [
        "%cd /content/HairCLIP\n",
        "\n",
        "from pydrive.auth import GoogleAuth\n",
        "from pydrive.drive import GoogleDrive\n",
        "from google.colab import auth\n",
        "from oauth2client.client import GoogleCredentials\n",
        "\n",
        "download_with_pydrive = True\n",
        "\n",
        "class Downloader(object):\n",
        "    def __init__(self, use_pydrive):\n",
        "        self.use_pydrive = use_pydrive\n",
        "        current_directory = os.getcwd()\n",
        "        self.save_dir = \"/content/HairCLIP/pretrained_models\"\n",
        "        os.makedirs(self.save_dir, exist_ok=True)\n",
        "        if self.use_pydrive:\n",
        "            self.authenticate()\n",
        "\n",
        "    def authenticate(self):\n",
        "        auth.authenticate_user()\n",
        "        gauth = GoogleAuth()\n",
        "        gauth.credentials = GoogleCredentials.get_application_default()\n",
        "        self.drive = GoogleDrive(gauth)\n",
        "\n",
        "    def download_file(self, file_id, file_name):\n",
        "        file_dst = f'{self.save_dir}/{file_name}'\n",
        "        if os.path.exists(file_dst):\n",
        "            print(f'{file_name} already exists!')\n",
        "            return\n",
        "        if self.use_pydrive:\n",
        "            downloaded = self.drive.CreateFile({'id':file_id})\n",
        "            downloaded.FetchMetadata(fetch_all=True)\n",
        "            downloaded.GetContentFile(file_dst)\n",
        "        else:\n",
        "            !gdown --id $file_id -O $file_dst\n",
        "\n",
        "downloader = Downloader(download_with_pydrive)\n",
        "downloader.download_file(file_id=\"1cUv_reLE6k3604or78EranS7XzuVMWeO\", file_name=\"e4e_ffhq_encode.pt\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QYmO4couoYhz",
        "outputId": "94808ff6-31de-4207-9395-6f001d0000f3"
      },
      "outputs": [],
      "source": [
        "%cd /content/HairCLIP\n",
        "\n",
        "if not os.path.exists(\"./pretrained_models/hairclip.pt\"):\n",
        "  gdown.download('https://drive.google.com/uc?id=1hqZT6ZMldhX3M_x378Sm4Z2HMYr-UwQ4', \"./pretrained_models/hairclip.pt\", quiet=False)\n",
        "if not os.path.exists(\"./pretrained_models/stylegan2-ffhq-config-f.pt\"):\n",
        "  gdown.download('https://drive.google.com/uc?id=1pts5tkfAcWrg4TpLDu6ILF5wHID32Nzm', \"./pretrained_models/stylegan2-ffhq-config-f.pt\", quiet=False)\n",
        "if not os.path.exists(\"./pretrained_models/model_ir_se50.pth\"):\n",
        "  gdown.download('https://drive.google.com/uc?id=1FS2V756j-4kWduGxfir55cMni5mZvBTv', \"./pretrained_models/model_ir_se50.pth\", quiet=False)\n",
        "\n",
        "\n",
        "# if not os.path.exists(\"./pretrained_models/test_faces.pt\"):\n",
        "#   gdown.download('https://drive.google.com/uc?id=1j7RIfmrCoisxx3t-r-KC02Qc8barBecr', \"./pretrained_models/test_faces.pt\", quiet=False)\n",
        "\n",
        "if not os.path.exists(\"./pretrained_models/shape_predictor_68_face_landmarks.dat.bz2\"):\n",
        "  !wget -c http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2 \\\n",
        "        -O ./pretrained_models/shape_predictor_68_face_landmarks.dat.bz2\n",
        "  !bzip2 -dk ./pretrained_models/shape_predictor_68_face_landmarks.dat.bz2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IYP8fIDCjiNl"
      },
      "source": [
        "# テスト画像のセットアップ"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iAm3NJEyHypU"
      },
      "source": [
        "## テスト画像のダウンロード\n",
        "[使用画像](https://www.pakutaso.com/20190117010post-18492.html)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ljLvp7i_HLQc",
        "outputId": "13556e97-196b-4465-b64e-0bbe017637c6"
      },
      "outputs": [],
      "source": [
        "%cd /content/HairCLIP\n",
        "!mkdir demo\n",
        "\n",
        "!wget -c https://www.pakutaso.com/shared/img/thumb/model10211041_TP_V4.jpg \\\n",
        "      -O ./demo/model10211041_TP_V4.jpg"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gxplRMd50YUc"
      },
      "source": [
        "# Edit Hair"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Td9f98KGqaUs"
      },
      "source": [
        "## Set parametor"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6HUFDPiGGofy",
        "outputId": "1b929261-ba75-4519-ac9b-9b9c2a100638"
      },
      "outputs": [],
      "source": [
        "%cd /content/HairCLIP\n",
        "\n",
        "# @markdown 入力画像\n",
        "image_path = \"/content/HairCLIP/demo/model10211041_TP_V4.jpg\" #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "# @markdown editing_type=\"both\"でrandomに生成<br>\n",
        "# @markdown randomの場合下記の設定は反映されません。\n",
        "IsRandom = True #@param {type:\"boolean\"}\n",
        "random_num = 30 #@param {type:\"integer\"}\n",
        "\n",
        "# @markdown randomではない場合は以下設定<br>\n",
        "# @markdown 編集タイプ colorのみ、styleのみ、両方\n",
        "editing_type = \"both\" #@param[\"hairstyle\", \"color\", \"both\"]\n",
        "# @markdown 髪型選択\n",
        "hairstyle_description = \"crew cut hairstyle\" #@param[\"afro hairstyle\", \"bob cut hairstyle\", \"bowl cut hairstyle\", \"braid hairstyle\", \"caesar cut hairstyle\", \"chignon hairstyle\", \"cornrows hairstyle\", \"crew cut hairstyle\", \"crown braid hairstyle\", \"curtained hair hairstyle\", \"dido flip hairstyle\", \"dreadlocks hairstyle\", \"extensions hairstyle\", \"fade hairstyle\", \"fauxhawk hairstyle\", \"finger waves hairstyle\", \"french braid hairstyle\", \"frosted tips hairstyle\", \"full crown hairstyle\", \"harvard clip hairstyle\", \"high and tight hairstyle\", \"hime cut hairstyle\", \"hi-top fade hairstyle\",\"jewfro hairstyle\", \"jheri curl hairstyle\", \"liberty spikes hairstyle\", \"marcel waves hairstyle\", \"mohawk hairstyle\", \"pageboy hairstyle\", \"perm hairstyle\", \"pixie cut hairstyle\", \"psychobilly wedge hairstyle\", \"quiff hairstyle\", \"regular taper cut hairstyle\", \"ringlets hairstyle\", \"shingle bob hairstyle\", \"short hair hairstyle\", \"slicked-back hairstyle\", \"spiky hair hairstyle\",\"surfer hair hairstyle\", \"taper cut hairstyle\", \"the rachel hairstyle\", \"undercut hairstyle\", \"updo hairstyle\"]\n",
        "# @markdown 髪色選択\n",
        "color_description = \"yellow\" #@param[\"purple\", \"red\", \"orange\", \"yellow\", \"green\", \"blue\", \"gray\", \"brown\", \"black\", \"white\", \"blond\", \"pink\"]\n",
        "\n",
        "# 出力先ディレクトリ作成\n",
        "!mkdir outputs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YbaienSSuk9H"
      },
      "source": [
        "## Define functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h3uBaoakLK3Z"
      },
      "outputs": [],
      "source": [
        "def run_alignment(image_path):\n",
        "  predictor = dlib.shape_predictor(\"/content/HairCLIP/pretrained_models/shape_predictor_68_face_landmarks.dat\")\n",
        "  aligned_image = align_face(filepath=image_path, predictor=predictor)\n",
        "  print(\"Aligned image has shape: {}\".format(aligned_image.size))\n",
        "  return aligned_image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kHrd9G11Lja9"
      },
      "outputs": [],
      "source": [
        "def run_on_batch_e4e(inputs, net):\n",
        "  images, latents = net(\n",
        "      inputs.to(\"cuda\").float(), randomize_noise=False, return_latents=True\n",
        "      )\n",
        "  return images, latents"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2MP0KYACLw_7"
      },
      "outputs": [],
      "source": [
        "def run_on_batch(\n",
        "    inputs,\n",
        "    hairstyle_text_inputs,\n",
        "    color_text_inputs,\n",
        "    hairstyle_tensor_hairmasked,\n",
        "    color_tensor_hairmasked,\n",
        "    net,\n",
        "):\n",
        "    w = inputs\n",
        "    with torch.no_grad():\n",
        "        w_hat = w + 0.1 * net.mapper(\n",
        "            w,\n",
        "            hairstyle_text_inputs,\n",
        "            color_text_inputs,\n",
        "            hairstyle_tensor_hairmasked,\n",
        "            color_tensor_hairmasked,\n",
        "        )\n",
        "        x_hat, w_hat = net.decoder(\n",
        "            [w_hat],\n",
        "            input_is_latent=True,\n",
        "            return_latents=True,\n",
        "            randomize_noise=False,\n",
        "            truncation=1,\n",
        "        )\n",
        "        x, _ = net.decoder(\n",
        "            [w], input_is_latent=True, randomize_noise=False, truncation=1\n",
        "        )\n",
        "        result_batch = (x_hat, w_hat, x)\n",
        "    return result_batch\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NBKIidbVo-2N"
      },
      "outputs": [],
      "source": [
        "def predict(\n",
        "    edit_t, hair_d, color_d, \n",
        "    ck, im_path, trans, device):\n",
        "    editing_type_ = edit_t\n",
        "    hairstyle_description_ = hair_d\n",
        "    color_description_ = color_d\n",
        "\n",
        "    if editing_type_ == \"both\":\n",
        "      assert (\n",
        "          hairstyle_description_ is not None and color_d is not None\n",
        "          ), (\"Please provide description \" \"for both hairstyle and color.\")\n",
        "    elif editing_type_ == \"hairstyle\":\n",
        "      assert (\n",
        "          hairstyle_description_ is not None\n",
        "          ), \"Please provide description for hairstyle.\"\n",
        "    else:\n",
        "      assert (\n",
        "          color_description_ is not None\n",
        "          ), \"Please provide description for color.\"\n",
        "\n",
        "    opts = ck[\"opts\"]\n",
        "    opts = Namespace(**opts)\n",
        "    opts.editing_type = editing_type_\n",
        "    opts.input_type = \"text\"\n",
        "    opts.color_description = color_description_\n",
        "\n",
        "    if hair_d is not None:\n",
        "      with open(\"/content/HairCLIP/outputs/hairstyle_description.txt\", \"w\") as file:\n",
        "        file.write(hairstyle_description_)\n",
        "      opts.hairstyle_description = \"/content/HairCLIP/outputs/hairstyle_description.txt\"\n",
        "\n",
        "    opts.checkpoint_path  = \"/content/HairCLIP/pretrained_models/hairclip.pt\"\n",
        "    opts.parsenet_weights = \"/content/HairCLIP/pretrained_models/parsenet.pth\"\n",
        "    opts.stylegan_weights = \"/content/HairCLIP/pretrained_models/stylegan2-ffhq-config-f.pt\"\n",
        "    opts.ir_se50_weights  = \"/content/HairCLIP/pretrained_models/model_ir_se50.pth\"\n",
        "    net = HairCLIPMapper(opts)\n",
        "    net.eval()\n",
        "    net.cuda()\n",
        "\n",
        "    # 顔部分のalignment, transform\n",
        "    input_image = run_alignment(str(im_path))\n",
        "    resize_dims = (256, 256)\n",
        "    input_image.resize(resize_dims)\n",
        "    transformed_image = trans(input_image)\n",
        "\n",
        "    with torch.no_grad():\n",
        "      images, latents = run_on_batch_e4e(\n",
        "          transformed_image.unsqueeze(0), e4e_net\n",
        "          )\n",
        "      print(\"Latent code calculated!\")\n",
        "\n",
        "    dataset = LatentsDatasetInference(latents=latents.cpu(), opts=opts)\n",
        "    dataloader = DataLoader(dataset)\n",
        "\n",
        "    average_color_loss = (\n",
        "        average_lab_color_loss.AvgLabLoss(opts).to(device).eval()\n",
        "        )\n",
        "\n",
        "    out_filename = editing_type_ + \"_\" \\\n",
        "      + hairstyle_description_ + \"_\" \\\n",
        "      + color_description_ + \".png\"\n",
        "    out_path = os.path.join(\"/content/HairCLIP/outputs\", out_filename)\n",
        "    print(\"output path:\", out_path)\n",
        "\n",
        "    for input_batch in tqdm(dataloader):\n",
        "      with torch.no_grad():\n",
        "        w,hairstyle_text_inputs_list, color_text_inputs_list, selected_description_tuple_list, hairstyle_tensor_list, color_tensor_list = input_batch\n",
        "\n",
        "        hairstyle_text_inputs = hairstyle_text_inputs_list[0]\n",
        "        color_text_inputs = color_text_inputs_list[0]\n",
        "\n",
        "        selected_description = selected_description_tuple_list[0][0]\n",
        "        hairstyle_tensor = hairstyle_tensor_list[0]\n",
        "        color_tensor = color_tensor_list[0]\n",
        "\n",
        "        w = w.cuda().float()\n",
        "        hairstyle_text_inputs = hairstyle_text_inputs.cuda()\n",
        "        color_text_inputs = color_text_inputs.cuda()\n",
        "        hairstyle_tensor = hairstyle_tensor.cuda()\n",
        "        color_tensor = color_tensor.cuda()\n",
        "        if hairstyle_tensor.shape[1] != 1:\n",
        "          hairstyle_tensor_hairmasked = (\n",
        "              hairstyle_tensor * average_color_loss.gen_hair_mask(hairstyle_tensor))\n",
        "        else:\n",
        "          hairstyle_tensor_hairmasked = torch.Tensor([0]).unsqueeze(0).cuda()\n",
        "\n",
        "        if color_tensor.shape[1] != 1:\n",
        "          color_tensor_hairmasked = (\n",
        "              color_tensor * average_color_loss.gen_hair_mask(color_tensor))\n",
        "        else:\n",
        "          color_tensor_hairmasked = torch.Tensor([0]).unsqueeze(0).cuda()\n",
        "          result_batch = run_on_batch(\n",
        "              w, hairstyle_text_inputs, color_text_inputs,\n",
        "              hairstyle_tensor_hairmasked, color_tensor_hairmasked, net,)   \n",
        "\n",
        "        if (hairstyle_tensor.shape[1] != 1) and (color_tensor.shape[1] != 1):\n",
        "          img_tensor = torch.cat([hairstyle_tensor, color_tensor], dim=3)\n",
        "        elif hairstyle_tensor.shape[1] != 1:\n",
        "          img_tensor = hairstyle_tensor\n",
        "        elif color_tensor.shape[1] != 1:\n",
        "          img_tensor = color_tensor\n",
        "        else:\n",
        "          img_tensor = None   \n",
        "\n",
        "        if img_tensor is not None:\n",
        "          if img_tensor.shape[3] == 1024:\n",
        "            couple_output = torch.cat(\n",
        "                [result_batch[2][0].unsqueeze(0), result_batch[0][0].unsqueeze(0), img_tensor,])\n",
        "          elif img_tensor.shape[3] == 2048:\n",
        "            couple_output = torch.cat(\n",
        "                [result_batch[2][0].unsqueeze(0), result_batch[0][0].unsqueeze(0),\n",
        "                 img_tensor[:, :, :, 0:1024], img_tensor[:, :, :, 1024::], ])\n",
        "            couple_output = torch.cat(\n",
        "                [result_batch[2][0].unsqueeze(0), result_batch[0][0].unsqueeze(0),\n",
        "                 img_tensor[:, :, :, 0:1024], img_tensor[:, :, :, 1024::], ])\n",
        "        else:\n",
        "            couple_output = torch.cat(\n",
        "                [result_batch[2][0].unsqueeze(0),result_batch[0][0].unsqueeze(0),])\n",
        "\n",
        "        torchvision.utils.save_image(\n",
        "            couple_output, str(out_path), normalize=True, range=(-1, 1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JIv-IwXRhbMJ"
      },
      "source": [
        "## Predict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gzj5iYNEFOY_",
        "outputId": "81085434-4b51-481d-c2ba-8602dc4b9af4"
      },
      "outputs": [],
      "source": [
        "with open(\"/content/HairCLIP/mapper/hairstyle_list.txt\") as infile:\n",
        "  HAIRSTYLE_LIST = sorted([line.rstrip() for line in infile])\n",
        "COLORSTYLE_LIST = [\"purple\", \"red\", \"orange\", \"yellow\", \"green\", \"blue\", \"gray\", \"brown\", \"black\", \"white\", \"blond\", \"pink\"]\n",
        "\n",
        "device = \"cuda:0\"\n",
        "\n",
        "# load e4e ffhq model\n",
        "e4e_model_path = \"/content/HairCLIP/pretrained_models/e4e_ffhq_encode.pt\"\n",
        "e4e_ckpt = torch.load(e4e_model_path, map_location=\"cpu\")\n",
        "e4e_opts = e4e_ckpt[\"opts\"]\n",
        "e4e_opts[\"checkpoint_path\"] = e4e_model_path\n",
        "e4e_opts = Namespace(**e4e_opts)\n",
        "\n",
        "e4e_net = pSp(e4e_opts)\n",
        "e4e_net.eval()\n",
        "e4e_net.cuda()\n",
        "print(\"e4e model successfully loaded!\")\n",
        "\n",
        "# set transforms\n",
        "img_transforms = transforms.Compose(\n",
        "    [\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",
        "# load hairclip model\n",
        "checkpoint_path = \"/content/HairCLIP/pretrained_models/hairclip.pt\"\n",
        "ckpt = torch.load(checkpoint_path, map_location=\"cpu\")\n",
        "\n",
        "if IsRandom == True:\n",
        "  for i in range(random_num):\n",
        "    hairstyle_index = random.randrange(0, (len(HAIRSTYLE_LIST)-1), 1)\n",
        "    colorstyle_index = random.randrange(0, (len(COLORSTYLE_LIST)-1), 1)\n",
        "    predict(\n",
        "        \"both\", HAIRSTYLE_LIST[hairstyle_index], COLORSTYLE_LIST[colorstyle_index], \n",
        "        ckpt, image_path, img_transforms, device)\n",
        "\n",
        "else:\n",
        "  predict(\n",
        "      editing_type, hairstyle_description, color_description, \n",
        "      ckpt, image_path, img_transforms, device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qsk2Zd5NhhPW"
      },
      "source": [
        "## Show result"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 319
        },
        "id": "jObdg7PJelcK",
        "outputId": "dc4b4dc1-6389-438a-f92b-a491b04da307"
      },
      "outputs": [],
      "source": [
        "def generate_mp4(out_name, images, kwargs):\n",
        "  writer = imageio.get_writer(out_name + '.mp4', **kwargs)\n",
        "  for image in images:\n",
        "    writer.append_data(image)\n",
        "  writer.close()\n",
        "\n",
        "def show_mp4(filename, width):\n",
        "  mp4 = open(filename + '.mp4', 'rb').read()\n",
        "  data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
        "  display(HTML(\"\"\"\n",
        "  <video width=\"%d\" controls autoplay loop>\n",
        "    <source src=\"%s\" type=\"video/mp4\">\n",
        "  </video>\n",
        "  \"\"\" % (width, data_url)))\n",
        "\n",
        "res_list = glob.glob(\"/content/HairCLIP/outputs/*.png\")\n",
        "images = []\n",
        "for img_path in res_list:\n",
        "  images.append(np.array(Image.open(img_path)))\n",
        "\n",
        "kwargs = {'fps': 2}\n",
        "\n",
        "gif_path = os.path.join(\"/content/HairCLIP/outputs\", \"animation\")\n",
        "generate_mp4(gif_path, images, kwargs)\n",
        "show_mp4(gif_path, width=514)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "HairCLIP_demo.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 0
}