PoseNet Pythonでヘッドスピン動画を棒人間に変換してみる

8 min read読了の目安(約7300字

こんなサイトが話題になっていますね

このサイトを見た時に疑問に思ったのが「PoseNetはヘッドスピンのように逆さまになって回転している状態の人間を対象にちゃんと骨格検知できるのか?」ということです。気になったのでPoseNetのPython実装PoseNet Pythonを使って自分がヘッドスピンしている動画の棒人間化に挑戦してみました。

なお筆者は機械学習関連の知識が無いので技術的な詳細は解説しません。以後紹介するコードも適当ですがご容赦ください

環境

今回利用した環境です

  • OS: Mac OS X 10.15.5
  • Python: 3.6.5
  • Tensorflow: 1.12.0
  • OpenCV-Python:3.4.5.20
  • SciPy: 1.5.4
  • NumPy: 1.19.5

コードの準備

まずはPoseNet Pythonのリポジトリをクローンします。

$ git clone https://github.com/rwightman/posenet-python
Cloning into 'posenet-python'...
remote: Enumerating objects: 119, done.
remote: Total 119 (delta 0), reused 0 (delta 0), pack-reused 119
Receiving objects: 100% (119/119), 36.56 KiB | 4.57 MiB/s, done.
Resolving deltas: 100% (68/68), done.

続いて仮想環境を作成し、必要な依存ライブラリをインストールします。

$ cd posenet-python/
$ python -m venv .venv
$ source .venv/bin/activate
$ pip install tensorflow==1.12.0 opencv-python==3.4.5.20 scipy pyyaml

Collecting tensorflow==1.12.0
  Cache entry deserialization failed, entry ignored
  ...略

依存ライブラリがインストールできたらリポジトリに含まれているサンプルコードが動作することを確認しましょう。

$ python webcam_demo.py

骨格が検出できていればOKです

ヘッドスピンの動画を棒人間に変換してみる

先程のサンプルコードwebcam_demo.pyをコピー&加工し、カメラデバイスから取得した映像 or 動画ファイルから骨格を検出してmp4ファイルに出力してみます。まずはコードをコピー

$ cp webcam_demo.py conv_stick_human_movie.py

コピー後のconv_stick_human_movie.py を編集します。コード全体は以下のようになりました。

conv_stick_human_movie.py
import tensorflow as tf
import cv2
import time
import argparse
import numpy as np

import posenet

parser = argparse.ArgumentParser()
parser.add_argument('--model', type=int, default=101)
parser.add_argument('--cam_id', type=int, default=0)
parser.add_argument('--cam_width', type=int, default=1280)
parser.add_argument('--cam_height', type=int, default=720)
parser.add_argument('--scale_factor', type=float, default=0.7125)
parser.add_argument('--file', type=str, default=None, help="Optionally use a video file instead of a live camera")
args = parser.parse_args()


def main():
    with tf.Session() as sess:
        model_cfg, model_outputs = posenet.load_model(args.model, sess)
        output_stride = model_cfg['output_stride']

        if args.file is not None:
            cap = cv2.VideoCapture(args.file)
        else:
            cap = cv2.VideoCapture(args.cam_id)
        cap.set(3, args.cam_width)
        cap.set(4, args.cam_height)

        start = time.time()
        frame_count = 0

        output_video = None
        is_first = True
        try:
            while True:
                input_image, display_image, output_scale = posenet.read_cap(
                    cap, scale_factor=args.scale_factor, output_stride=output_stride)

                if is_first:
                    fmt = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
                    fps = cap.get(cv2.CAP_PROP_FPS)
                    output_video = cv2.VideoWriter('out.mp4', fmt, fps, (display_image.shape[1], display_image.shape[0]))
                    is_first = False


                heatmaps_result, offsets_result, displacement_fwd_result, displacement_bwd_result = sess.run(
                    model_outputs,
                    feed_dict={'image:0': input_image}
                )

                pose_scores, keypoint_scores, keypoint_coords = posenet.decode_multi.decode_multiple_poses(
                    heatmaps_result.squeeze(axis=0),
                    offsets_result.squeeze(axis=0),
                    displacement_fwd_result.squeeze(axis=0),
                    displacement_bwd_result.squeeze(axis=0),
                    output_stride=output_stride,
                    max_pose_detections=10,
                    min_pose_score=0.15)

                keypoint_coords *= output_scale

                skeleton_img = np.zeros(display_image.shape, dtype=np.uint8)
                skeleton_img = posenet.draw_skeleton(
                                skeleton_img, pose_scores, keypoint_scores, keypoint_coords,
                                min_pose_confidence=0.1, min_part_confidence=0.1)

                # 検出元の画像と検出後の画像を両方表示する
                # cv2.imshow('posenet', skeleton_img)
                # cv2.imshow('src', display_image)
                frame_count += 1
                output_video.write(skeleton_img)


                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break

        finally:
            if output_video:
                output_video.release()


if __name__ == "__main__":
    main()

順を追って解説していきます。まず

            if is_first:
                fmt = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
                fps = cap.get(cv2.CAP_PROP_FPS)
                output_video = cv2.VideoWriter('out.mp4', fmt, fps, (display_image.shape[1], display_image.shape[0]))
                is_first = False

の部分でカメラ or 動画ファイルから映像を読み込む初回ループ時のみout.mp4という名前で出力先のmp4ファイルを開いています。出力するファイルのFPSとサイズ(縦/横)は映像のソースから取得しcv2.VideoWriter()の引数で指定しています。

続いて

            skeleton_img = np.zeros(display_image.shape, dtype=np.uint8)
            skeleton_img = posenet.draw_skeleton(
                            skeleton_img, pose_scores, keypoint_scores, keypoint_coords,
                            min_pose_confidence=0.1, min_part_confidence=0.1)

            # 検出元の画像と検出後の画像を両方表示する
            # cv2.imshow('posenet', skeleton_img)
            # cv2.imshow('src', display_image)
            frame_count += 1
            output_video.write(skeleton_img)

の部分で映像のソースから骨格を検出して棒人間化し、出力先のmp4ファイルに書き出します。まずskeleton_img = np.zeros(display_image.shape, dtype=np.uint8)の部分で真っ黒な背景画像を用意し、続いて

skeleton_img = posenet.draw_skeleton(
	skeleton_img, pose_scores, keypoint_scores, keypoint_coords,
	min_pose_confidence=0.1, min_part_confidence=0.1)

で背景画像の上に骨格を描画します。min_pose_confidencemin_part_confidenceは0.1と低めの値に設定し、骨格が検出されやすいように設定しています。

コメントアウトしている

# cv2.imshow('posenet', skeleton_img)
# cv2.imshow('src', display_image)

の部分はソースの映像と検出された棒人間を同時に表示する処理です。デバッグ時に利用します。

最後の

finally:
  if output_video:
	  output_video.release()

で出力先のmp4ファイルのreleaseが確実に実行されるようにハンドリングしています。ソースの映像を読み込むposenet.read_capを使って動画ファイルを読み込んだ際、動画の最後まで読み込むとIOErrorraiseされるので、その対策として組み込んでいます。

コードが準備できたら適当な動画を引数に指定してスクリプトを実行します。※もしくはファイルを指定せずに実行し、カメラの前でヘッドスピンを回っても良いです

$ python conv_stick_human_movie.py --file <適当な動画ファイル>

実行が完了するとout.mp4というファイルに棒人間のヘッドスピン動画が出力されているはずです。

$ file out.mp4 
out.mp4: ISO Media, MP4 Base Media v1 [IS0 14496-12:2003]

再生してみましょう!!

お、いい感じか??

...と思ったのですが、逆さまになったあたりから全く骨格が検出されなくなったのか、延々と真っ黒な背景が続きました...

先程のデバッグ用コードをコメントアウトして再実行してみます

cv2.imshow('posenet', skeleton_img)
cv2.imshow('src', display_image)



逆さになったあたりから正しく骨格検知できてなくなっているようです。。。ヘッドスピンの形やスピード、服なんかも関係しているんでしょうかね??とにかく、Pythonのスクリプト1発で棒人間のヘッドスピン動画を作成するのはそう簡単では無さそうです。