himashi92 / co-bionet

[Nature Machine Intelligence Journal] Official pytorch implementation for Uncertainty-Guided Dual-Views for Semi-Supervised Volumetric Medical Image Segmentation

Co-BioNet: Uncertainty-Guided Dual-Views for Semi-Supervised Volumetric Medical Image Segmentation

This repo contains the supported pytorch code and configuration files to reproduce the results of Uncertainty-Guided Dual-Views for Semi-Supervised Volumetric Medical Image Segmentation Article.

Abstract

Deep learning has led to tremendous progress in the field of medical artificial intelligence. However, training deep-learning models usually require large amounts of annotated data. Annotating large-scale datasets is prone to human biases and is often very laborious, especially for dense prediction tasks such as image segmentation. Inspired by semi-supervised algorithms that use both labelled and unlabelled data for training, we propose a dual-view framework based on adversarial learning for segmenting volumetric images. In doing so, we use critic networks to allow each view to learn from high-confidence predictions of the other view by measuring a notion of uncertainty. Furthermore, to jointly learn the dual-views and the critics, we formulate the learning problem as a min–max problem. We analyse and contrast our proposed method against state-of-the-art baselines, both qualitatively and quantitatively, on four public datasets with multiple modalities (for example, computerized topography and magnetic resonance imaging) and demonstrate that the proposed semi-supervised method substantially outperforms the competing baselines while achieving competitive performance compared to fully supervised counterparts. Our empirical results suggest that an uncertainty-guided co-training framework can make two neural networks robust to data artefacts and have the ability to generate plausible segmentation masks that can be helpful for semi-automated segmentation processes.

Link to full paper:

Published in Nature Machine Intelligence : Link

Proposed Architecture

System requirements

Under this section, we provide details on the environmental setup and dependencies required to train/test the Co-BioNet model. This software was originally designed and run on a system running Ubuntu (Compatible with Windows 11 as well).
All the experiments are conducted on Ubuntu 20.04 Focal version with Python 3.8.
To train Co-BioNet with the given settings, the system requires a GPU with at least 40GB. All the experiments are conducted on Nvidia A40 single GPU. (Not required any non-standard hardware)
To test the model's performance on unseen Pancreas CT and LA MRI test data, the system requires a GPU with at least 4 GB.

Create a virtual environment

pip install virtualenv
virtualenv -p /usr/bin/python3.8 venv
source venv/bin/activate

Installation guide

• Install torch :

pip3 install torch==1.10.2+cu113 torchvision==0.11.3+cu113 torchaudio==0.10.2+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html

• Install other dependencies :

pip install -r requirements.txt

Typical Install Time

This depends on the internet connection speed. It would take around 15-30 minutes to create environment and install all the dependencies required.

Dataset Preparation

The experiments are conducted on two publicly available datasets,

• National Institutes of Health (NIH) Panceas CT Dataset : https://wiki.cancerimagingarchive.net/display/Public/Pancreas-CT
• 2018 Left Atrial Segmentation Challenge Dataset : http://atriaseg2018.cardiacatlas.org
• MSD BraTS Dataset : http://medicaldecathlon.com/

Pre-processed data can be found in folder data.

Figshare Project Page

All the pre-trained models, figures, evaluations, a video on how the training pipeline works, and the source code are included in this project page link

• DOI : https://doi.org/10.6084/m9.figshare.22140194.v5

Trained Model Weights

Download trained model weights from this shared drive link, and put it under folder code/model or code_msd_brats/model

Running Demo

Demonstration is created on generating segmentation masks on a sample of unseen Pancreas CT with trained torch models on 10% and 20% Labeled Pancreas CT and Left Atrial MRI data. You can run the given python notebook in the demo folder.

Train Model

• To train the model for Pancreas CT dataset on 10% Lableled data

cd code
nohup python train_cobionet_semi.py --dataset_name Pancreas_CT --labelnum 6 --lamda 1.0 --consistency 1.0 --mu 0.01 --t_m 0.2 --max_iteration 15000 &> pa_10_perc.out &

• To train the model for Pancreas CT dataset on 20% Lableled data

cd code
nohup python train_cobionet_semi.py --dataset_name Pancreas_CT --labelnum 12 --lamda 1.0 --consistency 1.0 --mu 0.01 --t_m 0.2 --max_iteration 15000 &> pa_20_perc.out &

• To train the model for Left Atrial MRI dataset on 10% Lableled data

cd code
nohup python train_cobionet_semi.py --dataset_name LA --labelnum 8 --lamda 0.7 --consistency 1.0 --mu 0.01 --t_m 0.4 --max_iteration 15000 &> la_10_perc.out &

• To train the model for Left Atrial MRI dataset on 20% Lableled data

cd code
nohup python train_cobionet_semi.py --dataset_name LA --labelnum 16 --lamda 0.7 --consistency 1.0 --mu 0.01 --t_m 0.4 --max_iteration 15000 &> la_20_perc.out &

• To train the model for MSD BraTS MRI dataset on 10% Lableled data

cd code_msd_brats
nohup python train_cobionet_semi.py --dataset_name MSD_BRATS --labelnum 39 --lamda 1.0 --consistency 1.0 --mu 0.01 --t_m 0.25 --max_iteration 10000 &> msd_10_perc.out &

• To train the model for MSD BraTS MRI dataset on 20% Lableled data

cd code_msd_brats
nohup python train_cobionet_semi.py --dataset_name MSD_BRATS --labelnum 77 --lamda 1.0 --consistency 1.0 --mu 0.01 --t_m 0.25 --max_iteration 10000 &> msd_20_perc.out &

It would take around 5 hours to complete model training for Pancreas and Left Atrium datasets. For MSD BraTS dataset, it will take around 12 hours to complete training. You can try out different hyper-parameter settings and further improve the accuracy.

Hyperparameter Setting and Experimental Results for different data splits.

Test Model

• To test the Co-BioNet ensemble model for Pancreas CT dataset on 10% Lableled data

cd code
python eval_3d_ensemble.py --dataset_name Pancreas_CT --labelnum 6

Acknowledgements

This repository makes liberal use of code from SASSNet, UAMT, DTC and MC-Net

Citing Co-BioNet

If you find this repository useful, please consider giving us a star ⭐ and cite our work:

     @article{peiris2023uncertainty,
        title={Uncertainty-guided dual-views for semi-supervised volumetric medical image segmentation},
        author={Peiris, Himashi and Hayat, Munawar and Chen, Zhaolin and Egan, Gary and Harandi, Mehrtash},
        journal={Nature Machine Intelligence},
        pages={1--15},
        year={2023},
        publisher={Nature Publishing Group UK London
      }
}

      Peiris, Himashi (2023): Project Contributions. figshare. Journal contribution. https://doi.org/10.6084/m9.figshare.22140194.v5

co-bionet's Issues

Other data sets(ACDC)

Hello Author! Have you used the ACDC dataset before in that paper model? Can I use ACDC dataset for 3D medical image segmentation? I don't know about the memory usage.

comparison experiment

Hello Author! I'd like to ask you how you did the comparison experiment, did you use the same dataset to reproduce the other articles individually or did you just replace the two segmentation models of your paper with the segmentation network of the other papers?

Is "train_cobionet_semi.py" full supervision by default?

Is "train_cobionet_semi.py" full supervision by default?

code for processing pancreas

Can you provide the code for processing pancreas?

Entropy estimation uncertainty

Hello, I am reading through your paper and I have a question that I am not sure can be answered:
The model says that it is trained by estimating uncertainty through entropy, may I ask how exactly the uncertainty is estimated through entropy in the modeling process? Can you elaborate on the process including the code implementation part, thank you very much.

Thank you very much for your excellent work! I have two questions:

Thank you very much for your excellent work! I have two questions:
(1). In the paper, where specifically does the term 'uncertainty' in the title manifest? Is it in step 6? I didn't see a step where uncertainty is calculated.
(2). How is Equation 1 derived?
Thank you.

train_cobionet_sup.py

Excuse me, what is the purpose of the train_cobionet_sup.py file in the project?

Hello author, I would like to ask directly run your code "train_cobionet_semi.py" where to set the percentage of tagged data?

Hello author, I would like to ask directly run your code "train_cobionet_semi.py" where to set the percentage of tagged data?

Test LA dataset

Hello, I'm using eval_3d_ensamble.py to test LA dataset, I get error as below image, I set the condition --dataset_name LA --labelnum 8,what's the problem please?