Skip to content

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

Notifications You must be signed in to change notification settings

xmed-lab/GenericSSL

Repository files navigation

[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.

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:

  • convert .h5 files to .npy.
  • generate the labeled/unlabeled splits

2.2 Pre-process Synapse dataset

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

  • resize the images and convert to .npy for faster loading;
  • generate the train/test splits (use labeled data for validation);
  • generate the labeled/unlabeled splits.

2.3 Pre-process MMWHS dataset

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

  • reorient to the same orientation, RAI;
  • convert to continuous labels;
  • crop centering at the heart region;
  • for each 3D cropped image top 2/% of its intensity histogram was cut off for alleviating artifacts;
  • resize and convert to .npy;
  • generate the train/validation/test splits.

2.4 Pre-process M&Ms dataset

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

  • split the original 4D data to 3D along the time axis;
  • crop and resize;
  • save to .npy;
  • generate the train/test splits (use labeled data for validation);
  • generate the labeled/unlabeled splits.

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.

Releases

No releases published

Packages

No packages published
pFad - Phonifier reborn

Pfad - The Proxy pFad of Β© 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy