SSL-DG: Rethinking and Fusing Semi-supervised Learning and Domain Generalization in Medical Image Segmentation
This code requires PyTorch 1.10 and python 3+. Please install dependencies by
pip install -r requirements.txt
We conduct datasets preparation following CSDG
1) Abdominal MRI
-
Download Combined Healthy Abdominal Organ Segmentation dataset and put the
/MR
folder under./data/CHAOST2/
directory -
Converting downloaded data (T2 SPIR) to
nii
files in 3D for the ease of reading.
run ./data/abdominal/CHAOST2/s1_dcm_img_to_nii.sh
to convert dicom images to nifti files.
run ./data/abdominal/CHAOST2/png_gth_to_nii.ipynp
to convert ground truth with png
format to nifti.
- Pre-processing downloaded images
run ./data/abdominal/CHAOST2/s2_image_normalize.ipynb
run ./data/abdominal/CHAOST2/s3_resize_roi_reindex.ipynb
The processed dataset is stored in ./data/abdominal/CHAOST2/processed/
1) Abdominal CT
- Download Synapse Multi-atlas Abdominal Segmentation dataset and put the
/img
and/label
folders under./data/SABSCT/CT/
directory
1.Pre-processing downloaded images
run ./data/abdominal/SABS/s1_intensity_normalization.ipynb
to apply abdominal window.
run ./data/abdominal/SABS/s2_remove_excessive_boundary.ipynb
to remove excessive blank region.
run ./data/abdominal/SABS/s3_resample_and_roi.ipynb
to do resampling and roi extraction.
The details for cardiac datasets will be given later.
The processed datasets can be found at processed datasets. Download and unzip the file where the folder structure should look this: To maintain anonymity, the detailed link will be updated later.
SSL-DG
├── ...
├── data
│ ├── abdominal
│ │ ├── CHAOST2
│ │ │ ├── processed
│ │ ├── SABSCT
│ │ │ ├── processed
│ ├── cardiac
│ │ ├── processed
│ │ │ ├── bSSFP
│ │ │ ├── LGE
├── ...
Download the pretrained model and unzip the file where the folder structure should look this: To maintain anonymity, the detailed link will be updated later.
SSL-DG
├── ...
├── logs
│ ├── 2023-xxxx-xx-xx-xx
│ │ ├── checkpoints
│ │ │ ├── latest.pth
│ │ ├── configs
│ │ │ ├── xx.yaml
│ │── visuals
│ │── test_visual
├── ...
1) Example for Cross-sequence Cardiac Dataset
For direction bSSFP -> LEG with 50% annotated samples (DICE 85.87), run the command
python test.py -r logs/2023-07-31T10-47-53_seed22_efficientUnet_bSSFP_to_LEG_labelnum_0.5
For direction LEG -> BSSFP with 20% annotated samples (DICE 83.15), run the command
test.py -r logs/2023-08-01T19-14-19_seed22_efficientUnet_LEG_to_BSSFP_labelnum_0.2
Each log contains visual results of the test images.
And more Pretrained models will be given later.
To reproduce the performance, you need one 3080 GPU
1) Cross-modality Abdominal Dataset
For direction CT -> MRI, run the command
python main.py --base configs/efficientUnet_SABSCT_to_CHAOS.yaml --seed 22 --labeled_bs 0.5 --labelnum 0.1/0.2/0.5
For direction MRI -> CT, run the command
python main.py --base configs/efficientUnet_CHAOS_to_SABSCT.yaml --seed 22 --labeled_bs 0.5 --labelnum 0.1/0.2/0.5
2) Cross-sequence Cardiac Dataset
For direction bSSFP -> LEG, run the command
python main.py --base configs/efficientUnet_bSSFP_to_LEG.yaml --seed 22 --labeled_bs 0.5 --labelnum 0.1/0.2/0.5
For direction LEG -> bSSFP, run the command
python main.py --base configs/efficientUnet_LEG_to_bSSFP.yaml --seed 22 --labeled_bs 0.5 --labelnum 0.1/0.2/0.5
Here, we open source the initial version of SSL-DG with the segmentation model containing two parallel networks. Later, we will further open source more versions of SSL-DG to facilitate community collaboration and advancement. Subsequent versions will include a 2D segmentation model with three parallel networks as well as 3D image segmentation applications.
Our codes are built upon CSDG, SLAug, and MC-Net, thanks for their contribution to the community and the development of researches!