AvatarCLIP_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o38wi96CDr-0"
      },
      "source": [
        "論文<br>\n",
        "https://arxiv.org/abs/2205.08535<br>\n",
        "<br>\n",
        "GitHub<br>\n",
        "https://github.com/hongfz16/AvatarCLIP<br>\n",
        "<br>\n",
        "<a href=\"https://colab.research.google.com/github/kaz12tech/ai_demos/blob/master/AvatarCLIP_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": "Dl0geaqIDr-5"
      },
      "source": [
        "# 環境セットアップ"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I0fVT2GdDr-6"
      },
      "source": [
        "## GPU確認"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2PcII5fLdHnJ"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zshu3r5PDr-7"
      },
      "source": [
        "## GitHubからコード取得"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R-Y5auuvnnnT"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "\n",
        "!git clone https://github.com/hongfz16/AvatarCLIP.git\n",
        "!git clone https://github.com/hongfz16/neural_renderer.git"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rii6XmjZDr-9"
      },
      "source": [
        "## ライブラリのインストール"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LScsvO3UppqT"
      },
      "outputs": [],
      "source": [
        "%cd /content/AvatarCLIP\n",
        "\n",
        "# Install Pytorch\n",
        "!pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html\n",
        "# Install lib\n",
        "!pip install -r requirements.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## neural_rendererファイルの修正"
      ],
      "metadata": {
        "id": "r4dm6JBJTqOf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/neural_renderer/neural_renderer\n",
        "\n",
        "with open('./perspective.py', 'r') as f:\n",
        "  lines = f.readlines()\n",
        "lines.insert(19, '    x[z<=0] = 0\\n')\n",
        "lines.insert(20, '    y[z<=0] = 0\\n')\n",
        "lines.insert(21, '    z[z<=0] = 0\\n')\n",
        "with open('./perspective.py', 'w') as f:\n",
        "  for l in lines:\n",
        "    f.write(l)"
      ],
      "metadata": {
        "id": "LO_xN37bTxNY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## neural_rendererのコンパイル"
      ],
      "metadata": {
        "id": "z8eapwdoUD8Q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/neural_renderer\n",
        "\n",
        "!python3 setup.py install"
      ],
      "metadata": {
        "id": "mimCNHQUUOpl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8olrgLpYEzWq"
      },
      "source": [
        "## SAMLのダウンロード\n",
        "\n",
        "{SAML}(https://smpl.is.tue.mpg.de/index.html)に登録<br>\n",
        "Download version 1.1.0 for Python 2.7 (female/male/neutral, 300 shape PCs)<br>をダウンロード後、Google Driveにアップロード\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xwyzDJ1wE1G6"
      },
      "outputs": [],
      "source": [
        "%cd /content/AvatarCLIP\n",
        "\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')\n",
        "\n",
        "# Google DriveのファイルをGoogle Colabにコピー\n",
        "!mkdir -p ./smpl_models/smpl\n",
        "!cp \"/content/gdrive/MyDrive/Colab Notebooks/SMPL_python_v.1.1.0.zip\" \"./smpl_models/smpl/SMPL_python_v.1.1.0.zip\"\n",
        "!unzip -d ./smpl_models/smpl/ ./smpl_models/smpl/SMPL_python_v.1.1.0.zip\n",
        "!cp ./smpl_models/smpl/SMPL_python_v.1.1.0/smpl/models/basicmodel_neutral_lbs_10_207_0_v1.1.0.pkl ./smpl_models/smpl/SMPL_NEUTRAL.pkl"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "raise Exception(\"Please restart runtime\")"
      ],
      "metadata": {
        "id": "iq2R4HEPbnyD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cW-nMTDnG5M3"
      },
      "source": [
        "## ライブラリのインポート\n",
        "以降のセル実行前に<br>\n",
        "ランタイム→ランタイムの再起動\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ],
      "metadata": {
        "id": "65VkQCjnc6eM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/AvatarCLIP/AvatarGen/AppearanceGen\n",
        "import os\n",
        "import time\n",
        "import logging\n",
        "import argparse\n",
        "import random\n",
        "import numpy as np\n",
        "import cv2 as cv\n",
        "import trimesh\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "from torchvision import transforms\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "from shutil import copyfile\n",
        "from icecream import ic\n",
        "from tqdm import tqdm\n",
        "from pyhocon import ConfigFactory\n",
        "from models.dataset import Dataset\n",
        "from models.dataset import SMPL_Dataset\n",
        "from models.fields import RenderingNetwork, SDFNetwork, SingleVarianceNetwork, NeRF\n",
        "from models.renderer import NeuSRenderer\n",
        "from models.utils import lookat, random_eye, random_at, render_one_batch, batch_rodrigues\n",
        "from models.utils import sphere_coord, random_eye_normal, rgb2hsv, differentiable_histogram\n",
        "from models.utils import my_lbs, readOBJ\n",
        "import clip\n",
        "from smplx import build_layer\n",
        "import imageio\n",
        "import glob\n",
        "\n",
        "import os\n",
        "from tqdm import tqdm\n",
        "import numpy as np\n",
        "from IPython import display\n",
        "from PIL import Image\n",
        "from base64 import b64encode\n",
        "\n",
        "to8b = lambda x : (255*np.clip(x,0,1)).astype(np.uint8)\n",
        "from main import Runner"
      ],
      "metadata": {
        "id": "B84BPPX3bRSr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# AppearanceDescription設定"
      ],
      "metadata": {
        "id": "IRx5quyBbffx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "AppearanceDescription = \"Fat Wizard\" #@param {type:\"string\"}\n",
        "\n",
        "torch.set_default_tensor_type('torch.cuda.FloatTensor')\n",
        "FORMAT = \"[%(filename)s:%(lineno)s - %(funcName)20s() ] %(message)s\"\n",
        "logging.basicConfig(level=logging.INFO, format=FORMAT)\n",
        "conf_path = 'confs/examples_small/example.conf'\n",
        "f = open(conf_path)\n",
        "conf_text = f.read()\n",
        "f.close()\n",
        "conf_text = conf_text.replace('{TOREPLACE}', AppearanceDescription)\n",
        "\n",
        "#@markdown google driveに保存する場合はTrue\n",
        "save_gdrive = True #@param {type:\"boolean\"}\n",
        "#@markdown save_gdrive指定時の出力先\n",
        "custom_exp_dir = \"/content/gdrive/MyDrive/Colab Notebooks\" #@param {type:\"string\"}\n",
        "if save_gdrive:\n",
        "  conf_text = conf_text.replace('base_exp_dir = ./exp/smpl/example', 'base_exp_dir = ' + custom_exp_dir + '/exp/smpl/example')\n",
        "\n",
        "\n",
        "# print(conf_text)\n",
        "conf = ConfigFactory.parse_string(conf_text)\n",
        "print(\"Prompt: {}\".format(conf.get_string('clip.prompt')))\n",
        "print(\"Face Prompt: {}\".format(conf.get_string('clip.face_prompt')))\n",
        "print(\"Back Prompt: {}\".format(conf.get_string('clip.back_prompt')))"
      ],
      "metadata": {
        "id": "zbdXtva9bX-H"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Avatar Generation"
      ],
      "metadata": {
        "id": "4QdgE1B_cCfS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "runner = Runner(conf_path, 'train_clip', 'smpl', False, True, conf)\n",
        "runner.init_clip()\n",
        "runner.init_smpl()\n",
        "runner.train_clip()"
      ],
      "metadata": {
        "id": "UPOVSARgb23E"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 最適化プロセスのビデオ生成"
      ],
      "metadata": {
        "id": "449pUC9BcWdf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "image_folder = 'exp/smpl/example/validations_extra_fine'\n",
        "video_path = 'video.mp4'\n",
        "if save_gdrive:\n",
        "  image_folder = os.path.join(custom_exp_dir, image_folder)\n",
        "  video_path = os.path.join(custom_exp_dir, video_path)\n",
        "image_fname = glob.glob(os.path.join(image_folder, \"*.png\"))\n",
        "\n",
        "init_frame = 1\n",
        "last_frame = len(image_fname)\n",
        "min_fps = 10\n",
        "max_fps = 60\n",
        "total_frames = last_frame - init_frame\n",
        "\n",
        "length = 15 #Desired time of the video in seconds\n",
        "\n",
        "frames = []\n",
        "for i in range(init_frame, last_frame): #\n",
        "    frames.append(Image.open(image_fname[i]))\n",
        "\n",
        "#fps = last_frame/10\n",
        "fps = np.clip(total_frames/length,min_fps,max_fps)\n",
        "\n",
        "from subprocess import Popen, PIPE\n",
        "p = Popen(['ffmpeg', '-y', '-f', 'image2pipe', '-vcodec', 'png', '-r', str(fps), '-i', '-', '-vcodec', 'libx264', '-r', str(fps), '-pix_fmt', 'yuv420p', '-crf', '17', '-preset', 'veryslow', video_path], stdin=PIPE)\n",
        "for im in tqdm(frames):\n",
        "    im.save(p.stdin, 'PNG')\n",
        "p.stdin.close()\n",
        "p.wait()\n",
        "mp4 = open(video_path,'rb').read()\n",
        "data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
        "\n",
        "display.HTML(\"\"\"\n",
        "<video width=400 controls>\n",
        "      <source src=\"%s\" type=\"video/mp4\">\n",
        "</video>\n",
        "\"\"\" % data_url)"
      ],
      "metadata": {
        "id": "Xg-BtvwrcTQw"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "r4dm6JBJTqOf",
        "z8eapwdoUD8Q"
      ],
      "name": "AvatarCLIP_demo.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 0
}