SegHeD+: Segmentation of Heterogeneous Data for Multiple Sclerosis Lesions with Anatomical Constraints and Lesion-aware Augmentation

TL;DR

SegHeD+ addresses MS lesion segmentation across heterogeneous multi-site MRI datasets by enabling simultaneous learning from cross-sectional and longitudinal data and multiple annotation types. It integrates longitudinal, volumetric, and spatial anatomical constraints into a 3D V-Net with four heads, and couples this with lesion-level augmentation via LesionMix to expand training data and balance lesion types. The method demonstrates strong performance on five MS datasets, achieving state-of-the-art results for all-lesion segmentation and competitive performance for new-lesion segmentation, while substantially improving vanishing-lesion detection through augmentation and temporal constraints. By jointly modeling all, new, and vanishing lesions and enforcing plausible temporal evolution, SegHeD+ offers a practical path toward robust, scalable MS lesion analysis across diverse imaging cohorts and timepoints.

Abstract

Assessing lesions and tracking their progression over time in brain magnetic resonance (MR) images is essential for diagnosing and monitoring multiple sclerosis (MS). Machine learning models have shown promise in automating the segmentation of MS lesions. However, training these models typically requires large, well-annotated datasets. Unfortunately, MS imaging datasets are often limited in size, spread across multiple hospital sites, and exhibit different formats (such as cross-sectional or longitudinal) and annotation styles. This data diversity presents a significant obstacle to developing a unified model for MS lesion segmentation. To address this issue, we introduce SegHeD+, a novel segmentation model that can handle multiple datasets and tasks, accommodating heterogeneous input data and performing segmentation for all lesions, new lesions, and vanishing lesions. We integrate domain knowledge about MS lesions by incorporating longitudinal, anatomical, and volumetric constraints into the segmentation model. Additionally, we perform lesion-level data augmentation to enlarge the training set and further improve segmentation performance. SegHeD+ is evaluated on five MS datasets and demonstrates superior performance in segmenting all, new, and vanishing lesions, surpassing several state-of-the-art methods in the field.

Figures (6)

• Figure 1: Visualisation of the proposed framework. SegHeD+ learns from heterogeneous datasets varying in image and label formats. SegHeD+ takes up to four inputs, and can analyse cross-sectional and longitudinal data. Four segmentation heads provide binary-class segmentation to all lesions (red) in two timepoints, and new (green) and vanishing (dashed blue) lesions in a second timepoint.
• Figure 2: Example outputs of LesionMix augmentation. Red labels denote lesions which are present in the original image; green labels denote new-lesion generated in the augmented image; blue labels denote vanishing-lesions which have been inpainted from the original image. Images best viewed online.
• Figure 3: Visual illustration of the size of original training data and LesionMix-augmented training data.
• Figure 4: Qualitative comparison of all-lesion (top row) and new-lesion (bottom row) segmentation performance. Yellow regions denote false positive segmentations, whereas cyan regions denote false negative segmentations. SegHeD+ exhibits fewer false segmentations. Best viewed online.
• Figure 5: SegHeD+ is capable of simultaneous multi-task segmentation (Rows 3 to 6). Some tasks do not show new/vanishing-lesions predictions as they are not present at the given slice. "Not available" denotes no ground truth annotation for comparison. A: Dataset where all-lesion labels are available for first and second timepoints. B: Dataset where first timepoint all-lesion label and second timepoint new-lesion label are available. C: Dataset where second timepoint vanishing-lesion label is available.