CateNorm

Paper

This repository provides the official implementation of training CateNorm in the following paper:

CateNorm: Categorical Normalization for Robust Medical Image Segmentation
Junfei Xiao¹, Lequan Yu², Zongwei Zhou¹, Yutong Bai¹,
Lei Xing³, Alan Yuille¹, Yuyin Zhou^4
1 Johns Hopkins University, ² The University of Hong Kong,
³ Stanford University, ⁴ UC Santa Cruz
MICCAI Workshop on Domain Adaptation and Representation Transfer (DART), 2022
${\color{red} {\textbf{Best Paper Award Honourable Mention}}}$
paper | code | slides

Install/Check dependencies:

 pip install requirements.txt

Prepare Dataset

1. Download dataset

1. Prostate: We use the preprocessed multi-site dataset for prostate MRI segmentation.
2. Abdominal: We use the BTCV and TCIA datasets. For single domain(site) experiments, we directly use BTCV with the official annotations. For multiple domain(site) experiments, we use the annotation published in here.

2. Preprocess data

1. Prostate: Datasets have already been preprocessed.
2. Abdominal: There are two jupyter notebooks in ./preprocess for preprocessing data with different settings.

3. Generate text file for each site

Each site-wise folder needs a text file(all_list.txt) including paths of all cases. data_list_generator.ipynb is offered to help you generate.

4. Overview of a dataset folder

A dataset folder should look like this:

Dataset/Prostate_Multi/
├── Site-A
│   ├── all_list.txt
│   ├── Case00.nii.gz
│   ├── Case00_segmentation.nii.gz
│   ├── Case01.nii.gz
│   ├── Case01_segmentation.nii.gz
│   ...
│
├── Site-B
│   ├── all_list.txt
│   ├── Case00.nii.gz
│   ├── Case00_segmentation.nii.gz
│   ├── Case01.nii.gz
│   ├── Case01_segmentation.nii.gz
│   ...
│
├── Site-C
│   ├── all_list.txt
│   ├── Case00.nii.gz
│   ├── Case00_segmentation.nii.gz
│   ├── Case01.nii.gz
│   ├── Case01_segmentation.nii.gz
│   ...
│

5. Set up the data and result paths

Please modify "data_dir" and "save_dir" in train.py & test.py with your own configuration.

 data_dir = {'local-prostate': 'G:/Dataset/Prostate_Multi_Site',
             'local-ABD-8': 'G:/Dataset/Abdominal_Single_Site_8organs',
             'local-ABD-6': 'G:/Dataset/Abdominal_Multi_Site_6organs',
             }

 save_dir = {'local-prostate': 'G:/DualNorm-Unet/',
             'local-ABD-8': 'G:/DualNorm-Unet/',
             'local-ABD-6': 'G:/DualNorm-Unet/',
             }

Training

1. Baseline (single site):

   # Here we use prostate site A as an example
   
   python train.py  --save-fold=Prostate-Single-A --batch-size=4 --sitename A --epochs=20

2. Baseline (Multiple site):

   # Here we use prostate sites A,B,C as an example
   
   python train.py  --save-fold=Prostate-Multi-ABC --batch-size=6 --sitename ABC --epochs=20

3. Load Pretraining Model

   # Here we use prostate sites A,B,C as an example
   
   python train.py  --save-fold=Prostate-Multi-ABC-Pretrained --batch-size=6 --sitename ABC --epochs=10 \
   --load=G:/DualNorm-Unet/checkpoints/xxx/xxx/Epochs_10_Aug_True_Zoom_False_Nonlinear_relu_Norm_BN

4. DualNorm-Unet :

   # Here we use prostate sites A,B,C with DualNorm blocks 1-4(inc, down1, down2, down3) as an example
   
   python train.py  --save-fold=Prostate-Multi-ABC-Pretrained --batch-size=6 --sitename ABC --epochs=10 \
   --load=G:/DualNorm-Unet/checkpoints/xxx/xxx/Epochs_10_Aug_True_Zoom_False_Nonlinear_relu_Norm_BN \
   --spade-aux-blocks inc down1 down2 down3

Testing

To evaluate the model and save predictions, run:

python test.py  --save-fold=Prostate-Multi-ABC-Test --batch-size=6 --sitename ABC \
--load=G:/DualNorm-Unet/checkpoints/xxx/xxx/Epochs_10_Aug_True_Zoom_False_Nonlinear_relu_Norm_BN \
--save-prediction=True

All the predictions are saved as .nii files in the prediction_nii folder e.g., G:/DualNorm-Unet/prediction_nii.

Reference:

• https://github.com/milesial/Pytorch-UNet
• https://github.com/liuquande/MS-Net

Citations

@inproceedings{xiao2022catenorm,
  title={CateNorm: Categorical Normalization for Robust Medical Image Segmentation},
  author={Xiao, Junfei and Yu, Lequan and Zhou, Zongwei and Bai, Yutong and Xing, Lei and Yuille, Alan and Zhou, Yuyin},
  booktitle={MICCAI Workshop on Domain Adaptation and Representation Transfer},
  pages={129--146},
  year={2022},
  organization={Springer}
}