DeepLabV3でSemantic Segmentation(gluoncv)

初めに

「deeplab_resnet152_voc」のデモのみ実行してみた。
以前も同様の記事を書いたが最新のmxnetではコードを一部書き換える必要があった。

環境

Windows10 Pro
NVIDIA GeForce GTX1080
CUDA 10.1
Python 3.6.8

バージョンの確認(pip freeze)

certifi==2019.6.16
chardet==3.0.4
cycler==0.10.0
gluoncv==0.5.0b20190812
graphviz==0.8.4
idna==2.6
kiwisolver==1.1.0
matplotlib==3.1.1
mxnet-cu101==1.6.0b20190812
numpy==1.16.4
Pillow==6.1.0
pyparsing==2.4.2
python-dateutil==2.8.0
requests==2.18.4
scipy==1.3.1
six==1.12.0
tqdm==4.33.0
urllib3==1.22

実際にインストールしたのは「mxnet」と「gluoncv」のみ。

pip install mxnet-cu101 --pre
pip install gluoncv --pre

実行ファイル

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

ctx = mx.gpu()
filename = 'example.jpg'

#画像をNDArryで読み込む
img = image.imread(filename)

#データの正規化
transform_fn = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([.485, .456, .406], [.229, .224, .225])
])
img = transform_fn(img)
img = img.expand_dims(0).as_in_context(ctx)

#モデルを読み込む
#初回時に(default)/.mxnet/modelsに保存される
#2回目以降はそこから読み込む
model= gluoncv.model_zoo.get_model('deeplab_resnet152_voc', pretrained=True)
model.collect_params().reset_ctx(ctx)
#モデルの適応
#人たけを抽出する(class:15)
output = model.predict(img)
predict = mx.nd.squeeze(mx.nd.argmax(output, 1)).asnumpy()
a = np.where(predict ==15, 255, 0)
b = Image.fromarray(a).convert('L')

#画像を改めてPILで読み込み、結果と重ね合わせる
img = Image.open(filename)
img.putalpha(b)
img.save('result.png')

結果

f:id:touch-sp:20190717192024p:plain