2. 使用 PSPNet 预训练模型进行测试

这是一个在 ADE20K 数据集上使用 GluonCV PSPNet 模型的快速演示。如果您尚未安装 MXNet 和 GluonCV，请按照安装指南进行安装。

import mxnet as mx
from mxnet import image
from mxnet.gluon.data.vision import transforms
import gluoncv
# using cpu
ctx = mx.cpu(0)

准备图像

下载示例图像

url = 'https://github.com/zhanghang1989/image-data/blob/master/encoding/' + \
    'segmentation/ade20k/ADE_val_00001142.jpg?raw=true'
filename = 'ade20k_example.jpg'
gluoncv.utils.download(url, filename, True)

输出

Downloading ade20k_example.jpg from https://github.com/zhanghang1989/image-data/blob/master/encoding/segmentation/ade20k/ADE_val_00001142.jpg?raw=true...

  0%|          | 0/52 [00:00<?, ?KB/s]
100%|##########| 52/52 [00:00<00:00, 28186.07KB/s]

加载图像

img = image.imread(filename)

from matplotlib import pyplot as plt
plt.imshow(img.asnumpy())
plt.show()

使用数据集均值归一化图像

from gluoncv.data.transforms.presets.segmentation import test_transform
img = test_transform(img, ctx)

加载预训练模型并进行预测

获取预训练模型

model = gluoncv.model_zoo.get_model('psp_resnet101_ade', pretrained=True)

输出

self.crop_size 480
Downloading /root/.mxnet/models/psp_resnet101_ade-240a4758.zip from https://apache-mxnet.s3-accelerate.dualstack.amazonaws.com/gluon/models/psp_resnet101_ade-240a4758.zip...

  0%|          | 0/266773 [00:00<?, ?KB/s]
  0%|          | 1031/266773 [00:00<00:31, 8392.73KB/s]
  2%|2         | 5384/266773 [00:00<00:10, 23844.55KB/s]
  5%|4         | 12902/266773 [00:00<00:05, 44461.37KB/s]
  8%|7         | 20452/266773 [00:00<00:04, 55776.28KB/s]
 11%|#         | 28197/266773 [00:00<00:03, 63195.43KB/s]
 13%|#3        | 34845/266773 [00:00<00:03, 64272.16KB/s]
 16%|#6        | 42949/266773 [00:00<00:03, 69625.04KB/s]
 19%|#9        | 50747/266773 [00:00<00:03, 71495.26KB/s]
 22%|##2       | 59338/266773 [00:00<00:02, 75919.75KB/s]
 25%|##5       | 66987/266773 [00:01<00:02, 75831.04KB/s]
 28%|##8       | 75307/266773 [00:01<00:02, 77917.21KB/s]
 31%|###1      | 83128/266773 [00:01<00:02, 77256.28KB/s]
 34%|###4      | 91099/266773 [00:01<00:02, 77492.01KB/s]
 37%|###7      | 99477/266773 [00:01<00:02, 79362.91KB/s]
 40%|####      | 107426/266773 [00:01<00:02, 77197.04KB/s]
 44%|####3     | 116084/266773 [00:01<00:01, 79848.09KB/s]
 47%|####6     | 124090/266773 [00:01<00:01, 78043.56KB/s]
 50%|####9     | 132070/266773 [00:01<00:01, 78554.55KB/s]
 53%|#####2    | 140081/266773 [00:01<00:01, 79010.15KB/s]
 55%|#####5    | 147995/266773 [00:02<00:01, 78087.21KB/s]
 59%|#####8    | 156270/266773 [00:02<00:01, 79458.08KB/s]
 62%|######1   | 164226/266773 [00:02<00:01, 77305.85KB/s]
 65%|######4   | 172466/266773 [00:02<00:01, 78789.35KB/s]
 68%|######7   | 180362/266773 [00:02<00:01, 76982.85KB/s]
 71%|#######   | 188577/266773 [00:02<00:01, 78172.79KB/s]
 74%|#######3  | 196629/266773 [00:02<00:00, 78856.90KB/s]
 77%|#######6  | 204527/266773 [00:02<00:00, 78590.58KB/s]
 80%|#######9  | 212769/266773 [00:02<00:00, 79715.43KB/s]
 83%|########2 | 220748/266773 [00:03<00:00, 78518.69KB/s]
 86%|########5 | 228788/266773 [00:03<00:00, 79071.20KB/s]
 89%|########8 | 236915/266773 [00:03<00:00, 79711.55KB/s]
 92%|#########1| 244892/266773 [00:03<00:00, 79623.70KB/s]
 95%|#########4| 252859/266773 [00:03<00:00, 77590.81KB/s]
 98%|#########7| 260850/266773 [00:03<00:00, 77826.00KB/s]
100%|##########| 266773/266773 [00:03<00:00, 74201.69KB/s]

使用单一尺度进行预测

output = model.predict(img)
predict = mx.nd.squeeze(mx.nd.argmax(output, 1)).asnumpy()

添加调色板用于可视化

from gluoncv.utils.viz import get_color_pallete
import matplotlib.image as mpimg
mask = get_color_pallete(predict, 'ade20k')
mask.save('output.png')

显示预测的掩码

mmask = mpimg.imread('output.png')
plt.imshow(mmask)
plt.show()