【MXNet】Pose Estimation （ファイルからファイルへ）

Pose Estimation

はじめに

「gluoncv」のベータ版（0.9.0b）をインストールするといろいろなものが一緒にインストールされる。
その中に「decord」と言うものがある。
動画を読み込む時に使うもので「opencv-python」の代わりとして使える。
今回は「decord」を使ってみるためにPose Estimationのスクリプトを書いた。
ちなみに「opencv-python」も「gluoncv」と一緒にインストールされる。

環境

エディション	Windows 10 Pro
バージョン	2004
OS ビルド	20257.1

Ubuntu 18.04 on WSL2
Python 3.7.5

GTX 1080
NVIDA driver 460.20
CUDA Toolkit 11.0


Python環境にインストールしたのは「mxnet-cu110」と「gluoncv」の二つのみ。
その他は勝手についてきた。

pip install mxnet-cu110==1.9.0b20201113 -f https://dist.mxnet.io/python/cu110
pip install gluoncv --pre

attrs==20.3.0
autocfg==0.0.6
autogluon.core==0.0.15b20201114
autograd==1.3
bcrypt==3.2.0
boto3==1.16.18
botocore==1.19.18
certifi==2020.11.8
cffi==1.14.3
chardet==3.0.4
click==7.1.2
cloudpickle==1.6.0
ConfigSpace==0.4.10
cryptography==3.2.1
cycler==0.10.0
Cython==3.0a6
dask==2.30.0
decord==0.4.2
distributed==2.30.1
future==0.18.2
gluoncv==0.9.0b20201114
graphviz==0.8.4
HeapDict==1.0.1
idna==2.10
importlib-metadata==2.0.0
iniconfig==1.1.1
jmespath==0.10.0
joblib==0.17.0
kiwisolver==1.3.1
matplotlib==3.3.3
msgpack==1.0.0
mxnet-cu110==1.9.0b20201113
numpy==1.19.4
opencv-python==4.4.0.46
packaging==20.4
pandas==1.1.4
paramiko==2.7.2
Pillow==8.0.1
pkg-resources==0.0.0
pluggy==0.13.1
portalocker==2.0.0
protobuf==3.14.0
psutil==5.7.3
py==1.9.0
pyaml==20.4.0
pycparser==2.20
PyNaCl==1.4.0
pyparsing==3.0.0b1
pytest==6.1.2
python-dateutil==2.8.1
pytz==2020.4
PyYAML==5.3.1
requests==2.25.0
s3transfer==0.3.3
scikit-learn==0.23.2
scikit-optimize==0.8.1
scipy==1.5.4
six==1.15.0
sortedcontainers==2.3.0
tblib==1.7.0
tensorboardX==2.1
threadpoolctl==2.1.0
toml==0.10.2
toolz==0.11.1
tornado==6.1
tqdm==4.51.0
typing==3.7.4.3
urllib3==1.26.2
yacs==0.1.8
zict==2.0.0
zipp==3.4.0

動画ファイルのダウンロード

動画はこちらからダウンロードした。
サイズは640×360を指定して「dance.mp4」に名前を変更した。

Pythonスクリプト

動画ファイルの読み込みに「decord」を使用した。
動画ファイルへの書き出しには「opencv-python」を使用した。

import decord
import cv2
import mxnet as mx
import gluoncv
from gluoncv.model_zoo import get_model
from gluoncv.data.transforms.pose import detector_to_alpha_pose, heatmap_to_coord
from gluoncv.utils.viz import cv_plot_keypoints

ctx = mx.gpu()
video_fname = 'dance.mp4'

vr = decord.VideoReader(video_fname)
decord.bridge.set_bridge('mxnet')
fps = vr.get_avg_fps()
height, width = vr[0].shape[0:2]

# Define the codec and create VideoWriter object
fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
out = cv2.VideoWriter('output.mp4',fourcc, fps, (width,height))

detector = get_model('ssd_512_mobilenet1.0_coco', pretrained=True, ctx=ctx, root='./models')
detector.reset_class(classes=['person'], reuse_weights={'person':'person'})
detector.hybridize()

estimator = get_model('alpha_pose_resnet101_v1b_coco', pretrained=True, ctx=ctx, root='./models')
estimator.hybridize()

for each_frame in vr:

    x, frame = gluoncv.data.transforms.presets.ssd.transform_test(each_frame, short=512)
    class_IDs, scores, bounding_boxs = detector(x.as_in_context(ctx))
    pose_input, upscale_bbox = detector_to_alpha_pose(frame, class_IDs, scores, bounding_boxs)

    if upscale_bbox is not None:
        predicted_heatmap = estimator(pose_input.as_in_context(ctx))
        pred_coords, confidence = heatmap_to_coord(predicted_heatmap, upscale_bbox)

        #bounding_boxを表示しないためのダミーデータ
        scores = mx.nd.zeros(shape=(1,100,1))
        
        img = cv_plot_keypoints(frame, pred_coords, confidence, class_IDs, bounding_boxs, scores,
                                box_thresh=0.5, keypoint_thresh=0.2)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        img = cv2.resize(img, (width, height))
        out.write(img)
    else:
        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        frame = cv2.resize(img, (width, height))
        out.write(frame)

out.release()

結果

「output.mp4」というファイルが作成される。
GIFに変換したものをのせておく。

ffmpeg -i output.mp4 -vf scale=320:-1 output.gif

f:id:touch-sp:20201114204639g:plain