StridedTransformer_Pose3D.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G9fueCPLFKin"
      },
      "source": [
        "GitHub<br>\n",
        "https://github.com/Vegetebird/StridedTransformer-Pose3D<br>\n",
        "論文<br>\n",
        "https://arxiv.org/abs/2103.14304<br>\n",
        "<br>\n",
        "<a href=\"https://colab.research.google.com/github/kaz12tech/ai_demos/blob/master/StridedTransformer_Pose3D.ipynb\" target=\"_blank\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "894c0QeEFKit"
      },
      "source": [
        "# 環境セットアップ"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iWBmluQQx-MK"
      },
      "source": [
        "## GPU確認"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qyLiFmW6x_3w",
        "outputId": "ded71052-9eb5-44b7-c365-3ea21f6394b0"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TmL3LoMMGcNF"
      },
      "source": [
        "## Githubからソースコード取得"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tthUF7D6FKiu",
        "outputId": "24260972-c149-429b-917f-3ae42839d279"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "\n",
        "!git clone https://github.com/Vegetebird/StridedTransformer-Pose3D.git"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2-2Fh1sOHdnn"
      },
      "source": [
        "## ライブラリのインストール"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "AP8SsHYvHfkR",
        "outputId": "66abe431-1aae-4908-cfa8-efeb774593dd"
      },
      "outputs": [],
      "source": [
        "%cd /content/StridedTransformer-Pose3D\n",
        "\n",
        "!pip install --upgrade gdown\n",
        "!pip install yacs\n",
        "!pip install filterpy\n",
        "!pip install einops\n",
        "!pip install yt-dlp moviepy\n",
        "!pip install matplotlib==3.0.3\n",
        "!pip install imageio==2.4.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mgANPXUSy_MJ"
      },
      "source": [
        "## ライブラリのインポート"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hq7ZoY0oh260",
        "outputId": "7a984ed5-9f59-4166-d93d-9d635f04ecde"
      },
      "outputs": [],
      "source": [
        "%cd /content/StridedTransformer-Pose3D\n",
        "\n",
        "import os\n",
        "from yt_dlp import YoutubeDL\n",
        "\n",
        "from moviepy.video.fx.resize import resize\n",
        "from moviepy.editor import VideoFileClip, AudioFileClip, ImageSequenceClip, CompositeAudioClip\n",
        "from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_subclip"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jRBoab25xtgl"
      },
      "source": [
        "# 学習済みモデルのセットアップ"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5cSMMFMwx2P5",
        "outputId": "35f4f8de-b4c0-41d1-bc7a-a859007b3a9e"
      },
      "outputs": [],
      "source": [
        "%cd /content/StridedTransformer-Pose3D\n",
        "!mkdir -p ./checkpoint/pretrained\n",
        "\n",
        "if not os.path.exists('checkpoint/pretrained/refine_4365.pth'):\n",
        "  !gdown https://drive.google.com/uc?id=1aDLu0SB9JnPYZOOzQsJMV9zEIHg2Uro7 -O checkpoint/pretrained/refine_4365.pth\n",
        "if not os.path.exists('checkpoint/pretrained/no_refine_4365.pth'):\n",
        "  !gdown https://drive.google.com/uc?id=1l63AI9BsNovpfTAbfAkySo9X2MOWgYZH -O checkpoint/pretrained/no_refine_4365.pth\n",
        "\n",
        "if not os.path.exists('demo/lib/checkpoint/yolov3.weights'):\n",
        "  !gdown https://drive.google.com/uc?id=1gWZl1VrlLZKBf0Pfkj4hKiFxe8sHP-1C -O demo/lib/checkpoint/yolov3.weights\n",
        "if not os.path.exists('demo/lib/checkpoint/pose_hrnet_w48_384x288.pth'):\n",
        "  !gdown https://drive.google.com/uc?id=1CpyZiUIUlEjiql4rILwdBT4666S72Oq4 -O demo/lib/checkpoint/pose_hrnet_w48_384x288.pth"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_BB4TY8OyFkw"
      },
      "source": [
        "# テスト動画のセットアップ"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "O6-h-pFo_jHI",
        "outputId": "83708f0a-410f-4f71-dc34-dd8edf29e608"
      },
      "outputs": [],
      "source": [
        "%cd /content/StridedTransformer-Pose3D/demo/video\n",
        "\n",
        "video_url = 'https://www.youtube.com/watch?v=VSSPwzSflr0' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown 動画の切り抜き範囲(秒)を指定してください。\\\n",
        "#@markdown 30秒以上の場合OOM発生の可能性が高いため注意\n",
        "start_sec =  1#@param {type:\"integer\"}\n",
        "end_sec =  10#@param {type:\"integer\"}\n",
        "\n",
        "(start_pt, end_pt) = (start_sec, end_sec)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QpQvx6Oo8gpr",
        "outputId": "66a4da12-2cec-4c89-e2b2-639f62501ef6"
      },
      "outputs": [],
      "source": [
        "download_resolution = 360\n",
        "full_video_path = '/content/StridedTransformer-Pose3D/demo/video/full_video.mp4'\n",
        "file_name = 'input_clip.mp4'\n",
        "input_clip_path = '/content/StridedTransformer-Pose3D/demo/video/' + file_name\n",
        "\n",
        "# 動画ダウンロード\n",
        "ydl_opts = {'format': f'best[height<={download_resolution}]', 'overwrites': True, 'outtmpl': full_video_path}\n",
        "with YoutubeDL(ydl_opts) as ydl:\n",
        "    ydl.download([video_url])\n",
        "\n",
        "# 指定区間切り抜き\n",
        "with VideoFileClip(full_video_path) as video:\n",
        "    subclip = video.subclip(start_pt, end_pt)\n",
        "    subclip.write_videofile(input_clip_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dOmybPEWnefE"
      },
      "source": [
        "# 3D Human Pose Estimation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jqp2Ev6NnhXW",
        "outputId": "294cd7a8-06da-4d6b-d580-dc0a1b7434f5"
      },
      "outputs": [],
      "source": [
        "%cd /content/StridedTransformer-Pose3D\n",
        "\n",
        "!python demo/vis.py --video {file_name}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true,
          "base_uri": "https://localhost:8080/"
        },
        "id": "fvYjXBfB-RHJ",
        "outputId": "92711c25-6e21-4f61-a557-00ad322ee639"
      },
      "outputs": [],
      "source": [
        "clip = VideoFileClip('/content/StridedTransformer-Pose3D/demo/output/input_clip/input_clip.mp4')\n",
        "clip = resize(clip, height=420)\n",
        "clip.ipython_display()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "StridedTransformer-Pose3D.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}