Summary

This repository contains an implementation of the MICCAI 2022 Young Scientist Award Finalist paper 'Bayesian Pseudo Labels: Expectation Maximization and Maximization for Robust and Efficient Semi-Supervised Segmentation', and its journal extension 'Expectation Maximization Pseudo Labels'. This code base was written by Moucheng Xu.

Abstract:

We study pseudo-labelling. Pseudo-labelling employs raw inferences on unlabelled data as pseudo-labels for self-training. We elucidate the empirical successes of pseudo-labelling by establishing a link between this technique and the Expectation Maximisation algorithm. Through this, we realise that the original pseudo-labelling serves as an empirical estimation of its more comprehensive underlying formulation. Following this insight, we present a full generalisation of pseudo-labels under Bayes’ theorem, termed Bayesian Pseudo Labels. Subsequently, we introduce a variational approach to generate these Bayesian Pseudo Labels, involving the learning of a threshold to automatically select high-quality pseudo labels. In the remainder of the paper, we showcase the applications of pseudo-labelling and its generalised form, Bayesian Pseudo-Labelling, in the semi-supervised segmentation of medical images. Specifically, we focus on: (1) 3D binary segmentation of lung vessels from CT volumes; (2) 2D multi-class segmentation of brain tumours from MRI volumes; (3) 3D binary segmentation of whole brain tumours from MRI volumes; and (4) 3D binary segmentation of prostate from MRI volumes. We further demonstrate that pseudo-labels can enhance the robustness of the learned representations.

Pseudo Labelling as Expectation Maximization (EM)

We revisit pseudo labelling as a latent variable model. The original pseudo labelling is an empirical estimation of E-step for estimating the latent variables (pseudo labels), updating model parameters using pseudo labells is the M-step. We further simplify the graphical model by using only the confidence threshold as a latent variable in the later section. More details of the proposed graphic model can be found in our paper 'Expectation Maximization Pseudo Labels':

Bayesian Pseudo Labels

We propose Bayesian pseudo-labels, a probabilistic generalization of standard pseudo-labels. In binary segmentation, standard pseudo-labels are generated by binarizing network inferences (after applying the sigmoid function) using a fixed confidence threshold, typically 0.5. Our Bayesian pseudo-label learns an optimal threshold for unlabeled data through variational inference.

Exemplar Results on binary segmentation of whole tumour using Task01 Brain Tumour

We use the first 8 cases from the original training data from Task01 from Medical Decathlon, as labelled training cases and the rest as the testing data. We use the original unlabelled testing data as unlabelled training data.

Test Sizing	32^3	64^3	96^3	128^3
Superivsed	61.07	66.94	70.13	72.09
SegPL-VI	64.44	71.43	73.07	74.48

The beneath picture is a plot of the learnt threshold with a prior of an univariate Gaussian (mean=0.9, std=0.1) on the unlabelled data of Task01 Brain Tumour. X-axis: training iterations; Y-axis: threshold.

Installation

This repository is based on PyTorch 1.4. To use this code, please first clone the repo and install the enviroment.

Prepare the datasets

This code base is for 3D volumetric data sets. If each of your scans is in (H x W x D), please set up the flag "transpose_dim" in your config yaml file as 1. If each scan is in (D x H x W), set "transpose_dim" as 0. This code base works either numpy or nifti data formats. To use numpy, set up the flag "dataset.data_format" as 'npy', to use nifity, set up the flag "dataset.data_format" as 'nii'. You also need to prepare your dataset in a structure following:

path_to_dataset
└───labelled
|   └───imgs # labelled scans. 
|   └───lbls # labels 
└───unlabelled
|   └───imgs # unlabelled scans
└───test
    └───imgs # testing scans
    └───lbls # testing labels

Training and config files:

Run the following with your own custom yaml config file:

python main.py \
-c configs/exp.yaml

Here is an example of the configuration .yaml file:

dataset:
  name: brain
  num_workers: 8 
  data_dir: '/SAN/medic/PerceptronHead/data/Task01_BrainTumour'
  data_format: 'nii'

logger:
  tag: 'cluster_brain'

seed: 1024

model:
  input_dim: 4
  output_dim: 1
  width: 10
  depth: 3

train:
  transpose_dim: 1 # use 1 for transposing input if input is in: D x H x W. For example, for Task06_Lung from medicaldecathlon, this should be 1
  optimizer:
    weight_decay: 0.0005
  lr: 0.001
  iterations: 25000 # number of training iterations, it's worth to mention this is different from epoch
  batch: 2 # batch size of labelled volumes
  temp: 1.0 # temperature scaling on output, default as 1
  contrast: False # random contrast augmentation
  crop_aug: True # random crop augmentation
  gaussian: False # random gaussian noise augmentation
  new_size_d: 32 # crop size on depth (number of slices)
  new_size_w: 128 # crop size on width
  new_size_h: 128 # crop size on height
  batch_u: 1 # this has to be zero in supervised setting, if set up larger than 0, semi-supervised learning will be used
  pri_mu: 0.7 # mean of prior
  pri_std: 0.15 # std of prior
  flag_post_mu: 0 # flag for mean of posterior
  flag_post_std: 0 # flag for std of posterior
  flag_pri_mu: 0 # flag for mean of prior
  flag_pri_std: 0 # flag for std of prior
  alpha: 1.0 # weight on the unsupervised learning part if semi-supervised learning is used
  warmup: 0.1 # ratio between warm-up iterations and total iterations
  warmup_start: 0.4 # ratio between warm-up starting iteration and total iterations
  beta: 1 # weight for pseudo supervision loss
  conf_lower: 0.0 # lower bound for confidence threshold

checkpoint:
  resume: False # resume training or not
  checkpoint_path: '/some/path/to/saved/model' # checkpoint path

Hyper-parameters of the K-L loss of Bayesian Pseudo Labels

To use semi-supervised learning, please enable the batch of unlabelled images larger than 0: batch_u: # positive integer. In kl_loss, we now included more flexible controls and different implementations of kl loss of the threshold of the pseudo labels. See the exact KL loss implementations in libs.Loss.kld_loss or kl_loss_test.py

1. To use the original implementation of bayesian pseudo labels from the MICCAI version. set up following flags:

   1. flag_post_mu: 0
   2. flag_post_std: 0
   3. flag_pri_mu: 0
   4. flag_pri_std: 0

2. We provide an alternative of use current model prediction confidence as prior so you don't need to tune pri_mu and pri_std. To test this configuration, please set up the following flags as:

   1. flag_post_mu: 0
   2. flag_post_std: 0
   3. flag_pri_mu: 1
   4. flag_pri_std: 1

3. To use random sampling from prior dist for the threshold, without learning the kl, please set up the following flags as:

   1. flag_post_mu: 2
   2. flag_post_std: 2
   3. flag_pri_mu: 0
   4. flag_pri_std: 0

4. To use supervised learning, simply set up batch_u: 0

Example uses:

This code base has been tested on whole tumour segmentation on Task 01 brain tumour data set downloaded from Medical Segmentation Decathlon: http://medicaldecathlon.com/

Citation

If you find our paper or code useful for your research, please consider citing us. Thank you!

@inproceedings{xu2022bpl,
     title={Bayesian Pseudo Labels: Expectation Maximization and Maximization for Robust and Efficient Semi-Supervised Segmentation},
     author={Xu, Moucheng and Zhou, Yukun and Jin, Chen and deGroot, Marius and Alexander, Daniel C. and Oxtoby, Neil P. and Hu, Yipeng and Jacob, Joseph},
     booktitle = {International Conference on Medical Image Computing and Computer Assisted Interventions (MICCAI)},
     year = {2022} }

@article{xu2024bpl,
     title={Expectation Maximization Pseudo Labels},
     author={Xu, Moucheng and Zhou, Yukun and Jin, Chen and deGroot, Marius and Alexander, Daniel C. and Oxtoby, Neil P. and Hu, Yipeng and Jacob, Joseph},
     booktitle = {Medical Image Analysis},
     doi = {https://doi.org/10.1016/j.media.2024.103125},
     volume = {94},
     pages = {103125},
     year = {2024}}

Questions

Please contact 'xumoucheng28@gmail.com' for any questions.

Ackwnoledgement

Massive thanks to my amazing colleagues at UCL and GSK including Yukun Zhou, Jin Chen, Marius de Groot, Fred Wilson, Danny Alexander, Neil Oxtoby, Yipeng Hu and Joe Jacob.