Skip to content

Latest commit

 

History

History
115 lines (94 loc) · 4.9 KB

README.md

File metadata and controls

115 lines (94 loc) · 4.9 KB

U-Noise - PyTorch Implementation

U-Noise: Learnable Noise Masks for Interpretable Image Segmentation
Teddy Koker, Fatemehsadat Mireshghallah, Tom Titcombe, Georgios Kaissis

Image + Mask U-Noise Large Occlusion Sensitivity Grad-CAM

Download Data/Pre-trained Models

The dataset can be created by downloading and un-taring Task07_Pancreas.tar from the Medical Segmentation Decathlon into the data/ directory. Once there, run the prepare_data.py script within the directory.

./download.sh

This will download the following files:

models
├── unoise_small.ckpt   # small U-Noise model
├── unoise_medium.ckpt  # medium U-Noise model
├── unoise_large.ckpt   # large U-Noise model TODO
├── unoise_small_pretrained.ckpt  # small U-Noise model pretrained TODO 
├── unoise_medium_pretrained.ckpt # medium U-Noise model pretrained TODO
├── unoise_large_pretrained.ckpt  # large U-Noise model pretrained TODO
└── utility.ckpt        # pretrained utility model

Alternatively, each model can be downloaded individually:

Note: each U-Noise model contains a copy of the utility model

Model # parameters weights
Utility 34M utility.ckpt
U-Noise Small 10K unoise_small.ckpt
U-Noise Medium 130K unoise_medium.ckpt
U-Noise Large 537K unoise_large.ckpt
U-Noise Small (pretrained) 10K unoise_small_pretrained.ckpt
U-Noise Medium (pretrained) 130K unoise_medium_pretrained.ckpt
U-Noise Large (pretrained) 537K unoise_large_pretrained.ckpt
U-Nets Params Depth Channels
Utility 34M 5 1, 64, 128, 256, 512, 1024, 512, ...
Small 28K 2 1, 16, 32, 16, 1
Medium 130K 3 1, 16, 32, 64, 32, 16, 1
Medium 537K 4 1, 16, 32, 64, 128, 64, 32, 16, 1

Reproducing Results

Train Utility model (~5 hours on NVIDIA 2070 Super):

python src/train_util.py

Train U-Noise Small:

python src/train_noise.py --depth 2 --channel_factor 4 --batch_size 8

Train U-Noise Medium:

python src/train_noise.py --depth 3 --channel_factor 4 --batch_size 8

Train U-Noise Large:

python src/train_noise.py --depth 4 --channel_factor 4 --batch_size 8

Train U-Noise Small (Pretrained):

python src/train_util.py --depth 2 --channel_factor 4 --batch_size 8
python src/train_noise.py --depth 2 --channel_factor 4 --batch_size 8 \
 --pretrained /path/to/pretrained --learning_rate 1e-3

Train U-Noise Medium (Pretrained):

python src/train_util.py --depth 3 --channel_factor 4 --batch_size 8
python src/train_noise.py --depth 3 --channel_factor 4 --batch_size 8 \
 --pretrained /path/to/pretrained --learning_rate 1e-3

Train U-Noise Large (Pretrained):

python src/train_util.py --depth 4 --channel_factor 4 --batch_size 8
python src/train_noise.py --depth 4 --channel_factor 4 --batch_size 8 \
 --pretrained /path/to/pretrained --learning_rate 1e-3

Citation

If you found this work helpful, please cite:

@misc{koker2021unoise,
      title={U-Noise: Learnable Noise Masks for Interpretable Image Segmentation}, 
      author={Teddy Koker and Fatemehsadat Mireshghallah and Tom Titcombe and Georgios Kaissis},
      year={2021},
      eprint={2101.05791},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}