【Realtime Segmentation】【GluonCV】【Streamlit】Webカメラの動画に対してぼかし強度をインタラクティブに変更しながらリアルタイムに人物以外をぼかす

GluonCV segmentation opencv-pythonで動画を扱う Streamlit

はじめに

前回背景を消す（白く塗りつぶす）ことをしました。
touch-sp.hatenablog.com
今回はぼかしを入れてみます。それだけだとほとんどスクリプトが変わらないのでさらにStreamlitを使ってインタラクティブにぼかし強度を変更できるようにしました。

Pythonスクリプト

import streamlit as st
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])
])

@st.cache(allow_output_mutation=True)
def load_model():
    model = get_model('deeplab_resnet152_voc', pretrained=True, root='./models', ctx=ctx)
    return model
model = load_model()

strength = st.slider('Blur Strength', 10, 30, 15, step=5)
imagelocation = st.empty()

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

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, 1)[...,np.newaxis]
    blurred_img = cv2.blur(frame, (strength, strength))
    result_img = (frame * mask_1 + blurred_img * mask_2).astype('uint8')
    result_img = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
    imagelocation.image(result_img)

実行する時は以下のようにして下さい。
上のスクリプトの名前を「script.py」とした場合です。

streamlit run script.py

終了する時はターミナル上で「Ctrl」+「c」を押して下さい。

動作環境

Windows 10 PC with GTX 1080
CUDA 10.2
Python 3.7.9

Python環境にインストールしたのは「mxnet」「gluoncv」「streamlit」の3つだけです。
すべてpipでインストール可能です。ただしWindows用のGPU版mxnetはダウンロード先を指定する必要があります。

pip install mxnet-cu102 -f https://dist.mxnet.io/python/cu102
pip install gluoncv
pip install streamlit

mxnet-cu102==1.7.0
gluoncv==0.10.0
streamlit==0.78.0

スクリプトの解説

opencv-pythonとgluoncv(mxnet)の基本的な使い方は前回の記事を参照して下さい。

背景をぼかすためのスクリプトは以下の数行です。

mask_1 = np.where(predict == 15, 1, 0)[...,np.newaxis]
mask_2 = np.where(predict == 15, 0, 1)[...,np.newaxis]
blurred_img = cv2.blur(frame, (strength, strength))
result_img = (frame * mask_1 + blurred_img * mask_2).astype('uint8')

cv2.blurでぼかし画像を作ります。strenghtの部分を変更することによってぼかし強度が変わります。
元の画像の人物以外の部分に0を掛けています。
ぼかし画像の人物部分に0を掛けています。
両者を足し合わせることで人物以外にぼかしが入った画像（numpy array）が完成します。

Streamlitを使うために変更・追加した部分は以下の数行です。

import streamlit as st

@st.cache(allow_output_mutation=True)
def load_model():
    model = get_model('deeplab_resnet152_voc', pretrained=True, root='./models', ctx=ctx)
    return model
model = load_model()

strength = st.slider('Blur Strength', 10, 30, 15, step=5)

blurred_img = cv2.blur(frame, (strength, strength))

imagelocation = st.empty()

imagelocation.image(result_img)

スライダーを変更するたびにモデルを読み込むのを避けるためload_modelという関数を定義しています。
StreamlitのimageはRGB配列のnumpy arrayを入力するだけで表示してくれます。

2021年6月9日追記

以下の環境でも動作しました。

Python 3.7.9

gluoncv==0.10.2
mxnet-cu102==1.7.0
streamlit==0.82.0