MMD自動トレースで OpenPose の代わりに tf‐pose‐estimation を使う - errno-mmd/tf-pose-estimation GitHub Wiki
概要
MMD自動トレース をDockerで動かせるようにした が、GPUを使わずCPUで処理を行うことになり、一連の処理の中で特にOpenPoseが非常に遅かった。そこで、OpenPoseの代わりにtf-pose-estimationを使うことで処理を高速化した。
背景
動画に写っている人の動きを自動的に解析して、MMD用のモーションデータを生成する MMDモーショントレース自動化(以下、MMD自動トレース)の開発がmiuさん(+協力者)により行われています。 MMD自動トレースでは、まずOpenPoseによる二次元姿勢推定を行い、その結果に基づいて三次元姿勢推定、センター位置推定、VMDデータへの変換等を行います。
各処理に用いるツールやライブラリのインストールが面倒なので、Dockerコンテナイメージに纏めた のですが、Docker for WindowsはGPUでの処理に対応していない(nvidia-dockerが使えない)ので、CPUで処理するしかありません。 ところが、OpenPoseはGPUを使わないと非常に遅いという問題があります。PCの性能や割り当てるCPUコア数によりますが、Core i7-6700 3.4GHzの4仮想コアを使った場合、1分の動画の処理に2時間弱かかり、そのうち9割がOpenPoseの処理時間でした。
そこで、このOpenPoseの処理を何とか高速化できないかと調べていたところ、tf-pose-estimation というプログラムが見つかりました。
tf-pose-estimationとは
tf-pose-estimationは OpenPose と同じく、動画に写った人の二次元のポーズを推定するプログラムです。 元々OpenPoseと同じニューラルネットワークをTensorflow向けに実装したものだったのですが、新たにMobileNet v1 、v2 のニューラルネットワークモデルが順次追加され、より高速に処理できるようになっています。
下の図はOpenPoseとtf-pose-estimation(MobileNet v2)の推定結果の比較です。若干の差異はあるのですが、どちらも同じような結果が得られています。
ただし、OpenPoseと同じ機能が全てtf-pose-estimationに実装されているわけではありません。tf-pose-estimationをOpenPoseの代わりに使うには、MMD自動トレースで使われているOpenPoseの機能を tf-pose-estimation に追加する必要がありました。
tf-pose-estimationへの機能追加
MMD自動トレースで使う下記のオプションに対応しました。上から順に改造内容を説明していきます。
--write_json 推定した二次元の関節位置情報をJSONフォーマットでファイルに出力します。
--write_video ポーズ推定の結果を動画ファイルとして出力します。
--number_people_max ポーズ推定の対象となる人数の最大値を指定します。
--frame_first 解析対象となる最初のフレームの番号を指定します。
--no_display 処理中に推定結果の画面表示を行わないようにします。
--write_json オプション
ArashHosseiniによる改造をベースに、JSONファイルの出力をOpenPoseの新しいフォーマットに合わせるよう更に改造を行いました。なお ArashHosseini のリポジトリは更新が止まっているため、私のリポジトリは、オリジナルのildoonet/tf-pose-estimationからforkし、ArashHosseini の改造は cherry-pick で取り込むことにしました。
JSONファイル出力処理は tf_pose/estimator.py の draw_humans() 関数に埋め込まれています。出力自体にはPython標準ライブラリの json.dump() が使われています。
429 if output_json_dir:
430 with open(os.path.join(output_json_dir, '{0}_keypoints.json'.format(str(frame).zfill(12))), 'w') as outfile:
431 json.dump(dc, outfile)
その前の、ポーズ推定結果を表示する処理の中で辞書 dc に関節位置などの情報が格納されています。 余談ですが、他人の書いた機械学習等のコードをしっかり理解することなく結果を取り出して利用する場合、 このように結果を画面などに出力するところで出力対象の情報を横取りするのが楽だと思います。
--write_video オプション
run_video.py の main 処理の中で、先ほどの draw_humans()関数を呼んでポーズ推定結果を描画している箇所があります。
73 image = TfPoseEstimator.draw_humans(image, humans, imgcopy=False, frame=frame, output_json_dir=args.write_json)
したがって、この image を動画ファイルに書き出す機能さえ追加すれば改造完了です。OpenCV を使っているので、動画ファイル出力は非常に簡単です。
55 fourcc = cv2.VideoWriter_fourcc(*'XVID')
56 out = cv2.VideoWriter(args.write_video, fourcc, 30.0, (width, height))
で動画ファイル出力の準備をして、1フレームずつ作成した image を出力していきます。
79 out.write(image)
最後に release() を呼んでファイルを閉じます。
87 out.release()
これだけ。
--number_people_max オプション
映像に人っぽい物が写っていたり、奥の方に鏡があったりすると、TfPoseEstimator の inference() が本来推定してほしい人数より多くの結果を返すことがあります。そのため、--number_people_max オプションで指定された人数より後の結果を捨てる処理を入れてあります。
69 humans = e.inference(image, resize_to_default=(w > 0 and h > 0), upsample_size=args.resize_out_ratio)
70 del humans[args.number_people_max:]
--frame_first オプション
フレーム番号が frame_first に達するまで、処理をスキップするだけです。
64 if frame < args.frame_first:
65 frame += 1
66 continue
--no_display オプション
run_video.py はデフォルトでポーズ推定結果を画面表示するのですが、特にコンテナで実行する場合に邪魔なので、このオプションを付けました。
76 if not args.no_display:
77 cv2.imshow('tf-pose-estimation result', image)
その他改造した点
あと、OpenPoseは処理中にメッセージを出さないため、動いているのか止まっているのか分かりにくかったのですが、tf-pose-estimationでは処理が何フレーム目まで進んだか表示するようにしました。
実際に機能追加を行ったソースコードの差分を見たい方は、errno-mmd/tf-pose-estimationの develop ブランチを参照してください。
処理速度の比較
評価環境:
- OS: Windows 10 Pro 64bit
- Docker for Windows: Docker Desktop 2.0.0.3
- CPU: Core i7-6700 3.4GHzの4仮想コアをDockerに割り当て
- Memory: 8GB をDockerに割り当て
- バッチ: errno-mmd/motion_trace_bulk の Docker/MotionTraceBulk_Docker.bat
使用した動画:
- 再生時間: 1分4秒
- サイズ: 1280x720
- フレームレート: 30fps
- 動画に写っている人の数: 1人
結果は下記のとおり。単位は[hour]
旧バージョン | 新バージョン | |
---|---|---|
OpenPose | tf-pose-estimation | |
2Dポーズ推定時間 | 1.61 | 0.25 |
その他の時間 | 0.18 | 0.18 |
合計 | 1.79 | 0.42 |
2Dポーズ推定の処理速度が6.5倍になり、その結果トータルの処理速度は4.2倍になりました。
使い方(1) Docker版を使う場合
既に Docker imageおよび 一括処理バッチは OpenPose から tf-pose-estimation に変更済みなので、Docker版の紹介動画 で説明したとおりの手順で動かすことができます。
具体的には、Docker for Windowsをインストール、上の紹介動画で説明したとおり共有ドライブなどの設定をします。 Dockerの準備ができたら、https://github.com/errno-mmd/motion_trace_bulk から一括処理バッチの ZIP ファイルをダウンロード、展開して、Docker/MotionTraceBulk_Docker.bat を実行します。 バッチ実行時に入力する内容は、バッチのZIPファイルに同梱されているREADME.mdを参照してください。
使い方(2) ネイティブの環境で使う場合
errno-mmd/tf-pose-estimation の README.mdの Install セクションに従ってインストールを行います。
MMD自動トレース用に2Dポーズ推定を行う場合、下記のようにオプションを付けて run_video.pyを実行します。
python3 run_video.py --video test1.mp4 --model mobilenet_v2_large --output_json output/test1/json --no_display --number_people_max 1
なお下記のように --help オプションを付けて実行すると、オプションの説明が表示されます。
python3 run_video.py --help