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