Home

Awesome

[NeurIPS-2023] Towards Generic Semi-Supervised Framework for Volumetric Medical Image Segmentation

This repo is the official implementation of Towards Generic Semi-Supervised Framework for Volumetric Medical Image Segmentation which is accepted at NeurIPS-2023.

πŸš€ The significance of this work lies in its ability to encourage semi-supervised medical image segmentation methods to address more complex real-world application scenarios, rather than just developing frameworks in ideal experimental environments. Furthermore, we have consolidated all four settings within this single codebase, enabling the execution of any task using a single bash file by merely adjusting the arguments.

<p align = "center"> <img src="./images/analysis.png" width="70%"/> </p> <p align = "center"> <img src="./images/Framework.png" width="80%"/> </p>

Online Presentation Video is available for brief introduction.

1. Environment

First, create a new environment and install the requirements:

conda create -n genericssl python=3.8
conda activate genericssl
cd GenericSSL/
pip install -r requirements.txt

[πŸ“ŒIMPORTANT] Then, before running the code, set the PYTHONPATH to pwd:

export PYTHONPATH=$(pwd)/code:$PYTHONPATH

2. Data Preparation

First, download the datasets and put them under the Datasets folder:

The file structure should be:

.
β”œβ”€β”€ Datasets
β”‚Β Β  β”œβ”€β”€ LASeg
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ 2018LA_Seg_Training Set
β”‚Β Β  β”‚Β Β  β”‚Β Β  β”œβ”€β”€ 0RZDK210BSMWAA6467LU
β”‚Β Β  β”‚Β Β  β”‚Β Β  β”‚Β Β  β”œβ”€β”€ mri_norm2.h5
β”‚Β Β  β”‚Β Β  β”‚Β Β  β”œβ”€β”€ 1D7CUD1955YZPGK8XHJX
β”‚Β Β  β”‚Β Β  β”‚Β Β  └── ...
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ test.list
β”‚Β Β  β”‚Β Β  └── train.list
β”‚Β Β  β”œβ”€β”€ MMWHS
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ CT
β”‚Β Β  β”‚Β Β  β”‚Β Β  β”œβ”€β”€ imagesTr
β”‚Β Β  β”‚Β Β  β”‚Β Β  β”‚Β Β  β”œβ”€β”€ ct_train_1001_image.nii.gz
β”‚Β Β  β”‚Β Β  β”‚Β Β  β”‚Β Β  └── ...
β”‚Β Β  β”‚Β Β  β”‚Β Β  └── labelsTr
β”‚Β Β  β”‚Β Β  β”‚Β Β  β”‚Β Β  β”œβ”€β”€ ct_train_1001_label.nii.gz
β”‚Β Β  β”‚Β Β  β”‚Β Β  β”‚Β Β  └── ...
β”‚Β Β  β”‚Β Β  └── MR
β”‚Β Β  β”‚Β Β      β”œβ”€β”€ imagesTr
β”‚Β Β  β”‚Β Β      └── labelsTr
β”‚Β Β  β”œβ”€β”€ MNMs
β”‚Β Β  β”‚Β Β  └── Labeled
β”‚Β Β  β”‚Β Β      β”œβ”€β”€ VendorA
β”‚Β Β  β”‚Β Β   Β Β  β”‚Β Β  β”œβ”€β”€ A0S9V9
β”‚Β Β  β”‚Β Β   Β Β  β”‚Β Β  β”‚Β Β  β”œβ”€β”€ A0S9V9_sa.nii.gz
β”‚Β Β  β”‚Β Β   Β Β  β”‚Β Β  β”‚Β Β  β”œβ”€β”€ A0S9V9_sa_gt.nii.gz
β”‚Β Β  β”‚Β Β   Β Β  β”‚Β Β  β”œβ”€β”€ A1D9Z7
β”‚Β Β  β”‚Β Β   Β Β  β”‚Β Β  └── ...
β”‚Β Β  β”‚Β Β      β”œβ”€β”€ VendorB
β”‚Β Β  β”‚Β Β      β”œβ”€β”€ VendorC
β”‚Β Β  β”‚Β Β      └── VendorD
β”‚Β Β  β”œβ”€β”€ OpenDataset
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ Testing
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ Training
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ Validation
β”‚Β Β  β”‚Β Β  └── mnms_dataset_info.xls
β”‚Β Β  └── Synapse
β”‚Β Β      β”œβ”€β”€ imagesTr
β”‚Β Β   Β Β  β”‚Β Β  β”œβ”€β”€img0001.nii.gz
β”‚Β Β   Β Β  β”‚Β Β  └── ...
β”‚Β Β      └── labelsTr
β”‚Β Β   Β Β   Β Β  β”œβ”€β”€label0001.nii.gz
β”‚Β Β   Β Β   Β Β  └── ...

2.1 Pre-process LASeg dataset

Run python ./code/data/preprocess_la.py to:

2.2 Pre-process Synapse dataset

Run python ./code/data/preprocess_synapse.py to

2.3 Pre-process MMWHS dataset

Run python ./code/data/preprocess_mmwhs.py to:

2.4 Pre-process M&Ms dataset

Run python ./code/data/preprocess_mnms.py to:

For all the pre-processing, you can comment out the functions corresponding to splits and use our pre-split files.

Finally, you will get a file structure as follow:

.
β”œβ”€β”€ Synapse_data
β”‚   β”œβ”€β”€ npy
β”‚   β”‚   β”œβ”€β”€ <id>_image.npy
β”‚   β”‚   β”œβ”€β”€ <id>_label.npy
β”‚   β”‚   └── ...
β”‚   └── split_txts
β”‚       β”œβ”€β”€ labeled_0.2.txt
β”‚       β”œβ”€β”€ unlabeled_0.2.txt
β”‚       β”œβ”€β”€ train.txt
β”‚       β”œβ”€β”€ eval.txt
β”‚       β”œβ”€β”€ test.txt
β”‚       └── ...
β”œβ”€β”€ LA_data
β”‚   └── ...
β”œβ”€β”€ MMWHS_data
β”‚   └── ...
└── MNMS_data
    └── ...

3. Training & Testing & Evaluating

πŸ”₯πŸ”₯ This codebase allows train, test, and evaluate on all the four settings using one single bash file. πŸ”₯πŸ”₯

Run the following commands for training, testing and evaluating.

bash train.sh -c 0 -e diffusion -t <task> -i '' -l 1e-2 -w 10 -n 300 -d true 

Parameters:

-c: use which gpu to train

-e: use which training script, can be diffusion for train_diffusion.py, or diffusion_2d for train_diffusion_2d.py

-t: switch to different tasks:
Β  Β  Β  Β  Β  For SSL on 5% labeled LA dataset: la_0.05
Β  Β  Β  Β  Β  For IBSSL on 20% labeled Synapse dataset: synapse_0.2
Β  Β  Β  Β  Β  For UDA on MMWHS dataset: mmwhs_ct2mr for labeled CT and unlabeled MR, mmwhs_mr2ct in opposite
Β  Β  Β  Β  Β  For SemiDG on M&Ms dataset, 2% labeled B,C,D -> A setting: mnms_A_0.02; 5% labeled A,B,C -> D setting: mnms_D_0.05

-i: name of current experiment, can be whatever you like

-l: learning rate

-w: weight of unsupervised loss

-n: max epochs

-d: whether to train, if true, training -> testing -> evaluating; if false, testing -> evaluating

4. Results & Model Weights


🌟🌟 All trained model weights can be downloaded from this link. 🌟🌟


Put the logs directory under the root directory of this repo and set -d False, then you can test and evaluate the models.

4.1 Imbalance SSL on Synapse dataset

Result_IBSSL.png

4.2 SSL on LASeg dataset & UDA on MMWHS dataset

Result_SSL_UDA.png

4.3 SemiDG on M&Ms dataset

Result_SemiDG.png

Citations

If this code is helpful for your study, please cite:

@inproceedings{wang2023towards,
  title={Towards Generic Semi-Supervised Framework for Volumetric Medical Image Segmentation},
  author={Wang, Haonan and Li, Xiaomeng},
  booktitle={Thirty-seventh Conference on Neural Information Processing Systems},
  year={2023}
}

Contact

Haonan Wang (hwanggr@connect.ust.hk)

License

This repository is released under MIT License.