Comparative evaluation of training strategies using partially labelled datasets for segmentation of white matter hyperintensities and stroke lesions in FLAIR MRI

TL;DR

This study tackles the challenge of segmenting white matter hyperintensities and ischaemic stroke lesions in FLAIR MRI when annotations are partially available. It systematically evaluates six supervision strategies on 2052 volumes from 12 datasets, finding that pseudolabels provide the strongest overall gains, especially on datasets not represented in the fully labeled subset. The results show that leveraging partially labelled data can boost key metrics such as $AP$ and $DSC$, while also revealing evaluation challenges due to ground-truth variability and cross-dataset annotation policies. The findings offer practical guidance for training with partial labels in medical image segmentation and highlight the importance of pseudolabel quality for reliable cross-dataset deployment.

Abstract

White matter hyperintensities (WMH) and ischaemic stroke lesions (ISL) are imaging features associated with cerebral small vessel disease (SVD) that are visible on brain magnetic resonance imaging (MRI) scans. The development and validation of deep learning models to segment and differentiate these features is difficult because they visually confound each other in the fluid-attenuated inversion recovery (FLAIR) sequence and often appear in the same subject. We investigated six strategies for training a combined WMH and ISL segmentation model using partially labelled data. We combined privately held fully and partially labelled datasets with publicly available partially labelled datasets to yield a total of 2052 MRI volumes, with 1341 and 1152 containing ground truth annotations for WMH and ISL respectively. We found that several methods were able to effectively leverage the partially labelled data to improve model performance, with the use of pseudolabels yielding the best result.

Figures (5)

• Figure 1: Pictorial depiction of the various approaches to utilising partially labelled data, that are studied in this paper. The multiclass baseline simply trains a multiclass model on the fully labelled subset (FLS) of the training data. The multi-model approach trains two binary models, one for segmenting white matter hyperintensities (WMH) and one for segmenting ischaemic stroke lesions (ISL), using the respective partially labelled subsets (PLS) of the training data. The class-conditional model has two output heads, one for WMH and one for ISL; if both labels are present for a training sample the forward pass is run twice (once for each head) and the average loss is calculated. The pseudolabel-trained model is a standard multiclass model, but pseudolabels are used for the missing labels. The marginal loss-trained model uses the marginal loss formuation where missing classes are merged into the background. The class-adaptive loss-trained model calculates the loss on the available classes only. The two-phase-trained model is pre-trained on all the PLS data with the labels for WMH and ISL being merged into one "not background" label, and then the final layer is replaced and the model is fine-tuned on the FLS data. Note: $PLS_{all} = PLS_{WMH} \cup PLS_{ISL}$ and $FLS \subset PLS_{all}$.
• Figure 2: Test set predictions of the pseudolabels model, grouped into predictions that appear worse, different, or better than the ground truth labels.
• Figure 3: Average precision (AP) for (a) WMH; (b) ISL. Boxplot whiskers extend to 1.5 times the interquartile range.
• Figure 4:
• Figure 5: (a) An LBC1936 flair scan where it is clear that the ground truth (red) includes very small puncate areas of hyperintensity, often no larger than a single voxel, while the prediction (yellow) is more conservative; (b) An LBC1921 flair scan where the ground truth does not include small disconnected wmh, but rather smoothly segmented areas of hyperintensity.