2. Test with PSPNet Pre-trained Models

This is a quick demo of using GluonCV PSPNet model on ADE20K dataset. Please follow the installation guide 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

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)

Out:

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]

load the image

img = image.imread(filename)

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

normalize the image using dataset mean

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

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

Out:

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]

make prediction using single scale

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

Add color pallete for visualization

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

show the predicted mask

demo psp

Total running time of the script: ( 0 minutes 8.428 seconds)

Gallery generated by Sphinx-Gallery