Senyh / UCMT

[IJCAI 2023] Co-training with High-Confidence Pseudo Labels for Semi-supervised Medical Image Segmentation

[IJCAI 2023] UCMT

This repo is the PyTorch implementation of our paper:

"Co-training with High-Confidence Pseudo Labels for Semi-supervised Medical Image Segmentation"

Uncertainty-guided Collaborative Mean-Teacher (UCMT)

Usage

🔥🔥 the 3D version of UCMT has been uploaded. 🔥🔥

0. Requirements

The code is developed using Python 3.7 with PyTorch 1.11.0. All experiments in our paper were conducted on a single NVIDIA Quadro RTX 6000 with 24G GPU memory.

Install from the requirements.txt using:

pip install -r requirements.txt

1. Data Preparation

1.1. Download data

The original data can be downloaded in following links:

• ISIC Dataset - Link (Original)

1.2. Split Dataset

The ISIC dataset includes 2594 dermoscopy images and corresponding annotations. Split the dataset, resulting in 1815 images for training and 779 images for testing.

python data/split_dataset.py

Then, the dataset is arranged in the following format:

DATA/
|-- ISIC
|   |-- TrainDataset
|   |   |-- images
|   |   |-- masks
|   |-- TestDataset
|   |   |-- images
|   |   |-- masks

2. Training

2.1 Adopting DeepLabv3Plus as backbone:

python train.py --backbone DeepLabv3p

2.2 Adopting U-Net as backbone:

python train.py --backbone UNet

3. Evaluation

python eval.py

4. Visualization

python visualization.py

Citation

If you find this project useful, please consider citing:

@inproceedings{ijcai2023p467,
  title     = {Co-training with High-Confidence Pseudo Labels for Semi-supervised Medical Image Segmentation},
  author    = {Shen, Zhiqiang and Cao, Peng and Yang, Hua and Liu, Xiaoli and Yang, Jinzhu and Zaiane, Osmar R.},
  booktitle = {Proceedings of the Thirty-Second International Joint Conference on
               Artificial Intelligence, {IJCAI-23}},
  publisher = {International Joint Conferences on Artificial Intelligence Organization},
  editor    = {Edith Elkind},
  pages     = {4199--4207},
  year      = {2023},
  month     = {8},
  note      = {Main Track},
  doi       = {10.24963/ijcai.2023/467},
  url       = {https://doi.org/10.24963/ijcai.2023/467},
}

Contact

If you have any questions or suggestions, please feel free to contact me ([e-mail]).

Acknowledgements

• TorchSemiSeg
• UA-MT
• SSL4MIS