Awesome
Cut-Paste Consistency Learning
This is the official code repository for the WACV 2023 paper "Cut-Paste Consistency Learning for Semi-Supervised Lesion Segmentation".
Installation
python -m virtualenv -p 3.6 env
source env/bin/activate
pip install -r requirements.txt
python setup.py install
Downloads
Datasets
Pretrained weights (PyTorch)
Example model checkpoints for the lesion segmentation tasks in IDRiD are provided:
Model | Method | AUC-PR (%) |
---|---|---|
IDRiD-MA | Cut-Paste Consistency + Mean Teacher | 51.33 |
IDRiD-HE | Cut-Paste Consistency + Mean Teacher | 66.86 |
IDRiD-EX | Cut-Paste Consistency + Mean Teacher | 88.70 |
IDRiD-SE | Cut-Paste Consistency + Mean Teacher | 79.53 |
Training
List of supported datasets and learning methods:
data_module | model | Description |
---|---|---|
idrid , ich | unet | Supervised baseline |
idrid-base-cp , ich-base-cp | unet-pseudo | Cut-paste baseline |
idrid-st , ich-st | unet-pseudo | Self-training |
idrid-st-cp , ich-st-cp | unet-pseudo | Self-training + cut-paste |
idrid-semi , ich-semi | unet-mt | Mean Teacher |
idrid-semi , ich-semi | unet-classmix | ClassMix consistency |
idrid-semi , ich-semi | unet-cutmix | CutMix consistency |
idrid-cp , ich-cp | unet-cp | Cut-paste consistency |
Type python main.py <data_module> <model> --help
in the console for more details.
Example of cut-paste consistency learning on IDRiD-MA:
python main.py \
idrid-cp \
unet-cp \
--unlabeled_weight 0.01 \
--mean_teacher \
--base_ema 0.996 \
--seed 42 \
--num_workers 5 \
--batch_size 5 \
--synth_split 0.4 \
--num_synth 300 \
--mask_blur gaussian \
--background_blur gaussian \
--img_match \
--val_split 0.1 \
--data_dir data/IDRiD \
--gpus [0] \
--max_epochs 500 \
--check_val_every_n_epoch 1 \
--early_stopping_patience -1 \
--log_every_n_steps 10 \
--learning_rate 6e-4 \
--warmup_epochs 10 \
--optimizer adamw \
--weight_decay 1e-5 \
--lr_scheduler cosine \
--num_layers 5 \
--features_start 64 \
--preprocess resize \
--size 512 \
--inference_mode resize \
--inference_size 512 \
--checkpoint_monitor "val/aupr" \
--do_train \
--num_sanity_val_steps 0 \
--pos_weight 6.84 \
--default_root_dir "model/idrid-MA" \
--task_id MA
Example of cut-paste consistency learning on CT-ICH:
python main.py \
ich-cp \
unet-cp \
--unlabeled_weight 0.1 \
--mean_teacher \
--base_ema 0.996 \
--seed 42 \
--num_workers 5 \
--batch_size 8 \
--labeled_split 0.7 \
--synth_split 0.4 \
--img_match \
--mask_blur gaussian \
--background_blur none \
--data_dir "data/CT-ICH/data/fold-1" \
--default_root_dir "model/ich" \
--gpus [0] \
--max_epochs 50 \
--check_val_every_n_epoch -1 \
--early_stopping_patience -1 \
--log_every_n_steps 10 \
--learning_rate 3e-5 \
--warmup_epochs 10 \
--optimizer adamw \
--weight_decay 1e-5 \
--lr_scheduler cosine \
--num_layers 5 \
--features_start 64 \
--input_channels 1 \
--preprocess resize \
--size 512 \
--inference_mode resize \
--inference_size 512 \
--do_train \
--do_test \
--disable_aupr \
--num_sanity_val_steps 0 \
--pos_weight 7.08
Evaluation
Example of evaluating a trained model on IDRiD-MA:
python main.py \
idrid \
unet \
--num_workers 1 \
--data_dir "data/IDRiD" \
--num_layers 5 \
--features_start 64 \
--inference_mode resize \
--inference_size 512 \
--do_test \
--aupr_in_cpu \
--batch_size 1 \
--gpus 1 \
--default_root_dir "model/test/IDRiD-MA" \
--task_id MA \
--resume_from_checkpoint "model/IDRiD-MA/checkpoint.ckpt"
Citation
@InProceedings{Yap_2023_WACV,
author = {Yap, Boon Peng and Ng, Beng Koon},
title = {Cut-Paste Consistency Learning for Semi-Supervised Lesion Segmentation},
booktitle = {Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV)},
month = {January},
year = {2023},
pages = {6160-6169}
}