LST-AI: A deep learning ensemble for accurate MS lesion segmentation

Abstract

Automated segmentation of brain white matter lesions is crucial for both clinical assessment and scientific research in multiple sclerosis (MS). Over a decade ago, we introduced an engineered lesion segmentation tool, LST. While recent lesion segmentation approaches have leveraged artificial intelligence (AI), they often remain proprietary and difficult to adopt. As an open-source tool, we present LST-AI, an advanced deep learning-based extension of LST that consists of an ensemble of three 3D U-Nets. LST-AI explicitly addresses the imbalance between white matter (WM) lesions and non-lesioned WM. It employs a composite loss function incorporating binary cross-entropy and Tversky loss to improve segmentation of the highly heterogeneous MS lesions. We train the network ensemble on 491 MS pairs of T1-weighted and FLAIR images, collected in-house from a 3T MRI scanner, and expert neuroradiologists manually segmented the utilized lesion maps for training. LST-AI also includes a lesion location annotation tool, labeling lesions as periventricular, infratentorial, and juxtacortical according to the 2017 McDonald criteria, and, additionally, as subcortical. We conduct evaluations on 103 test cases consisting of publicly available data using the Anima segmentation validation tools and compare LST-AI with several publicly available lesion segmentation models. Our empirical analysis shows that LST-AI achieves superior performance compared to existing methods. Its Dice and F1 scores exceeded 0.62, outperforming LST, SAMSEG (Sequence Adaptive Multimodal SEGmentation), and the popular nnUNet framework, which all scored below 0.56. Notably, LST-AI demonstrated exceptional performance on the MSSEG-1 challenge dataset, an international WM lesion segmentation challenge, with a Dice score of 0.65 and an F1 score of 0.63-surpassing all other competing models at the time of the challenge. With increasing lesion volume, the lesion detection rate rapidly increased with a detection rate of >75% for lesions with a volume between 10 mm³ and 100 mm³. Given its higher segmentation performance, we recommend that research groups currently using LST transition to LST-AI. To facilitate broad adoption, we are releasing LST-AI as an open-source model, available as a command-line tool, dockerized container, or Python script, enabling diverse applications across multiple platforms.

Keywords: Artificial Intelligence; Deep Learning; Lesion Segmentation; Magnetic Resonance Imaging; Multiple Sclerosis; White Matter Lesions.

Fig. 1

The different processing steps of the holistic LST-AI tool are presented. First, a pair of T1w and FLAIR images is warped to MNI space, then skull-stripped, cropped, and intensity-normalized during preprocessing. The resulting images are used as input for the three 3D U-Nets of the ensemble network. Each U-Net provides a lesion probability map. To generate the binary lesion map, the three lesion probability maps are averaged and a threshold is subsequently applied. Finally, the binary lesion map is warped back to the subject image space (original space of the FLAIR image).

Fig. 2

Architecture of the 3D U-Nets which constitute the ensemble network of LST-AI. They comprise two channels (one for T1w images and one for FLAIR images) and consist of 5 encoder and 5 decoder blocks. Strided convolutions (stride 2) are used for downsampling and transposed convolutions are used for upscaling. Encoder and decoder blocks are connected via skip connections.

Fig. 3

Rationale behind the ensemble network of LST-AI. First, the three 3D U-Nets generate a lesion probability map. The mean of the three outputs is calculated and thresholded to generate the final binary lesion map. On the right-hand side, we show a slice of a FLAIR image and the corresponding manual segmentation (i.e., the ground truth). The orange arrow and circle highlight a false positive present in the lesion probability map of 3D U-Net 1, but not in the other lesion probability maps. The light blue arrow and circle highlight a false positive present in the lesion probability map of 3D U-Net 2, but not in the other lesion probability maps. The green arrow and circle highlight a false negative lesion in the lesion probability map of 3D U-Net 3, which is detected by 3D U-Net 1 and 2. Note how the output of the ensemble network is more accurate than the output of the individual networks, as it does not show the false positives and false negatives. (For interpretation of the references to colour in this figure legend, the reader is referred to the web version of this article.)

Fig. 4

MS-specific anatomical mask indicating four different brain regions: ventricles outlined in light gray (used to label lesions as periventricular), cortex outlined in dark gray (used to label lesions as juxtacortical), subcortical region outlined in gray (used to label lesions as subcortical), or infratentorial region (not visible in the image). Note that lesions are dilated using a 3x3x3 mm³ cube before overlaying with the anatomical mask, which is how lesions can overlap with ventricle or cortex regions, resulting in lesions labeled as periventricular or juxtacortical, respectively.

Fig. 5

Binary lesion maps generated by the different lesion segmentation methods applied in this study. As reference, the first row shows the underlying FLAIR image as well as the manual segmentation (which is the ground truth). Each method provides slightly different lesion maps, and, in the slice presented here, LST-AI appears to be the most accurate.

Fig. 6

First-level Dice similarity coefficient (DSC) (across all test datasets) of each lesion segmentation tool are provided for lesions in different brain regions: all lesions in the whole brain, infratentorial lesions, juxtacortical lesions, periventricular lesions, and subcortical lesions.

Fig. 7

This graph shows the distribution of lesions per volume. The bars and numbers indicate how many lesions are in each volume group. We divided the lesions into groups with a volume range of 10 mm³ and the first bar from the left shows the number of lesions with a volume between 3 mm³ and 10 mm³.

Fig. 8

These graphs illustrate the proportion of lesions that are detected in each volume group. We divided the lesions into groups according to their volume (on the logarithmic scale): 3–10 mm³, 11–100 mm³, 101–1000 mm³, 1001–10000 mm³, and larger than 10000 mm³. A) shows the number of lesions distribution across the volume groups; B) − F) show the lesion detection ratios of LST-AI, LST-LGA, LST-LPA, nnUNet, and SAMSEG for the different lesion volumes. Note, how the detection rate increases with increasing lesion volume for each segmentation, whereby LST-AI yields the highest detection rates. The detection rate is given in %.