Awesome

Robust Trust Region for Weakly Supervised Segmentation

This is an implementation for the paper:

• Dmitrii Marin and Yuri Boykov. Robust Trust Region for Weakly Supervised Segmentation. In International Conference on Computer Vision (ICCV), 2021

[paper page] [arxiv] [poster] [video]

If you find this code useful in your research, consider citing (bibtex):

@InProceedings{Marin_2021_ICCV,
    author    = {Marin, Dmitrii and Boykov, Yuri},
    title     = {Robust Trust Region for Weakly Supervised Segmentation},
    booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
    month     = {October},
    year      = {2021},
    pages     = {6608-6618}
}

Prerequisites

Environment

The code uses the following python packages: pytorch, tqdm, scipy, tensorboardX.

ScribbleSup dataset

Download original PASCAL VOC 2012 dataset: http://host.robots.ox.ac.uk/pascal/VOC/.

Download Scribble annotations: http://cs.uwaterloo.ca/~m62tang/rloss/pascal_2012_scribble.zip, the original dataset can be found here.

Building and installing python extensions

Alpha-expansion code (GCO) for GridCRF regularized loss with Robust Trust Region:

cd wrapper/alpha_expansion
swig -python -c++ alpha_expansion.i
python setup.py install

Bilateral filtering for the DenseCRF regularized loss:

cd wrapper/bilateralfilter
swig -python -c++ bilateralfilter.i
python setup.py install

Loading pre-trained models

All TR models require pretraining with partial cross-entropy. The models can be downloaded using the links:

Training with Robust Trust Region

Run the following script:

bash full_tr_error_const.sh 

The script is controlled by the following environmental variables:

• suffix — one of empty string, 0.8, 0.5, 0.3, 0. Determines the level of supervision, see pre-trained models description.
• ERROR_PROB — the robustness parameter ε. Must be between 0 and 1. See paper for details.
• LR — the base learning rate (default 0.0007).
• TR_WEIGHT — the weight of the unary term in alpha-expansion optimization. See stage-A in the paper.
• HIDDEN_UPDATE — the number of iterations of alpha expansion (defualt 5).

Parameter configuration for different levels of supervision

For example, the following will train a model using fulllength scribbles:

suffix= ERROR_PROB=0.94400 TR_WEIGHT=0.05 bash full_tr_error_const.sh