SPPSMNet: Spixel Pyramid Stereo Matching Network

This is a PyTorch implementation of the SPPSMNet proposed in our CVPR-20 paper:

Superpixel Segmentation with Fully Convolutional Network

Fengting Yang, Qian Sun, Hailin Jin, and Zihan Zhou

The code for the superpixel segmentation method (SpixelFCN) is available in this repository. Please contact Fengting Yang ([email protected]) if you have any questions.

Prerequisites

The training code was mainly developed and tested with Python 2.7, PyTorch 0.4.1, CUDA 9, and Ubuntu 16.04.

We also made a test with Python 3.6, Pytorch 1.6, and CUDA 10.2. The result drops slightly, but it is still close to what we reported on the paper.

Prior to running this code, we need a pre-trained SpixelFCN model to initialize the SpixelFCN part. Two pre-trained model with grid size 4 and 16 have been provided in preTrain_spixel_model folder. Please refer to SpixelFCN to pre-train a model with different grid size.

Data preparation

We train/fine-tune our model on three public dataset: SceneFlow, HRVS, and Middlebury-v2.

• For Sceneflow, we follow the official training and test set.
• For HRVS, we have listed the training and validation split in dataloader/HRVS_*.txt.
• For Middlebury-v2, we fine-tuned on both the training and the additional dataset for our submission.

Please follow the offical instruction of each dataset to download them.

Training

We first train the model on SceneFlow dataset. For joint training, run

python main.py --datapath <SCENEFLOW_PATH> --savemodel <SF_DUMP_PATH> --m_w 30

For SpixelFCN fixed training, run

python main_fixed.py --datapath <SCENEFLOW_PATH> --savemodel <SF_FIX_DUMP_PATH> 

The superpixel grid size is equal to 4 for the SceneFlow training, and the corresponding SpixelFCN checkpoint will be loaded automatically from preTrain_spixel_model.

Fine-tuning

Next, we can fine-tune the pre-trained SceneFlow model on HR_VS, by running

python finetune_HRVS_sp16.py --datapath <HRVS_PATH> --loadmodel <SF_CKPT> --m_w 30

Similarly, we also provide the code to fine-tune a fixed model in finetune_HRVS_sp16_fixed.py.

To fine-tune on Middlebury-v2, run

python finetune_mb_sp16.py --dataset <MB_PATH> --loadmodel <SF_CKPT> --m_w 60

The superpixel grid size is equal to 16 for these two datasets. The corresponding SpixelFCN checkpoint will be automatically loaded and overwrite the weight in <SF_CKPT>.

Test

We have provided our pre-trained SceneFlow model and HR_VS model online. To test on SceneFlow, run

python infer_sceneflow.py --dataset <SCENEFLOW_PATH> --loadmodel <SF_CKPT>

To test on HRVS, run

python infer_sceneflow.py --dataset <HRVS_PATH> --loadmodel <HRVS_CKPT>

The EPE score will be shown after the test.

We can also use submission_mb.py to create the submission package for the Middlebury-v2 evaluation.

Acknowledgement

This code is developed based on the PSMNet and SpixelFCN.