Awesome
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 (fuy34bkup@gmail.com) 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.