Bounding Box Tightness Prior for Weakly Supervised Image Segmentation

This project hosts the codes for the implementation of the paper Bounding Box Tightness Prior for Weakly Supervised Image Segmentation (MICCAI 2021) [miccai] [arxiv].

Dataset preprocessing

Download Promise12 dataset, and put it on the "data/prostate" folder.

Download Atlas dataset, and put it on the "data/atlas" folder.

Run the following codes for preprocessing:

# trainig and valid subsets for promise12 dataset
python preprocess/slice_promise_train_val.py
python preprocess/slice_promise_augment_train_val.py

# trainig and valid subsets for atlas dataset
python preprocess/slice_atlas.py

Training

#  The following experiments include full supervision (exp_no=0), MIL ablation study (exp_no=1), smooth maximum approximation ablation study (exp_no=2,3), and main experiments (exp_no=4,5)

# training for promise12 dataset, exp_no=0,1,2,3,4,5
CUDA_VISIBLE_DEVICES=0 python tools/train_promise_unetwithbox.py --n_exp exp_no
# training for atlas dataset, exp_no=0,1,2,3,4,5
CUDA_VISIBLE_DEVICES=0 python tools/train_atlas_unetwithbox.py --n_exp exp_no

Validation

# Dice validation results for promise12 dataset, exp_no=0,1,2,...,16
CUDA_VISIBLE_DEVICES=0 python tools/valid_promise_unetwithbox.py --n_exp exp_no
# Dice validation results for atlas dataset, exp_no=0,1,2,...,16
CUDA_VISIBLE_DEVICES=0 python tools/valid_atlas_unetwithbox.py --n_exp exp_no

Performance summary

python tools/report_promise_unetwithbox_paper.py
python tools/report_atlas_unetwithbox_paper.py

Citations

Please consider citing our paper in your publications if the project helps your research.

@inproceedings{wang2021bounding,
  title={Bounding Box Tightness Prior for Weakly Supervised Image Segmentation},
  author={Wang, Juan and Xia, Bin},
  booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
  pages={526--536},
  year={2021},
  organization={Springer}
}

Logs

1. 3/16/2022: a bug in _C_promise.py was fixed such that the training subset was used for training.
2. 3/17/2022: add data augmentation script for promise12 dataset