[2022 May 30th] Happy to announce that our paper got accepted at MIDL 2022 as an oral presentation (top 11.6%)
[2023 Sep 15th] We release a new 3D version of our implementation based on a new 3D dataset of left atrium and the extended version is published as a journal paper at IEEE TMI (impact factor 11): https://ieeexplore.ieee.org/stamp/stamp.jsp?tp=&arnumber=10121397
This repository is an implementation of the MIDL2022 paper: 'Learning Morphological Feature Perturbations for Calibrated Semi-Supervised Segmentation' and the IEEE TMI paper 'MisMatch: Calibrated Segmentation via Consistency on Differential Morphological Feature Perturbations With Limited Labels'. This code base was written and maintained by Moucheng Xu
Consistency regularisation with input data perturbations in semi-supervised classification works because of the cluster assumption. However, the cluster assumption does not hold in the data space in segmentation (https://arxiv.org/abs/1906.01916). Fortunately, the cluster assumption can be observed in the feature space for segmentation. Therefore, we propose to use consistency regularisation on feature perturbations for semi-supervised segmentation and we propose to learn feature perturbations end-to-end with network architecture manipulations based on differential morphological operations.
- We provide a new interperation of ERF (effective receptive field: https://arxiv.org/abs/1701.04128) as a theoretical foundation for incorporating differential morphological operations of features in neural networks;
- Based on our insight on the connection between ERF and morphological operations, we build a new encoder-decoder network architecture of semi-supervised segmentation with two decoders:
- positive attention decoder which enforces inductive bias to do differential dilation operations on the features;
- negative attention decoder which enforces another inductive bias to do differential erosion operations on the features.
- We then apply l1 normalisation along batch dimension on the two outputs which come from the dilaiton decoder and the erosion decoder respectively before we apply a consistency loss.
See our paper (https://arxiv.org/pdf/2110.12179.pdf) for more details.
| LR | Batch | Seed | Width | Consistency | Labels | Steps |
|---|---|---|---|---|---|---|
| 0.01 | 4 | 1337 | 8 | 1 | 2 | 5000 |
Results on the LA dataset between consistency on feature perturbations (Ours) and consistency on data perturbations (UA-MT)
| Models (5000 steps) | Dice (⬆) | Jaccard (⬆) | Hausdorff Dist. (⬇) | Average Surface Dist. (⬇) |
|---|---|---|---|---|
| MisMatch (Ours) | 0.73 | 0.58 | 32 | 10 |
| UA-MT | 0.70 | 0.55 | 37 | 12 |
This repository is based on PyTorch 1.4. To use this code, please first clone the repo and install the anaconda environments via:
git clone https://github.com/moucheng2017/Morphological_Feature_Perturbation_SSL
cd MisMatchSSL
conda env create -f midl.yml
To train the baseline on LA with default hyperparameters:, use:
cd MisMatchSSL/code # change directory to your working directory where you downloaded the github repo
python train_LA_meanteacher_certainty_unlabel.py To train our proposed model MisMatch on LA with default hyperparameters, use:
cd MisMatchSSL/code # change directory to your working directory where you downloaded the github repo
python train_LA_mismatch.py To train the models on other custom datasets or the lung tumour or the brain tumour, you have to first prepare your datasets following:
--labelled/
-- imgs/
-- some_case_1.nii.gz
-- some_case_2.nii.gz
-- ...
-- lbls/
-- some_case_1.nii.gz
-- some_case_2.nii.gz
-- ...
--unlabelled/
-- imgs/
-- some_case_5.nii.gz
-- some_case_6.nii.gz
-- ...
--test/
--imgs/
-- some_case_7.nii.gz
-- some_case_8.nii.gz
-- ...
--lbls/
-- some_case_7.nii.gz
-- some_case_8.nii.gz
-- ...
To train our model on your own datasets, please use the following template, but do remember to change the data directory, if you set up "witdh" as 8 and cropping size at 96 x 96 x 96 (default in the training code, but in config), then a 12GB GPU should be enough:
cd MisMatchSSL/code # change directory to your working directory where you downloaded the github repo
python train_3D_mismatch.py \
--root_path '/directory/to/your/datasets/Task01_BrainTumour' \
--exp 'MisMatch_brain' \
--max_iterations 6000 \
--batch_size 4 \
--labeled_bs 2 \
--base_lr 0.001 \
--seed 1337 \
--width 8 \
--consistency 1.0If you find our paper or code useful for your research, please consider citing:
@inproceedings{xmc2022midl,
title={Learning Morphological Feature Perturbations for Calibrated Semi-Supervised Segmentation},
author={Xu, Moucheng and Zhou, Yukun and Jin, Chen and deGroot, Marius and Wilson Frederick J. and Blumberg, Stefano B. and Alexander, Daniel C. and Oxtoby, Neil P. and Jacob, Joseph},
booktitle = {International Conference on Medical Imaging with Deep Learning (MIDL)},
year = {2022} }
The left atrium processed h5 dataset is in the data folder. You can refer the code in code/dataloaders/la_heart_processing.py to process your own data. If you use the LA segmentation data, please also consider citing:
@article{xiong2020global,
title={A Global Benchmark of Algorithms for Segmenting Late Gadolinium-Enhanced Cardiac Magnetic Resonance Imaging},
author={Xiong, Zhaohan and Xia, Qing and Hu, Zhiqiang and Huang, Ning and Vesal, Sulaiman and Ravikumar, Nishant and Maier, Andreas and Li, Caizi and Tong, Qianqian and Si, Weixin and others},
journal={Medical Image Analysis},
year={2020} }
The lung (Task_06) and brain tumour (Task_01) datasets are downloaded from the http://medicaldecathlon.com/
Please contact 'xumoucheng28@gmail.com'
Massive thanks to my amazing colleagues including Yukun Zhou, Jin Chen, Marius de Groot, Fred Wilson, Neil Oxtoby, Danny Alexander and Joe Jacob. This code base is built upon a previous public code base on consistency on data space perturbations: https://github.com/yulequan/UA-MT