[IJCAI 2023] UCMT

November 17, 2023 · View on GitHub

This repo is the PyTorch implementation of our paper:

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

<img src=docs/UCMT.png width=75% />

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:

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 (xxszqyy@gmail.com).

Acknowledgements