Home

Awesome

AuxSegNet

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

<p align="left"> <img src="mis/framework2.jpg" width="720" title="" > </p> <p align="left"> <img src="mis/example.jpg" width="720" title="" > </p>

Prerequisite

pip install -r prerequisite.txt

Usage

1. Prepare initial pseudo labels

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

<p align="left"> <img src="mis/sota_voc.png" width="300" title="" > </p>

Segmentation results on the MS COCO dataset

<p align="left"> <img src="mis/sota_coco.png" width="300" title="" > </p>

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}
}