3. 使用 DeepLabV3 预训练模型进行测试

这是关于在 ADE20K 数据集上使用 GluonCV DeepLabV3 模型的快速演示。如果您尚未安装 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_00001755.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_00001755.jpg?raw=true...

  0%|          | 0/13 [00:00<?, ?KB/s]
14KB [00:00, 8474.56KB/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('deeplab_resnet101_ade', pretrained=True)

输出

Downloading /root/.mxnet/models/deeplab_resnet101_ade-bf1584df.zip from https://apache-mxnet.s3-accelerate.dualstack.amazonaws.com/gluon/models/deeplab_resnet101_ade-bf1584df.zip...

  0%|          | 0/221969 [00:00<?, ?KB/s]
  0%|          | 101/221969 [00:00<04:36, 803.10KB/s]
  0%|          | 520/221969 [00:00<01:37, 2265.47KB/s]
  1%|          | 2182/221969 [00:00<00:30, 7234.27KB/s]
  3%|3         | 7550/221969 [00:00<00:09, 23029.09KB/s]
  7%|6         | 14610/221969 [00:00<00:05, 38583.15KB/s]
 10%|#         | 22401/221969 [00:00<00:03, 51093.04KB/s]
 14%|#3        | 30196/221969 [00:00<00:03, 59484.55KB/s]
 17%|#7        | 38051/221969 [00:00<00:02, 65368.22KB/s]
 21%|##        | 45798/221969 [00:00<00:02, 69068.17KB/s]
 24%|##4       | 53665/221969 [00:01<00:02, 71986.34KB/s]
 28%|##7       | 61377/221969 [00:01<00:02, 73348.15KB/s]
 31%|###1      | 69673/221969 [00:01<00:01, 76241.63KB/s]
 35%|###4      | 77347/221969 [00:01<00:01, 75721.66KB/s]
 39%|###8      | 85500/221969 [00:01<00:01, 77453.06KB/s]
 42%|####2     | 93272/221969 [00:01<00:01, 77342.55KB/s]
 46%|####5     | 101383/221969 [00:01<00:01, 78466.44KB/s]
 49%|####9     | 109244/221969 [00:01<00:01, 78259.35KB/s]
 53%|#####2    | 117167/221969 [00:01<00:01, 78548.28KB/s]
 56%|#####6    | 125174/221969 [00:01<00:01, 79002.29KB/s]
 60%|#####9    | 133080/221969 [00:02<00:01, 78589.53KB/s]
 64%|######3   | 141116/221969 [00:02<00:01, 78995.23KB/s]
 67%|######7   | 149019/221969 [00:02<00:00, 78701.69KB/s]
 71%|#######   | 156892/221969 [00:02<00:00, 78601.97KB/s]
 74%|#######4  | 164934/221969 [00:02<00:00, 79142.10KB/s]
 78%|#######7  | 172850/221969 [00:02<00:00, 79039.16KB/s]
 82%|########1 | 180953/221969 [00:02<00:00, 79632.97KB/s]
 85%|########5 | 188918/221969 [00:02<00:00, 62789.97KB/s]
 88%|########8 | 196298/221969 [00:02<00:00, 65563.78KB/s]
 92%|#########1| 203297/221969 [00:03<00:00, 66521.33KB/s]
 95%|#########4| 210326/221969 [00:03<00:00, 67551.98KB/s]
 98%|#########7| 217316/221969 [00:03<00:00, 67283.14KB/s]
100%|##########| 221969/221969 [00:03<00:00, 65635.86KB/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()