xulianuwa / AuxSegNet

AuxSegNet

The pytorch code for our ICCV 2021 paper Leveraging Auxiliary Tasks with Affinity Learning for Weakly Supervised Semantic Segmentation.

Prerequisite

• Ubuntu 18.04, with Python 3.6 and the following python dependencies.

pip install -r prerequisite.txt

• Download the PASCAL VOC 2012 development kit.

Usage

1. Prepare initial pseudo labels

• Off-the-shelf saliency maps used as the initial saliency pseudo labels. [DSS]
• Extract the class activation maps (CAM) from a pre-trained single-task classification network. [ResNet38]
• Generate the initial pseudo segmentation labels using the above saliency and CAM maps via [heuristic fusion].

2. Train the AuxSegNet

python train_AuxAff.py --img_path 'Path to the training images'\
                       --seg_pgt_path 'Path to the pseudo segmentation labels' \
                       --sal_pgt_path 'Path to the pseudo saliency labels' \
                       --init_weights 'Path to the initialization weights' \
                       --save_path 'Path to save the trained AuxSegNet model' 

3. Pseudo label updating

python gen_pgt.py --weights 'path to the trained AuxSegNet weights'\   
                  --img_path 'Path to the training images'\
                  --SALpath 'Path to the pre-trained saliency maps' \
                  --seg_pgt_path 'Path to save updated pseudo segmentation labels' \
                  --sal_pgt_path 'Path to save updated pseudo saliency labels' 

4. Iterate Step 2 and 3

(Optional) Integrated iterative model learning and label updating (Step 2-4)

bash iter_learn.sh

5. Inference

python infer_AuxAff.py --img_path 'Path to the training images'\
                       --weights 'Path to the trained AuxSegNet weights'\
                       --save_path 'Path to save the segmentation results'

Performance comparison with SOTA

Segmentation results on the PASCAL VOC 2012 dataset

Segmentation results on the MS COCO dataset

Citation

Please consider citing our paper if the code is helpful in your research and development.

@inproceedings{xu2021leveraging,
  title={Leveraging Auxiliary Tasks with Affinity Learning for Weakly Supervised Semantic Segmentation},
  author={Xu, Lian and Ouyang, Wanli and Bennamoun, Mohammed and Boussaid, Farid and Sohel, Ferdous and Xu, Dan},
  booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
  year={2021}
}