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