【Ubuntu 22.04】【MXNet】【GluonCV】Webカメラの動画に対してリアルタイムに人物以外の背景を消す

はじめに
結果
Pythonスクリプト
- 修正前
- 修正後
動作環境
環境構築

はじめに

以前に書いたスクリプトの動作確認になります。
touch-sp.hatenablog.com

今回はUbuntu 22.04 on WSL2を使用しています。

PythonもMXNetもGluonCVも新しいバージョンになっています。

結果

問題なく動作しました。

Pythonスクリプト

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

ctx = mx.gpu() if mx.context.num_gpus() >0 else mx.cpu()
transform_fn = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([.485, .456, .406], [.229, .224, .225])
])
model = get_model('deeplab_resnet152_voc', pretrained=True, root='./models', ctx=ctx)
cap = cv2.VideoCapture(0)
while True:
    ret, frame = cap.read()
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    img = transform_fn(mx.nd.array(frame_rgb))
    img = img.expand_dims(0).as_in_context(ctx)
    output = model.predict(img)
    predict = mx.nd.squeeze(mx.nd.argmax(output, 1)).asnumpy()
    mask_1 = np.where(predict == 15, 1, 0)[...,np.newaxis]
    mask_2 = np.where(predict == 15, 0, 255)[...,np.newaxis]
    result_img = (frame * mask_1 + mask_2).astype('uint8')
    cv2.imshow('result', result_img)
    if cv2.waitKey(1) & 0xFF == 27:
        break
cap.release()
cv2.destroyAllWindows()

1行だけWindows用からUbuntu用に書き換えました。

修正前

cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)

修正後

cap = cv2.VideoCapture(0)

動作環境

使用したPCはこちらです。

Ubuntu 22.04 on WSL2 (Windows 11)
Python 3.10.4

CUDA toolkit 11.4
cuDNN 8.4.0

autocfg==0.0.8
certifi==2022.5.18
charset-normalizer==2.0.12
cycler==0.11.0
fonttools==4.33.3
gluoncv==0.10.5.post0
graphviz==0.8.4
idna==3.3
kiwisolver==1.4.2
matplotlib==3.5.2
mxnet-cu112==1.9.1
numpy==1.22.3
opencv-python==4.5.5.64
packaging==21.3
pandas==1.4.2
Pillow==9.1.1
portalocker==2.4.0
pyparsing==3.0.9
python-dateutil==2.8.2
pytz==2022.1
PyYAML==6.0
requests==2.27.1
scipy==1.8.1
six==1.16.0
tqdm==4.64.0
urllib3==1.26.9
yacs==0.1.8

CUDA toolkit 11.4でmxnet-cu112が問題なく動いています。