【MXNet】動画に対してSegmentation（fcn_resnet101）

初めに

以前GPUなしのノートパソコンでやってみたがパワー不足でスムーズに動きませんでした。今回GPU搭載のデスクトップパソコンでやってみました。

環境

HP 870-281jp
CPU Intel(R) Core(TM) i7-7700K CPU @ 4.20GHz 4.20GHz
RAM 32.0 GB
NVIDIA GeForce GTX1080

Windows10 Pro
Python 3.7.7

カメラがついていないので「EpocCam」というアプリ（無料）を利用してiPadをカメラとして利用しています。
特に難しい設定はいりません。

バージョンの確認（pip freeze）

インストールが必要なのは「mxnet-cu101]と「gluoncv」と「opencv-python」のみです。

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

その他は勝手についてきます。

certifi==2020.4.5.1
chardet==3.0.4
cycler==0.10.0
gluoncv==0.7.0
graphviz==0.8.4
idna==2.6
kiwisolver==1.2.0
matplotlib==3.2.1
mxnet-cu101 @ https://repo.mxnet.io/dist/python/cu101/mxnet_cu101-1.6.0-py2.py3-none-win_amd64.whl
numpy==1.16.6
opencv-python==4.2.0.34
Pillow==7.1.2
portalocker==1.7.0
pyparsing==2.4.7
python-dateutil==2.8.1
pywin32==227
requests==2.18.4
scipy==1.4.1
six==1.14.0
tqdm==4.46.0
urllib3==1.22

実行ファイル

import numpy as np
import mxnet as mx
import gluoncv
from gluoncv.utils.viz import get_color_pallete, cv_merge_two_images, cv_plot_image
from mxnet.gluon.data.vision import transforms
import time
import cv2

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

# Load the model
model = gluoncv.model_zoo.get_model('fcn_resnet101_voc', pretrained=True)
model.collect_params().reset_ctx(ctx)
# Compile the model for faster speed
model.hybridize()

# Load the webcam handler
cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
# letting the camera autofocus
time.sleep(1)

while(True):
    # Load frame from the camera
    ret, frame = cap.read()

    # Image pre-processing
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    frame_mx = mx.nd.array(frame, ctx=ctx)
    img = transform_fn(frame_mx)
    img = img.expand_dims(0)
    # Run frame through network
    output = model.predict(img)
    predict = mx.nd.squeeze(mx.nd.argmax(output, 1)).asnumpy()

    mask = get_color_pallete(predict, 'pascal_voc')
    m = mask.convert('RGB')
    n = np.asarray(m)
    o = cv_merge_two_images(frame, n)
    cv_plot_image(o)

    # escを押したら終了
    if cv2.waitKey(1) == 27:
        break

cap.release()
cv2.destroyAllWindows()

補足

カメラがうまく検出されない場合には実行ファイルの中の次の部分を少し変えてみると良いと思います。

# Load the webcam handler
cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)

数字を「0」から「1」や「2」に変えてみる。