Semantic Segmentation 結果を元画像に表示（deeplab_resnet101）

初めに

チュートリアルに方法が書いていないので簡単そうでなかなか難しかった。

環境

Windows10 Pro
GPUなし
Python 3.8.2

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

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

pip install mxnet
pip install gluoncv

その他は勝手についてくる。

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==1.6.0
numpy==1.16.6
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
from PIL import Image
import mxnet as mx
from mxnet.gluon.data.vision import transforms
import gluoncv

filename = gluoncv.utils.download('https://github.com/dmlc/web-data/blob/master/' +
                          'gluoncv/detection/biking.jpg?raw=true',
                          path='biking.jpg')

#画像をPILで読み込む
img_pil = Image.open(filename)

#画像をmx.ndarrayに変換
img = mx.nd.array(np.asarray(img_pil))

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

#モデルを読み込む
#初回時に(default)/.mxnet/modelsに保存される
#2回目以降はそこから読み込む
model= gluoncv.model_zoo.get_model('deeplab_resnet101_voc', pretrained=True)

#モデルの適応
output = model.predict(img)
predict = mx.nd.squeeze(mx.nd.argmax(output, 1)).asnumpy()

mask = gluoncv.utils.viz.get_color_pallete(predict, 'pascal_voc')
m = mask.convert('RGB')
m.putalpha(128)

#元の画像と結果を重ね合わせる
img_pil.paste(m, m)
img_pil.save('result.png')

結果