PSPNet Tensorflow 2

Keras Pyramid Scene Parsing Network ported to tensorflow 2 from keras/tf_1.13.

• Caffe implementation: PSPNet
• Py35 Keras Tensorflow1.13 implementation: PSPNet-Keras-tensorflow

Dependencies

• Tensorflow 2 (tensorflow / tensorflow-gpu / tensorflow-cpu)
• OpenCV (opencv-python / opencv-contrib-python)

Pretrained weights

Pretrained weights can be downloaded here:

H5 Weights	Architecture	Numpy Weights
pspnet50_ade20k.h5	pspnet50_ade20k.json	pspnet50_ade20k.npy
pspnet101_cityscapes.h5	pspnet101_cityscapes.json	pspnet101_cityscapes.npy
pspnet101_voc2012.h5	pspnet101_voc2012.json	pspnet101_voc2012.npy

Download weights in

• .h5 and .json format and place them at weights/keras or
• .npy and place them at weights/npy

Find example notebook which demonstrates save and load.

Usage:

# python pspnet.py -m <model> -i <input_image>  -o <output_path> [-other_arguments]
python pspnet.py -m pspnet101_cityscapes -i example_images/cityscapes.jpg -o example_results/cityscapes.jpg -s -ms -f
python pspnet.py -m pspnet101_voc2012 -i example_images/pascal_voc.jpg -o example_results/pascal_voc.jpg -s -ms -f
python pspnet.py -m pspnet50_ade20k -i example_images/ade20k.jpg -o example_results/ade20k.jpg -s -ms -f

List of arguments:

 -m --model        - which model to use: 'pspnet50_ade20k', 'pspnet101_cityscapes', 'pspnet101_voc2012'
    --id           - (int) GPU Device id. Default 0
 -s --sliding      - Use sliding window
 -f --flip         - Additional prediction of flipped image
 -ms --multi_scale - Predict on multiscale images

Batch Predict on GPU, check source here

Keras results:

Input	Segmented	Blended	Probe

Implementation

• The interpolation layer is implemented as custom layer "Interp"
• Forward step takes about ~1 sec on single image
• Memory usage can be optimized with:

# before calling any of the tf functions
for gpu in tf.config.experimental.list_physical_devices('GPU'):
    tf.config.experimental.set_memory_growth(gpu, True)
    # if you want to restrict total memory you can try
    # tf.config.experimental.set_memory_growth(gpu, True)