kazuto1011 / pspnet-pytorch

PyTorch re-implementation of PSPNet
MIT License
52 stars 11 forks source link
pspnet pytorch semantic-segmentation

PSPNet with PyTorch

Unofficial implementation of "Pyramid Scene Parsing Network" (https://arxiv.org/abs/1612.01105). This repository is just for caffe to pytorch model conversion and evaluation.

Requirements

Preparation

Instead of building the author's caffe implementation, you can convert off-the-shelf caffemodels to pytorch models via the caffe.proto.

1. Compile the caffe.proto for Python API

This step can be skipped. FYI.
Download the author's caffe.proto into the libs, not the one in the original caffe.

# For protoc command
pip install protobuf
# This generates ./caffe_pb2.py
protoc --python_out=. caffe.proto

2. Model conversion

  1. Find the caffemodels on the author's page (e.g. pspnet50_ADE20K.caffemodel) and store them to the data/models/ directory.
  2. Convert the caffemodels to .pth file.
python convert.py -c <PATH TO YAML>

Demo

python demo.py -c <PATH TO YAML> -i <PATH TO IMAGE>

demo

Evaluation

PASCAL VOC2012 only. Please set the dataset path in config/voc12.yaml.

python eval.py -c config/voc12.yaml

88.1% mIoU (SS) and 88.6% mIoU (MS) on validation set.
NOTE: 3 points lower than caffe implementation. WIP

References