DeepLabV3でSemantic Segmentation(gluoncv)2

初めに

新たに「deeplab_resnet152_voc」が公開されたのでデモのみ実行してみた

環境

Windows 10 Pro
GPUなし
Python 3.6.6(venv使用)

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

astroid==2.0.4
certifi==2018.8.24
chardet==3.0.4
colorama==0.3.9
cycler==0.10.0
gluoncv==0.3.0b20181010
graphviz==0.8.4
idna==2.6
isort==4.3.4
kiwisolver==1.0.1
lazy-object-proxy==1.3.1
matplotlib==3.0.0
mccabe==0.6.1
mxnet-mkl==1.3.1b20180927
numpy==1.14.6
Pillow==5.3.0
pylint==2.1.1
pyparsing==2.2.2
python-dateutil==2.7.3
requests==2.18.4
scipy==1.1.0
six==1.11.0
tqdm==4.26.0
typed-ast==1.1.0
urllib3==1.22
wrapt==1.10.11

実行ファイル

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.cpu()
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)

#モデルの適応
#人たけを抽出する(class:15)
output = model.demo(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:20181010195827p:plain