The following is adapted from https://gluon-cv.mxnet.io/build/examples_segmentation/demo_psp.html#test-with-pspnet-pre-trained-models
In [1]:
!pip install mxnet
In [3]:
!pip install gluoncv
In [0]:
import mxnet as mx
from mxnet import image
from mxnet.gluon.data.vision import transforms
import gluoncv
# using cpu
ctx = mx.cpu(0)
In [5]:
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)
Out[5]:
In [6]:
img = image.imread(filename)
from matplotlib import pyplot as plt
plt.imshow(img.asnumpy())
plt.show()
In [0]:
transform_fn = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([.485, .456, .406], [.229, .224, .225])
])
img = transform_fn(img)
img = img.expand_dims(0).as_in_context(ctx)
In [8]:
model = gluoncv.model_zoo.get_model('psp_resnet50_ade', pretrained=True)
In [0]:
output = model.demo(img)
predict = mx.nd.squeeze(mx.nd.argmax(output, 1)).asnumpy()
In [0]:
from gluoncv.utils.viz import get_color_pallete
import matplotlib.image as mpimg
mask = get_color_pallete(predict, 'ade20k')
mask.save('output.png')
In [11]:
mmask = mpimg.imread('output.png')
plt.imshow(mmask)
plt.show()
In [0]: