【python / mxnet / gluoncv / jupyter notebook】基于mxnet和gluoncv的图像分割
程序环境为高性能集群：
CPU：Intel Xeon Gold 6140 Processor * 2（共36核心）
内存：512GB RAM
GPU：Tesla P100-PCIE-16GB * 2
Test with DeepLabV3 Pre-trained Models¶
This is a quick demo of using GluonCV DeepLabV3 model on ADE20K dataset. Please follow the installation guide <../index.html>_ to install MXNet and GluonCV if not yet.
In [2]:
import mxnet as mx
from mxnet import image
from mxnet.gluon.data.vision import transforms
import gluoncv
# using cpu
ctx = mx.cpu(0)
Prepare the image¶
download the example image
In [3]:
'''
url = 'https://github.com/zhanghang1989/image-data/blob/master/encoding/' + \
'segmentation/ade20k/ADE_val_00001755.jpg?raw=true'
filename = 'ade20k_example.jpg'
gluoncv.utils.download(url, filename, True)
'''
filename = 'data/2/nju-cs.jpg'
load the image
In [4]:
img = image.imread(filename)
from matplotlib import pyplot as plt
plt.imshow(img.asnumpy())
plt.show()
normalize the image using dataset mean
In [5]:
from gluoncv.data.transforms.presets.segmentation import test_transform
img = test_transform(img, ctx)
Load the pre-trained model and make prediction¶
get pre-trained model
In [6]:
model = gluoncv.model_zoo.get_model('deeplab_resnet101_ade', pretrained=True)
make prediction using single scale
In [7]:
output = model.predict(img)
predict = mx.nd.squeeze(mx.nd.argmax(output, 1)).asnumpy()
Add color pallete for visualization
In [8]:
from gluoncv.utils.viz import get_color_pallete
import matplotlib.image as mpimg
mask = get_color_pallete(predict, 'ade20k')
mask.save('data/4/output-nju-cs.png')
show the predicted mask
In [9]:
mmask = mpimg.imread('data/4/output-nju-cs.png')
plt.imshow(mmask)
plt.show()
More Examples¶
In [10]:
filename = 'data/2/1.jpg'
img = image.imread(filename)
plt.imshow(img.asnumpy())
plt.show()
img = test_transform(img, ctx)
model = gluoncv.model_zoo.get_model('deeplab_resnet101_ade', pretrained=True)
output = model.predict(img)
predict = mx.nd.squeeze(mx.nd.argmax(output, 1)).asnumpy()
mask = get_color_pallete(predict, 'ade20k')
mask.save('data/4/output1.png')
mmask = mpimg.imread('data/4/output1.png')
plt.imshow(mmask)
plt.show()
In [12]:
filename = 'data/2/train.jpg'
img = image.imread(filename)
plt.imshow(img.asnumpy())
plt.show()
img = test_transform(img, ctx)
model = gluoncv.model_zoo.get_model('deeplab_resnet101_ade', pretrained=True)
output = model.predict(img)
predict = mx.nd.squeeze(mx.nd.argmax(output, 1)).asnumpy()
mask = get_color_pallete(predict, 'ade20k')
mask.save('data/4/output-train.png')
mmask = mpimg.imread('data/4/output-train.png')
plt.imshow(mmask)
plt.show()
tz@croplab,hzau
