本記事では、機械学習手法StridedTransformer-Pose3Dを用いて、任意の動画に映る人物の3次元姿勢推定を行う方法をご紹介します。

出典: Vegetebird/StridedTransformer-Pose3D

StridedTransformer

概要

Strided Transformerは、Transformerベースの2D pose sequenceを活用し、3D pose sequenceを生成する姿勢推定手法です。

出典: Exploiting Temporal Contexts with Strided Transformer for 3D Human Pose Estimation

Strided Transformerでは、Vanilla Transformer Encoder(VTE)を使用し、2D poseシーケンスのlong-rangeの依存関係をモデル化します。この時、シーケンスの冗長性を減らすため、VTEに繋がるレイヤーは、シーケンスの長さを徐々に縮小し、ローカルコンテキストからの情報を集約するためstrided CNNに置き換えられています。

この変更されたVTEがStridedTransformerEncoder(STE)と呼ばれ、long-rangeな情報を効果的に集約し、計算コストの大幅な削減を実現しています。
このような構成で実現されるStridedTransformerは、ベンチマークデータセットHuman3.6MとHumanEva-IにおいてSOTAを達成しています。

出典: Exploiting Temporal Contexts with Strided Transformer for 3D Human Pose Estimation

詳細はこちらの論文をご参照ください。

本記事では上記手法を用いて、任意の動画の3次元姿勢推定を行います。

デモ(Colaboratory)

それでは、実際に動かしながら任意の動画の3次元姿勢推定を行います。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。

また、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

初めにGithubからソースコードを取得します。

%cd /content

!git clone https://github.com/Vegetebird/StridedTransformer-Pose3D.git

次にライブラリをインストールします。

%cd /content/StridedTransformer-Pose3D

!pip install --upgrade gdown
!pip install yacs
!pip install filterpy
!pip install einops
!pip install yt-dlp moviepy
!pip install matplotlib==3.0.3

matplotlibのみバージョンにご注意ください。repositoryのコード内でmatplotlib 3.1.0以上では廃止された関数を使用しているため3.0.3を明示的にインストールしています。

最後にライブラリをインポートします。

%cd /content/StridedTransformer-Pose3D

import os
from yt_dlp import YoutubeDL

from moviepy.video.fx.resize import resize
from moviepy.editor import VideoFileClip, AudioFileClip, ImageSequenceClip, CompositeAudioClip
from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_subclip

以上で環境セットアップは完了です。

学習済みモデルのセットアップ

次に論文発表元が公開する学習済みモデルをダウンロードします。

%cd /content/StridedTransformer-Pose3D
!mkdir -p ./checkpoint/pretrained

if not os.path.exists('checkpoint/pretrained/refine_4365.pth'):
  !gdown https://drive.google.com/uc?id=1aDLu0SB9JnPYZOOzQsJMV9zEIHg2Uro7 -O checkpoint/pretrained/refine_4365.pth
if not os.path.exists('checkpoint/pretrained/no_refine_4365.pth'):
  !gdown https://drive.google.com/uc?id=1l63AI9BsNovpfTAbfAkySo9X2MOWgYZH -O checkpoint/pretrained/no_refine_4365.pth

if not os.path.exists('demo/lib/checkpoint/yolov3.weights'):
  !gdown https://drive.google.com/uc?id=1gWZl1VrlLZKBf0Pfkj4hKiFxe8sHP-1C -O demo/lib/checkpoint/yolov3.weights
if not os.path.exists('demo/lib/checkpoint/pose_hrnet_w48_384x288.pth'):
  !gdown https://drive.google.com/uc?id=1CpyZiUIUlEjiql4rILwdBT4666S72Oq4 -O demo/lib/checkpoint/pose_hrnet_w48_384x288.pth

テスト動画のセットアップ

3D姿勢推定を行う入力動画をダウンロードします。
本記事では、Youtubeから動画をダウンロードします。

%cd /content/StridedTransformer-Pose3D/demo/video

video_url = 'https://www.youtube.com/watch?v=VSSPwzSflr0' #@param {type:"string"}

#@markdown 動画の切り抜き範囲(秒)を指定してください。\
#@markdown 30秒以上の場合OOM発生の可能性が高いため注意
start_sec =  1#@param {type:"integer"}
end_sec =  10#@param {type:"integer"}

(start_pt, end_pt) = (start_sec, end_sec)

download_resolution = 360
full_video_path = '/content/StridedTransformer-Pose3D/demo/video/full_video.mp4'
file_name = 'input_clip.mp4'
input_clip_path = '/content/StridedTransformer-Pose3D/demo/video/' + file_name

# 動画ダウンロード
ydl_opts = {'format': f'best[height<={download_resolution}]', 'overwrites': True, 'outtmpl': full_video_path}
with YoutubeDL(ydl_opts) as ydl:
    ydl.download([video_url])

# 指定区間切り抜き
with VideoFileClip(full_video_path) as video:
    subclip = video.subclip(start_pt, end_pt)
    subclip.write_videofile(input_clip_path)

Youtubeから動画をダウンロードし、指定の秒数の区間切り抜き、入力動画とします。

3D Human Pose Estimation

最後に3次元の姿勢推定を行います。

%cd /content/StridedTransformer-Pose3D

!python demo/vis.py --video {file_name}

clip = VideoFileClip('/content/StridedTransformer-Pose3D/demo/output/input_clip/input_clip.mp4')
clip = resize(clip, height=420)
clip.ipython_display()

出力結果は以下の通りです。

2Dの姿勢推定は非常に精巧です。一方で3Dの姿勢推定は、速い体の動きにややついていけていない箇所が見受けられます。

まとめ

本記事では、StridedTransformerを用いて任意の動画の3次元の姿勢推定を行いました。
2Dの姿勢推定結果から、3Dの姿勢推定を再構築するという発想が興味深いです。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。

また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - Exploiting Temporal Contexts with Strided Transformer for 3D Human Pose Estimation

2. GitHub - Vegetebird/StridedTransformer-Pose3D