【MXNet】Segmentation（動画ファイルにうつっている人を塗りつぶす）

はじめに

以前動画ファイルに対してPose Estimationをやった。
touch-sp.hatenablog.com
今回は動画ファイルに対してSegmentationをやってみた。

結果

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

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

精度はいまいち。
ただし後述するスクリプトをぜひ見てもらいたい。「こんなに短いのか？」と感じると思う。「MXNet」というフレームワークとそのエコシステムである「GluonCV」が非常に便利であることがわかる。

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

動画はこちらからダウンロードさせて頂いた。
（Video by Adesh Kumar Singh from Pixabay）
サイズは640×360を指定して「dance.mp4」に名前を変更した。

Pythonスクリプト

見やすくなるように空白行をいくつか入れているが、それを合わせてもたったの40数行。非常に短い。

import numpy as np
import mxnet as mx
from mxnet.gluon.data.vision import transforms
import gluoncv

import decord
import cv2

transform_fn = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([.485, .456, .406], [.229, .224, .225])
])

ctx = mx.gpu()
model = gluoncv.model_zoo.get_model('deeplab_resnet152_voc', pretrained=True, root='./models', ctx=ctx)

vr = decord.VideoReader( 'dance.mp4')
decord.bridge.set_bridge('mxnet')
all_frames = vr.get_batch(range(len(vr)))
fps = vr.get_avg_fps()
height, width = vr[0].shape[0:2]

imgs = transform_fn(all_frames).as_in_context(ctx)

fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
out = cv2.VideoWriter('output.mp4',fourcc, fps, (width,height))

for i in range(len(vr)):

    output = model.predict(imgs[i:i+1])
    
    predict = mx.nd.squeeze(mx.nd.argmax(output, 1)).asnumpy()
    a = np.repeat((predict==15)[:,:,None], 3, axis=2)
    
    original_img = vr[i].asnumpy()
    original_img[a] = 255
    img = cv2.cvtColor(original_img, cv2.COLOR_RGB2BGR)
    out.write(img)

out.release()

環境

エディション	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」の二つのみ。
「decord」や「opencv-python」などは勝手についてきた。
cuDNNはいれていない。（こちらを参照）

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

attrs==20.3.0
autocfg==0.0.6
autogluon.core==0.0.15b20201117
autograd==1.3
bcrypt==3.2.0
boto3==1.16.19
botocore==1.19.19
certifi==2020.11.8
cffi==1.14.3
chardet==3.0.4
click==7.1.2
cloudpickle==1.6.0
ConfigSpace==0.4.16
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.0b20201117
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.0b20201116
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.52.0
urllib3==1.26.2
yacs==0.1.8
zict==2.0.0
zipp==3.4.0

WSL2を使いたくないWindowsユーザーに

2020年11月17日現在WindowsではGluonCVのベータ版がインストールできない。
安定版のGluonCVをインストールして、別途decordとopencv-pythonをインストールする必要がある。

pip install mxnet-cu101 -f https://dist.mxnet.io/python/cu101
pip install gluoncv
pip install decord
pip install opencv-python

これで上記スクリプトは動作する。
「cu101」の部分はCUDA toolkit 10.1を表している。各自の環境に合わせて変更して下さい。

おまけ

GluonCVを使うと短いスクリプトでこんなことができます。そんな趣旨で書いた過去記事をのせておきます。
よかったら見て下さい。
touch-sp.hatenablog.com
touch-sp.hatenablog.com
touch-sp.hatenablog.com
touch-sp.hatenablog.com

2020年11月18日追記

Ubuntu on WSL2よりWindowsで実行する方が圧倒的に速い。
また、batchを使えばもっと速くなると思ってスクリプトを修正したがなぜかむしろ遅くなった。一応スクリプトを残しておく。

import numpy as np
import mxnet as mx
from mxnet.gluon.data.vision import transforms
import gluoncv

import decord
import cv2

transform_fn = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([.485, .456, .406], [.229, .224, .225])
])

ctx = mx.gpu()
model = gluoncv.model_zoo.get_model('deeplab_resnet152_voc', pretrained=True, root='./models', ctx=ctx)

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

dataset = mx.gluon.data.dataset.ArrayDataset(vr)
data_loader = mx.gluon.data.DataLoader(dataset, batch_size=10)

fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
out = cv2.VideoWriter('output_batch.mp4',fourcc, fps, (width,height))

for imgs_batch in data_loader:

    imgs = transform_fn(imgs_batch)
    output = model.predict(imgs.as_in_context(ctx))
    
    for i in range(output.shape[0]):
        predict = mx.nd.argmax(output[i], 0).asnumpy()
        a = np.repeat((predict==15)[:,:,None], 3, axis=2)
    
        original_img = imgs_batch[i].asnumpy()
        original_img[a] = 255
        img = cv2.cvtColor(original_img, cv2.COLOR_RGB2BGR)
        out.write(img)

out.release()

バッチ処理の際に使う「dataloader」については下記を参照。
touch-sp.hatenablog.com