/CateNorm

The official implementation of "CateNorm: Categorical Normalization for Robust Medical Image Segmentation"

Primary LanguagePython

CateNorm

Paper

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

CateNorm: Categorical Normalization for Robust Medical Image Segmentation
Junfei Xiao1, Lequan Yu2, Zongwei Zhou1, Yutong Bai1,
Lei Xing3, Alan Yuille1, Yuyin Zhou4
1 Johns Hopkins University, 2 The University of Hong Kong,
3 Stanford University, 4 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:

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