6DRepNet_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "論文  \n",
        "https://arxiv.org/abs/2202.12555<br>\n",
        "<br>  \n",
        "GitHub  \n",
        "https://github.com/thohemp/6drepnet<br>\n",
        "<br>\n",
        "<a href=\"https://colab.research.google.com/github/kaz12tech/ai_demos/blob/master/6DRepNet_demo.ipynb\" target=\"_blank\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 環境セットアップ"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## GPU確認"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2PcII5fLdHnJ",
        "outputId": "e94eeace-9878-47d7-a82c-20437e27baa7"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## GitHubからコード取得"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R-Y5auuvnnnT"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "\n",
        "!git clone https://github.com/thohemp/6DRepNet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## ライブラリのインストール"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LScsvO3UppqT"
      },
      "outputs": [],
      "source": [
        "%cd /content/6DRepNet\n",
        "\n",
        "!pip install --upgrade gdown\n",
        "!pip install git+https://github.com/elliottzheng/face-detection.git@master"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## ライブラリのインポート"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o2lxptJeo6kr"
      },
      "outputs": [],
      "source": [
        "from model import SixDRepNet\n",
        "import math\n",
        "import re\n",
        "from matplotlib import pyplot as plt\n",
        "import sys\n",
        "import os\n",
        "import argparse\n",
        "\n",
        "import numpy as np\n",
        "import cv2\n",
        "from google.colab.patches import cv2_imshow\n",
        "import matplotlib.pyplot as plt\n",
        "from numpy.lib.function_base import _quantile_unchecked\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision import transforms\n",
        "import torch.backends.cudnn as cudnn\n",
        "import torchvision\n",
        "import torch.nn.functional as F\n",
        "import utils\n",
        "import matplotlib\n",
        "from PIL import Image\n",
        "import time\n",
        "from face_detection import RetinaFace\n",
        "\n",
        "import glob\n",
        "from google.colab import files\n",
        "from tqdm import tqdm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 学習済みモデルのダウンロード"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QYmO4couoYhz"
      },
      "outputs": [],
      "source": [
        "%cd /content/6DRepNet\n",
        "!mkdir pretrained\n",
        "\n",
        "#https://drive.google.com/file/d/1vPNtVu_jg2oK-RiIWakxYyfLPA9rU4R4/view?usp=sharing\n",
        "pretrained_ckpt = 'pretrained/6DRepNet_300W_LP_AFLW2000.pth'\n",
        "if not os.path.exists(pretrained_ckpt):\n",
        "  !gdown --id 1vPNtVu_jg2oK-RiIWakxYyfLPA9rU4R4 \\\n",
        "          -O {pretrained_ckpt}\n",
        "\n",
        "snapshot_path = os.path.join(\"/content/6DRepNet/\", \"pretrained/6DRepNet_300W_LP_AFLW2000.pth\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# テスト動画のセットアップ"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iAm3NJEyHypU"
      },
      "source": [
        "## 動画のアップロード\n",
        "使用動画<br>\n",
        "https://www.pexels.com/ja-jp/video/3201691/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ljLvp7i_HLQc"
      },
      "outputs": [],
      "source": [
        "%cd /content/6DRepNet\n",
        "!rm -rf upload\n",
        "!mkdir -p upload/frames\n",
        "%cd upload\n",
        "\n",
        "uploaded = files.upload()\n",
        "uploaded = list(uploaded.keys())\n",
        "file_name = uploaded[0]\n",
        "\n",
        "upload_path = os.path.join(\"/content/6DRepNet/upload\", file_name)\n",
        "print(\"upload file here:\", upload_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "juTvlja6JAHM"
      },
      "source": [
        "## 動画をフレーム画像に分割"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yviWX14JJDAJ"
      },
      "outputs": [],
      "source": [
        "%cd /content/6DRepNet/upload\n",
        "\n",
        "!ffmpeg -i {upload_path} frames/%06d.png\n",
        "\n",
        "frames = glob.glob(\"/content/6DRepNet/upload/frames/*.png\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Head Pose Estimation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TriYzYbUqjcC"
      },
      "outputs": [],
      "source": [
        "%cd /content/6DRepNet\n",
        "\n",
        "!rm -rf output\n",
        "!mkdir -p output/frames\n",
        "\n",
        "cudnn.enabled = True\n",
        "gpu = 0\n",
        "\n",
        "print(\"Start model setup...\")\n",
        "# Modelのビルド\n",
        "model = SixDRepNet(\n",
        "    backbone_name='RepVGG-B1g2',\n",
        "    backbone_file='',\n",
        "    deploy=True,\n",
        "    pretrained=False)\n",
        "\n",
        "detector = RetinaFace(gpu_id=gpu)\n",
        "\n",
        "# Modelのロード\n",
        "saved_state_dict = torch.load(os.path.join(snapshot_path), map_location='cpu')\n",
        "\n",
        "if 'model_state_dict' in saved_state_dict:\n",
        "  model.load_state_dict(saved_state_dict['model_state_dict'])\n",
        "else:\n",
        "  model.load_state_dict(saved_state_dict)    \n",
        "  model.cuda(gpu)\n",
        "\n",
        "# Test the Model\n",
        "model.eval()  # Change model to 'eval' mode (BN uses moving mean/var).\n",
        "\n",
        "print(\"Complete model setup.\")\n",
        "\n",
        "print(\"loading \", len(frames), \" frames...\")\n",
        "\n",
        "process_start = time.time()\n",
        "with torch.no_grad():\n",
        "  for i in tqdm( range(len(frames)) ):\n",
        "    img_path = frames[i]\n",
        "    frame = np.array(Image.open(img_path))\n",
        "\n",
        "    faces = detector(frame)\n",
        "\n",
        "    for box, landmarks, score in faces:\n",
        "      # Print the location of each face in this image\n",
        "      if score < .95:\n",
        "          continue\n",
        "      x_min = int(box[0])\n",
        "      y_min = int(box[1])\n",
        "      x_max = int(box[2])\n",
        "      y_max = int(box[3])         \n",
        "      bbox_width = abs(x_max - x_min)\n",
        "      bbox_height = abs(y_max - y_min)\n",
        "\n",
        "      x_min = max(0,x_min-int(0.2*bbox_height))\n",
        "      y_min = max(0,y_min-int(0.2*bbox_width))\n",
        "      x_max = x_max+int(0.2*bbox_height)\n",
        "      y_max = y_max+int(0.2*bbox_width)\n",
        "\n",
        "      img = frame[y_min:y_max,x_min:x_max]\n",
        "\n",
        "      img = cv2.resize(img, (244, 244))/255.0\n",
        "      img = img.transpose(2, 0, 1)\n",
        "      img = torch.from_numpy(img).type(torch.FloatTensor)\n",
        "      img = torch.Tensor(img).cuda(gpu)\n",
        "      img=img.unsqueeze(0)\n",
        "\n",
        "      start = time.time()\n",
        "      R_pred = model(img)\n",
        "      end = time.time()\n",
        "      #print('Head pose estimation: %2f ms'% ((end - start)*1000.))\n",
        "\n",
        "      euler = utils.compute_euler_angles_from_rotation_matrices(R_pred)*180/np.pi\n",
        "      p_pred_deg = euler[:, 0].cpu()\n",
        "      y_pred_deg = euler[:, 1].cpu()\n",
        "      r_pred_deg = euler[:, 2].cpu()\n",
        "\n",
        "      utils.plot_pose_cube(frame,  y_pred_deg, p_pred_deg, r_pred_deg, x_min + int(.5*(x_max-x_min)), y_min + int(.5*(y_max-y_min)), size = bbox_width)\n",
        "\n",
        "    # 1フレーム完了毎に表示する場合はコメントアウト解除\n",
        "    #cv2_imshow(frame)\n",
        "    cv2.imwrite( os.path.join(\"/content/6DRepNet/output/frames\", os.path.basename(img_path)), frame)\n",
        "\n",
        "process_end = time.time()\n",
        "print('Complete All Head pose estimation: %2f s'% (process_end - process_start))\n",
        "print('Average %2f ms/ %06d frames'% (((process_end - process_start)*1000.)/len(frames), len(frames)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RU0jLi3rTUJz"
      },
      "source": [
        "## フレーム画像を動画に変換"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gaeX9vztTkU4"
      },
      "outputs": [],
      "source": [
        "!ffmpeg -i \"/content/6DRepNet/output/frames/%06d.png\" -c:v libx264 -vf \"format=yuv420p\" \"/content/6DRepNet/output/result.mp4\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Head Pose Estimationの結果を表示"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eSjzocGWT1w1"
      },
      "outputs": [],
      "source": [
        "from moviepy.editor import *\n",
        "from moviepy.video.fx.resize import resize\n",
        "clip = VideoFileClip(\"/content/6DRepNet/output/result.mp4\")\n",
        "clip = resize(clip, height=420)\n",
        "clip.ipython_display()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "6DRepNet_demo.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 0
}