Self-supervised brain lesion generation for effective data augmentation of medical images

Abstract

Accurate brain lesion delineation is important for planning neurosurgical treatment. Automatic brain lesion segmentation methods based on convolutional neural networks have demonstrated remarkable performance. However, neural network performance is constrained by the lack of large-scale well-annotated training datasets. In this manuscript, we propose a comprehensive framework to efficiently generate new samples for training a brain lesion segmentation model. We first train a self-supervised lesion generator based on the adversarial autoencoder to model lesion appearance and shape. Next, we utilize a novel image composition algorithm, Soft Poisson Blending, to seamlessly combine synthetic lesions and brain images to obtain training samples. Finally, to effectively train the brain lesion segmentation model with augmented images we introduce a new prototype consistence regularization to align real and synthetic features. Our framework is validated by extensive experiments on two public brain lesion segmentation datasets: ATLAS v2.0 and Shift MS. Our method outperforms existing brain image data augmentation schemes. For instance, our method improves the Dice from 50.36% to 60.23% compared to the UNet with conventional data augmentation techniques for the ATLAS v2.0 dataset.

1. Introduction

Brain lesions are often indicative of serious neurological conditions, from cancer to stroke. Magnetic Resonance (MR) Imaging is widely used to detect brain lesions as MR provides excellent soft-tissue contrast, allowing for a clear distinction between healthy and abnormal brain tissue. Accurate segmentation of brain lesions is crucial for quantitative analysis of lesion progression and planning surgical treatments. The current clinical standard is human delineation of the brain lesion boundary by an expert which is tedious, time-consuming, and costly. Neural networks have emerged as a promising technique to automate brain lesion segmentation (Pereira, Pinto, Alves, & Silva, 2016; Ronneberger, Fischer, & Brox, 2015). However, training neural networks requires large amounts of well-annotated images to ensure good performance. The need for large training datasets limits the development of automatic brain lesion segmentation models since the scale of brain lesion datasets is often limited.

Data augmentation is a widely used technique to increase training dataset size and diversity in order to improve model performance. Conventional data augmentation strategies include random rotations, brightness adjustment, etc. Although these simple spatial and appearance transformations improve segmentation performance to some extent, they do not fundamentally increase the diversity of the dataset and provide a smaller boost in performance compared to acquiring new data. Recently, methods to perform data augmentation based on image fusion (such as Mixup Zhang, Cisse, Dauphin, & Lopez-Paz, 2018, CutMix Yun et al., 2019, etc.) have been developed to increase training dataset diversity. However, these approaches may shift the distribution of the original dataset (Pinto, Yang, Lim, Torr, & Dokania, 2022), which is catastrophic for small datasets as the model learns non-representative features in the augmented images (Zhang et al., 2018). Neural network approaches, such as Generative Adversarial Networks (GANs) (Goodfellow et al., 2014) have been proposed to synthesize data for model training. However, the performance of GANs, similar to other neural networks, is constrained by the size of the training dataset, posing challenges in generating realistic images when only limited data is available. Regardless of the method used to create augmented samples, combining augmented and real samples does not guarantee the segmentation model will learn discriminative features to segment lesions on real samples as there is no supervision in the feature space. Therefore, we raise two questions: How can we generate realistic images that will not shift the original data distribution with limited training samples? How can we effectively use synthetic samples to train a segmentation model to perform well on real samples?

We present a comprehensive framework to effectively augment brain imaging data with synthetic lesions to train a lesion segmentation model. Our framework has three stages: (1) train a lesion generator in a self-supervised manner, and sample feasible latent vectors from a constrained embedding space to create realistic paired lesion images and masks; (2) blend lesion images into existing brain images leveraging our novel image composition technique called Soft Poisson Blending (SPB); (3) train the segmentation model using a prototype consistency regularization term to align real and synthetic lesion features for better performance. The main contributions of our work are:

• We develop a data augmentation method to enhance the performance of neural networks trained to perform segmentation tasks. It comprises a two-stage Adversarial Autoencoder (AAE), consisting of shape and intensity generation networks, to generate on-the-fly new foreground regions when training. Distinct from other GANs, we designed a self-supervised approach where we simulate image-label pairs for training our AAE. This enables our AAE to learn a larger distribution of lesions and improve the quality of synthetic results.

• We introduce Soft Poisson Blending (SPB), based on Poisson Image Editing (Pérez, Gangnet, & Blake, 2023), to ensure realistic and smooth boundaries when inserting generated lesions into a brain image. SPB computes a refined guidance vector field that adjusts intensity values to make the synthetic lesion appear similar to the surrounding brain tissues.

• We design a new loss term, prototype consistency regularization, to learn common features between synthetic and real training samples.

2. Related works

2.1. Data augmentation for data scarcity

To mitigate the problem of data scarcity when training neural networks, data augmentation techniques are used to increase the training dataset size thereby improving model performance. Conventional data augmentation techniques include shape transformations (flip, rotation, scaling, etc.) and appearance transformations (color jittering, brightness, and contrast adjustment, etc.) (Isensee, Jaeger, Kohl, Petersen, & Maier-Hein, 2021; Krizhevsky, Sutskever, & Hinton, 2012). However, the diversity of the dataset achieved by such transformations is limited. Furthermore, the intrinsic characteristics of the training samples are not fundamentally changed, limiting improvements in model performance. The development of GANs enabled more advanced data augmentation where entirely new samples are created by the model. GANs can generate realistic samples for both natural (Brock, Donahue, & Simonyan, 2018) and medical (Nie et al., 2017; Schlegl, Seeböck, Waldstein, Langs, & Schmidt-Erfurth, 2019) images. GANs can also perform image-to-image translation, where the generator learns to map between modalities. Nevertheless, GANs require large datasets when training the model to enable accurate image generation. Other generative models such as Variational Autoencoders (VAEs) (Kingma et al., 2013) and diffusion models (Ho, Jain, & Abbeel, 2020; Liu, Ren, Lin, & Zhao, 2022) have been used to synthesize medical images. VAEs aim to regularize the latent space using a Kullback–Leibler (KL) divergence loss to encourage structured sampling and enable image reconstruction. However, VAEs often produce blurry outputs and are prone to training instability, particularly when the dataset size is small. Diffusion models have recently emerged as a powerful alternative, generating high-fidelity images through iterative denoising processes. Despite their impressive visual quality in natural image synthesis, diffusion models typically require large datasets during training and are computationally expensive due to the iterative denoising steps, making them less suitable for low-resource medical imaging scenarios.

Recently, image-manipulation-based data augmentation techniques (Yun et al., 2019; Zhang et al., 2018, 2023) have been developed to increase training dataset size and variety by using a set of image manipulation rules to generate new samples. For instance, CutMix (Yun et al., 2019) cuts and fuses patches from different images to create new training samples. Such techniques must be carefully designed and may shift the distribution of the training dataset especially when the original training dataset size is small (Pinto et al., 2022). In this work, we choose to adopt Adversarial Autoencoders (AAEs) for lesion generation as they effectively combine the benefits of both autoencoders and adversarial training. Unlike standard VAEs, AAEs allow us to enforce a structured latent space while maintaining the capacity to generate sharper and more realistic outputs through adversarial loss. Compared to traditional GANs, AAEs are more stable to train on small medical datasets and more suitable for our self-supervised framework, as they can learn from pseudo lesion–mask pairs without relying on paired real annotations. This balance between latent space control and visual quality makes AAEs particularly well-suited for synthesizing diverse and plausible lesions in limited-data regimes.

2.2. Poisson blending in deep learning

Poisson blending is an image processing technique to seamlessly integrate regions of a source image into a target image. In the context of data augmentation, this typically involves identifying the foreground of a source image and integrating this region into a target image. The blending process is guided by the gradient of the source image and the intensity of the target image (Pérez et al., 2023). Poisson blending generates more coherent images compared to simpler fusion methods such as Copy–Paste (Ghiasi et al., 2021). Poisson blending has been used for data augmentation for a variety of natural images (Liu et al., 2021). In the context of medical imaging, Tan et al. (2021) developed a self-supervised learning strategy to detect abnormalities in chest X-ray images, and abnormalities were generated by fusing image patches into the target image using Poisson blending. Wang et al. (2022) utilized Poisson blending to generate retinal images with lesions and different appearances than in the original training dataset to improve the performance of a lesion segmentation model. Lee and Cho (2023) introduced Poisson blending to the gastric disease classification task for better model generalization ability. However, all of these applications were applied to 2D images, and to our knowledge, Poisson blending has not yet been applied to 3D medical images.

2.3. Prototype learning

Prototype networks (Snell, Swersky, & Zemel, 2017) were first presented for few-shot learning, where the network learns to symbolize each class using a prototype(a.k.a. a representative vector in the embedding space). Each class prototype is typically obtained by computing the average feature of samples belonging to the class. This method was initially applied in image classification where distances between different class prototypes were maximized across training samples (Snell et al., 2017). Prototype learning was extended to image segmentation by computing cosine similarity between the class prototypes and individual pixel features. In this context for the test dataset pixels the class with which they have the highest similarity to its prototype. Wang, Liew, Zou, Zhou, and Feng (2019) designed a bidirectional prototype alignment mechanism for the few-shot image segmentation task. Kuo, Ma, Huang, and Kira (2020) utilized class prototypes to augment training samples in the feature space. (Xu et al., 2022) introduced a cyclic prototype consistency framework for semi-supervised medical image segmentation. In our work, we draw on the idea of prototype consistency to introduce a regularization term during segmentation model training to align features between synthetic and real samples.

3. Methods

Fig. 2 illustrates our entire framework comprised of three stages: I. training the lesion generator, II. inserting synthetic lesions into brain images, and III. using the generated images for segmentation model training and incorporating prototype consistency regularization. Exemplar synthetic lesions and real lesions are shown in Fig. 1(a). The synthetic lesions have appearances similar to those of real lesions.

Fig. 2. Overview of our framework containing three stages, operating on 3D volumes (visualized with 2D slices for illustration purposes).

Fig. 1.

(a) 2D slice visualizations of 3D volumes from the ATLAS v2.0 (top) and MS Shift dataset (bottom) showing real lesions (green arrow) and synthetic lesions (orange arrow), illustrating that the synthetic lesions have a similar appearance to real ones. Note that all operations are performed on full 3D volumes, and the displayed images are 2D slices for illustration only. (b) t-SNE of the intensity embedding space for real lesions, synthetic lesions, and training samples. (c) t-SNE of the shape embedding space for real lesions, synthetic lesions, and training samples.

3.1. Self-supervised AAE for brain lesion synthesis

Training a 3D generative model with a small dataset is challenging and achieving a good model performance is unlikely, therefore, we decompose brain lesion generation into two sub-tasks to reduce model complexity and use a self-supervised learning strategy during training. We designed two models: a Shape Adversarial Autoencoder (Shape AAE) and an Intensity Adversarial Autoencoder (Intensity AAE), to first create lesion masks and then perform texture synthesis to generate lesion images corresponding to the masks. Both shape and intensity AAEs are trained in a self-supervised manner. For lesion synthesis real lesion images are used to define the data distribution in the latent space and synthetic lesion images are created by sampling latent vectors from only this distribution before going through the decoder block of the trained AAEs.

3.1.1. Shape and intensity AAE design

We follow the model architecture proposed in (Rombach, Blattmann, Lorenz, Esser, & Ommer, 2022) to design the shape and intensity AAEs. As shown in Fig. 2, each AAE contains an encoder E, a decoder G, an image discriminator D_x, and a latent discriminator D_z. Although the structures are similar, for the intensity AAE we introduce a mask embedding block (MEB) (Huo et al., 2022) to provide shape guidance when generating the lesion intensity. The detailed structure of MEB is shown in Fig. 3. It embeds the mask to the feature space to control the shape of synthetic lesions.

Fig. 3. The detailed structure of MEB block.

For training the AAE models we use a three-term loss: reconstruction loss ℒ_rec, latent adversarial loss ℒ_{adv_z}, and image adversarial loss ℒ_{adv_x} · ℒ_rec is computed as the mean absolute error (MAE) between the input image Ι and the reconstructed image $\widehat{{\rm I}}$:

$${ℒ}_{rec}=\Vert {\rm I}-\widehat{{\rm I}}{\Vert }_1.$$ (1)

ℒ_rec guarantees $\widehat{{\rm I}}$ and Ι look similar in general.

The latent adversarial loss ℒ_{adv_z} is formulated to ensure that the latent space of the lesions has a normal distribution:

$$\begin{array}{l}{ℒ}_{ad{v}_{-}z}(D)=\mathbb{E}\left[max\left(0,1+D_z(E({\rm I}))\right)\right]\\ +\mathbb{E}\left[max\left(0,1-D_z(\mathcal{N}(0,1))\right)\right],\end{array}$$ (2)

$${ℒ}_{ad{v}_{-}z}(G)=-\mathbb{E}\left[D_z(E({\rm I}))\right],$$ (3)

where 𝒩 (0, 1) is a normal distribution of mean 0 and standard deviation 1. Similarly, the image adversarial loss ℒ_{adv_x} is designed to ensure the reconstructed image is realistic in appearance. It is defined as:

$$\begin{array}{l}{ℒ}_{ad{v}_{-}x}(D)=\mathbb{E}\left[max\left(0,1+D_x(\widehat{{\rm I}})\right)\right]\\ +\mathbb{E}\left[max\left(0,1-D_x({\rm I})\right)\right],\end{array}$$ (4)

$${ℒ}_{ad{v}_{-}x}(G)=-\mathbb{E}\left[D_x(\widehat{{\rm I}})\right].$$ (5)

3.1.2. Training set generation

Unlike other generative models which train models with real image– mask pairs, we train our model in a self-supervised manner by generating lesion–mask pairs. As shown in Fig. 1(b) and (c), the latent distribution of real lesions (pink triangles) is a subset of the pre-generated lesions (gray dots), which indicates training on pre-generated lesion–mask pairs is sufficient to learn a representative feature embedding space for real lesion–mask pairs.

The pre-generated lesion–mask pairs are created as follows. Inspired by Hu et al. (2023), we first generate n | n ~ 𝒰 (1, 5) ellipsoids with overlap to simulate a general lesion shape. The half-axis lengths of three directions follow the uniform distribution 𝒰 (5, 15). Elastic deformations controlled by σ| σ ~ 𝒰 (3, 6) are applied to the ellipsoids to make the shape more natural and irregular. Finally, we add Perlin noise (Perlin, 1985) to make the boundary more irregular. A comparison between real lesion masks and those generated by this approach is shown in Fig. 4(a) and (b). Note the complexity of the pre-generated masks exceeds that of real lesion masks, this enhances the AAE’s ability to learn the reconstruction task.

Fig. 4.

Lesion mask in three views for (a) real masks and (b) pre-generated masks for the shape AAE model training. Lesion images in three views for (c) real images and (d) pre-generated images for the intensity AAE model training.

To generate lesion images, we randomly select a location within the brain image from the training set and then extract the intensity values for voxels inside the pre-generated mask within that region. To increase variation in the training set, we apply the foreground intensity perturbation (Huo et al., 2023) to randomly adjust the intensity values. Fig. 4(c) and (d) show real lesions and those generated by our approach. Note the styles of the images are similar, which ensures the pre-generated lesion–mask pairs are suitable for the model training.

3.1.3. Constrained sampling for new lesion synthesis

Once AAE model training is complete, we freeze both shape AAE and intensity AAE model weights. As we see in Fig. 1(b) and (c), the latent space of real lesions is a subset of the latent space of the lesions generated for training the AAE models. Therefore, we use a constrained sampling strategy to synthesize lesions that are more similar to real lesions in the latent space. Specifically, we first obtain the latent embedding vectors for real lesion masks and lesion intensity images using the shape encoder E_S and the intensity encoder E_I. Next, we apply Principal Component Analysis (PCA) to these vectors to obtain a dimensionality-reduced latent space representation. Note we apply PCA to the shape and intensity latent embedding vectors separately. We keep the top K principal components for each space which cover 90% of the latent embedding vector variance.

To create a new synthetic lesion we uniformly sample two vectors, one from the dimensionality-reduced shape latent space and one from the dimensionality-reduced intensity latent space. We then map these dimensionality-reduced latent vectors to the original latent embedding spaces. Finally, we use these two latent embedding vectors as the input to the shape decoder G_S and intensity decoder G_I, respectively, to generate new synthetic lesion masks and images.

3.2. Lesion image composition

After we generate the synthetic lesion images, we have to fuse the lesion images with a brain image to generate training samples for the segmentation model training. We first select a plausible location in the brain for the generated lesion, then create a composite image using our modified Poisson image editing method called Soft Poisson Blending (SPB) to ensure the boundary between the image and lesion is seamless. We detail this approach below.

3.2.1. Lesion location selection

The brain has a regular anatomical structure, with distinct regions including the ventricle and brain stem. The location of brain lesions depends on the underlying pathology (e.g.,, stroke and multiple sclerosis). For the datasets in this work, we use FastSurfer (Henschel et al., 2020) to segment the white matter area of the original patient’s brain as a region proposal for lesion location selection. After obtaining the mask for white matter areas, we apply morphological binary erosion operation to shrink the masked region so that the synthetic lesion area will not exceed the mask boundary. We randomly (following a uniform distribution) select a voxel in the masked area as the center of the synthetic lesion.

3.2.2. Soft Poisson blending

We developed Soft Poisson Blending (SPB), which is a modified implementation of Poisson Image Editing (Pérez et al., 2023). The key idea of Poisson Blending is to use the Poisson partial differential equation under the Dirichlet boundary condition to specify the intensity value at the boundary area. We first adapt the conventional Poisson Blending approach to apply to 3D images. Second, we refine the guidance vector field, to ensure that the lesion foreground exhibits a natural internal appearance while having edge consistency with the background image.

For a brain image s, the target region that we would like to blend with the lesion image g, we define as Ω. The boundary of Ω is defined as ∂Ω. f, the value function defined on Ω, has a value of f * at ∂Ω. We solve the optimization problem defined as:

$$\underset{f}{\min }{{\displaystyle \int \int }}_{\text{Ω}}|\nabla f-V(x){|}^2,{f|}_{\partial \text{Ω}}={f^{\ast }|}_{\partial \text{Ω}},x\in \text{Ω}$$ (6)

where V is the guidance vector field and ∇ is the gradient operator. This equation guarantees that: (i) The gradient of the foreground content is as close as possible to V. (ii) The boundary pixel values of the foreground are consistent with the existing s, i.e., a seamless transition. The solution under the Dirichlet boundary condition is the Poisson equation:

$$\text{Δ}f=\text{div}V(x),{f|}_{\partial \text{Ω}}={f^{\ast }|}_{\partial \text{Ω}},x\in \text{Ω}$$ (7)

where Δ is the Laplacian operator, and div is the divergence operator. The guidance vector field V (Fig. 5(f)) is calculated as the mixed gradient of the brain image s (Fig. 5(a)), and the synthetic lesion image g (Fig. 5(d)) by selecting (∇s, ∇g). However, using this definition of V to construct the blended image can lead to an unnatural appearance (Fig. 5(c)). This unnatural appearance is caused by the absolute value of ∇g on ∂Ω being much higher than ∇s since for g regions outside of the foreground are zero. Based on this observation, for Soft Poisson Blending, we modified the computation of V as follows:

$${V(x)|}_{x\in \text{Ω}}=\{\begin{array}{cc}\nabla s(x)& |\nabla s(x)|>|\nabla g(x)|\text{\&}x\in \partial \text{Ω},\\ \nabla g(x)& \text{otherwise}.\end{array}$$ (8)

Fig. 5.

(a) The target brain image used for the background image. (b) The composited image based on SPB. (c) The composited image based on the original Poisson Blending. (d) The synthetic brain lesion. (e) The guidance vector field used for SPB. (f) The guidance vector field used for the original Poisson Blending. The yellow arrow points to the gradient values on the region boundary. The boundary gradient values in the original Poisson blending are disproportionately small due to zero-padding outside the lesion mask, resulting in unnatural transitions and less realistic composited images.

This results in the blended image becoming more realistic (Fig. 5(e)) compared to the original Poisson Blending algorithm (Fig. 5(f)).

3.3. Prototype consistency for feature alignment

After synthetic lesions are blended into the brain images, the next step is training a segmentation model. Although the lesion generator is designed to produce synthetic lesions that share a similar latent distribution with real lesions, this alignment occurs in the generative space rather than the task-specific feature space learned by the segmentation model. Due to differences in spatial context, subtle intensity variations, or blending artifacts, synthetic lesions may still induce a feature shift during training. To address this, we introduce Prototype Consistency (PC) regularization to explicitly align the feature representations of real and synthetic lesions, promoting domain-invariant features. We hypothesize this will tend towards feature maps that are more general to the segmentation problem and less specific to features of the particular training dataset thereby increasing segmentation model robustness.

The feature map of real lesions is denoted as F_real = F ⋅ 𝟙 [M = 1], and the feature map of synthetic lesions is denoted F_syn = F ⋅ 𝟙 [M = 2]. Here F indicates the feature map of the composited image $\widehat{I}$, and 𝟙 (⋅) is an indicator function where the value is 1 if the condition is true, otherwise it is 0. Inspired by the prototypical network (Snell et al., 2017), we aim to force the segmentation model to learn similar feature distributions for F_real and F_syn via a class prototype. Specifically, we first obtain feature prototypes for both lesion types by averaging feature maps for the specific lesion type in the spatial dimension:

$${𝒫}_{real}=\frac{{\sum }_{x,y,z}{F}_{real}^{(x,y,z))}}{{\sum }_{x,y,z}𝟙[M^{(x,y,z)}=1]},$$ (9)

$${𝒫}_{syn}=\frac{{{\displaystyle \sum }}_{x,y,z}{F}_{syn}^{(x,y,z)}}{{{\displaystyle \sum }}_{x,y,z}𝟙\left[M^{(x,y,z)}=2\right]},$$ (10)

where (x, y, z) is the spatial coordinate. The loss of the prototype differences can then be computed as:

$${ℒ}_{pd}=\parallel {𝒫}_{real}-{𝒫}_{syn}{\parallel }_1,$$ (11)

where ‖ ⋅ ‖₁ is the L1-norm of a vector.

Only optimizing ℒ_pd neglects the intrinsic relation between class-specific features since it only minimizes the discrepancy between two class prototypes (Asadi, Davari, Mudur, Aljundi, & Belilovsky, 2023). To this end, we develop prototype relationship loss to maximize the consistency between relationship matrices constructed from the class prototype and class-specific features. Since the number of voxels in real and synthetic lesion areas can be different, we randomly sample k feature vectors within each class to obtain ℱ_real ∈ ℝ^n×c×k and ℱ_syn ∈ ℝ^n×c×k, where c is the number of feature channels. To measure consistency we compute cosine similarity:

$$\cos (𝒫,ℱ)=\frac{𝒫\cdot ℱ}{\Vert 𝒫{\Vert }_2\cdot \Vert ℱ{\Vert }_2},$$ (12)

where ‖ ⋅ ‖₂ denotes the L2-norm. The prototype relationship loss is computed as:

$$\begin{array}{c}{ℒ}_{prd}=\Vert \text{cos}\left({𝒫}_{real}^i,{ℱ}_{real}^i\right)-\text{cos}{\left({𝒫}_{real}^i,{ℱ}_{syn}^i\right)\Vert }_1\\ +\Vert \text{cos}\left({𝒫}_{syn}^i,{ℱ}_{real}^i\right)-\cos {\left({𝒫}_{syn}^i,{ℱ}_{syn}^i\right)\Vert }_1.\end{array}$$ (13)

The loss function for training the segmentation model is:

$$ℒ={ℒ}_{seg}+\underset{{ℒ}_{pc}}{\underbrace{{\lambda }_1{ℒ}_{pd}+{\lambda }_2{ℒ}_{prd}}},$$ (14)

where λ₁ and λ₂ are weight factors. ℒ_pc is the prototype consistency, comprised of a difference and relationship loss term, and ℒ_seg is the compound loss which comprises of the Dice and Cross-entropy loss functions.

4. Experiments

Our framework is evaluated on two public brain segmentation datasets: the ATLAS v2.0 dataset and the Shift MS segmentation dataset as described in Section 4.1. We compare our approach to existing methods (Section 4.3) to show the superiority in both lesion synthesis and segmentation tasks. We further conduct ablation studies to validate the effectiveness of each component in our framework (Section 4.4).

4.1. Dataset and preprocessing

4.1.1. ATLAS v2.0 dataset

The ATLAS v2.0 dataset (Liew et al., 2022) is a large stroke segmentation dataset that contains 655 T1-weighted brain images collected from 33 research cohorts. All images first have intensity standardization and are registered to the MNI-152 template (1 mm³ voxel spacing). A defacing step is applied to anonymize the scan. All lesion masks are annotated and then checked by two neurological experts. We randomly select 80% of the dataset as the training set, and keep the remaining 20% for model evaluation.

4.1.2. Shift MS dataset

The Shift MS dataset (Malinin et al., 2022) is a multi-center white matter multiple sclerosis segmentation dataset comprised of i.e., MSSEG-1 (Commowick et al., 2018), ISBI (Carass et al., 2017), PubMRI (Lesjak et al., 2018), and a private dataset collected from the University of Lausanne. We use the three public datasets for model training and evaluation since the private dataset is not publicly available. Detailed information on the dataset is shown in Table 1, the training set (Trn) is used for model training. The in-domain development set (Dev_in) is used for model validation. The out-domain development set (Dev_out) and in-domain evaluation set (Eva_in) serve as separate, unseen test sets used for evaluation of model generalization capability.

Table 1. Shift MS dataset details, including scanner location, type, magnetic field strength, and dataset split.

Dataset	Location	Scanner	Field	Trn	Devin	Devout	Evlin
MEESG-1	Rennes	S Verio	3.0T	8	2	0	5
	Bordeaux	GE Disk	3.0T	5	1	0	2
	Lyon	S Aera	1.5T	10	2	0	17
		P Ingenia	3.0T
ISBI	Best	P Medical	3.0T	10	2	0	9
PubMRI	Ljubljana	S Mag	3.0T	0	0	25	0

Each subject has two available modalities, FLAIR and T1 with contrast. In this work, we only use the FLAIR modality. All of the subjects have been preprocessed with image denoising, skull stripping, and bias field correction. We resample all images to 1 mm³ isotropic spacing. The ground-truth segmentation mask is determined by the consensus of annotations acquired from clinical experts.

4.2. Implementation details

The framework is implemented in PyTorch 1.13.1. All model training and validation were performed using an NVIDIA A100 40G GPU. For the lesion generator, the input mask and image size is 64 × 64 × 64. We used the AdamW optimizer to train the shape and the intensity models, with the learning rate set to 1e−5. The batch size was 4 and the total number of training epochs was 100. For the segmentation model, we use UNet (Ronneberger et al., 2015) as the backbone model. The input patch size is 128 × 128 × 128 and the batch size is 2. We used the AdamW optimizer for the segmentation model training with the learning rate set to 1e − 3 and consecutively reduced with a cosine annealing strategy. A total of 500 epochs were set for the ATLAS v2.0 dataset and 1000 epochs for the Shift MS dataset. The loss function coefficients are set to λ₁ = 1 and λ₂ = 50 for each dataset, respectively. These values were chosen empirically.

4.3. Model evaluation

We evaluated our framework on (1) its ability to generate lesions, (2) the performance of the segmentation model trained with images generated using our framework compared to other models and (3) comparing our framework with other data augmentation techniques for training the segmentation model.

4.3.1. Lesion synthesis performance

To evaluate the effectiveness of our generative model and self-supervised training strategy, we compare the synthetic lesions generated by our framework to other generative models. We use Peak Signal-to-Noise Ratio (PSNR) and Mean Absolute Error (MAE) to evaluate synthetic results. Structural Similarity (SSIM) is not suitable for this scenario because the large proportion of background dominates the small foreground area, leading to an inaccurate representation of the actual image quality. The measures are reported in Table 2. We compare our approach to both GAN and diffusion models. For each method, the model is trained either with only real lesions or only pre-generated image-mask pairs created using the self-supervised strategy (see Section 3.1.2). Across all models training with pre-generated data yields superior metrics compared to the real dataset, underscoring the ability of appropriate synthetic data to enhance model training. Additionally, our approach achieves the highest PSNR and lowest MAE, indicating its robustness and adaptability. Notably, the increased performance of our model is not attributable to the model’s size, as evidenced by the comparison with f-AnoGAN which has the largest number of parameters but the lowest quantitative performance. The performance gains of our method are attributed to its innovative architecture and the utilization of the self-supervised training strategy.

Table 2. Quantitative results of our method and other generative models for lesion synthesis.

Here ‘real’ indicates models are trained using only real images, and ‘syn’ indicates models are trained with the self-supervised strategy.

Methods	Training data type	#Param	ATLAS v2.0			Shifts MS
			PSNR↑	MAE↓		PSNR↑	MAE↓
AE	real	64.11M	31.72	0.0014		30.42	0.0017
Rumelhart, Hinton, and Williams (1986)	syn	64.11M	32.07	0.0011		31.56	0.0014
AAE	real	64.41M	30.66	0.0016		31.33	0.0014
Makhzani, Shlens, Jaitly, Goodfellow, and Frey (2015)	syn	64.41M	32.11	0.0010		32.84	0.0009
f-AnoGAN	real	78.15M	29.86	0.0019		28.35	0.0023
Schlegl et al. (2019)	syn	78.15M	30.22	0.0017		30.05	0.0018
DDPM	real	74.23M	32.53	0.0009		32.38	0.0009
Ho et al. (2020)	syn	74.23M	34.48	0.0007		34.23	0.0008
PNDM	real	74.23M	32.56	0.0009		32.32	0.0009
Liu et al. (2022)	syn	74.23M	34.87	0.0007		34.68	0.0007
Ours	real	67.08M	35.23	0.0006		34.38	0.0007
	syn	67.08M	37.42	0.0004		36.57	0.0005

4.3.2. Downstream segmentation model performance

We compare the segmentation model performance using our frame-work to existing segmentation models. All models were trained from scratch. Note that our framework can use any backbone model to perform the segmentation task, here we consider UNet as the base segmentation model. Segmentation performance for all models was evaluated using the Dice Similarity Coefficient (DSC), Average Surface Distance (ASD), and 95% Hausdorff Distance (HD95). Table 4 shows the segmentation model performance for the ATLAS v2.0 dataset. Our approach consistently demonstrates improved performance for all metrics compared to all competing models, irrespective of the base model. This consistent overperformance is attributed to two pivotal enhancements: the use of an augmented brain lesion dataset, which closely mimics real-world appearance for model training, and the incorporation of the prototype consistency loss to align model features for real and synthetic lesions. Our findings suggest that while architectural improvements enhance segmentation, the quality of the training dataset may have an even greater impact on model generalization. Our approach, by leveraging high-fidelity synthetic data and strategic feature alignment, effectively unleashes the potential of classical segmentation models like UNet, establishing that dataset augmentation and targeted modifications to the training loss function can improve segmentation models’ performances for brain stroke lesion segmentation.

Table 4. Quantitative results of our method and other segmentation models for the ATLAS v2.0 dataset.

* indicates that the p-value < 0.05 compared to the second-best model using a paired Student’s t-test.

Methods	DSC↑	ASD↓	HD95↓
UNet (Ronneberger et al., 2015)	50.36 ± 30.28	20.25 ± 25.46	39.42 ± 37.18
Attention UNet (Oktay et al., 2018)	52.64 ± 29.66	19.75 ± 23.70	43.65 ± 38.92
SwinUNETR (Tang et al., 2022)	53.19 ± 29.82	21.47 ± 51.44	37.33 ± 56.20
MedNeXt (Roy et al., 2023)	48.24 ± 32.79	21.67 ± 52.65	35.81 ± 54.92
nnUNet (Isensee et al., 2021)	56.25 ± 29.81	15.41 ± 21.74	28.67 ± 29.21
Ours	60.23 ± 29.48*	6.32 ± 13.68*	20.26 ± 25.81*

The quantitative metrics for model segmentation performance in the Shift MS Dataset are reported in Table 3. As with ATLAS v2.0, our framework has improved performance compared to the other models for the out-domain development set (Dev_out), which demonstrates our framework improves domain generalization ability.

Table 3. Quantitative results of our method compared to other segmentation models for the Shift MS dataset.

* indicates that the p-value < 0.05 compared to the second-best model performance computing with a paired Student’s t-test.

Methods	Shifts MS Devin				Shifts MS Devout				Shifts MS Evlin
	DSC↑	ASD↓	HD95↓		DSC↑	ASD↓	HD95↓		DSC↑	ASD↓	HD95↓
UNet (Ronneberger et al., 2015)	62.73 ± 21.04	7.71 ± 9.55	20.40 ± 18.91		61.58 ± 21.70	4.34 ± 8.57	14.52 ± 13.60		51.74 ± 20.29	14.01 ± 18.34	29.02 ± 23.83
Attention UNet (Oktay et al., 2018)	72.34 ± 13.24	2.46 ± 2.77	12.53 ± 16.88		64.80 ± 22.98	3.66 ± 8.41	12.83 ± 15.65		65.62 ± 15.21	3.34 ± 5.07	13.92 ± 16.64
SwinUNETR (Tang et al., 2022)	66.65 ± 14.34	3.54 ± 4.51	18.04 ± 17.42		61.39 ± 22.29	4.79 ± 8.24	15.27 ± 14.60		57.83 ± 13.85	5.38 ± 8.24	20.94 ± 20.77
MedNeXt (Roy et al., 2023)	69.21 ± 14.39	10.93 ± 12.96	53.96 ± 69.47		64.99 ± 22.99	2.72 ± 6.32	14.97 ± 17.89		59.20 ± 19.88	21.54 ± 37.39	54.39 ± 68.26
nnUNet (Isensee et al., 2021)	74.57 ± 14.86	2.14 ± 2.61	11.46 ± 15.74		64.75 ± 21.53	3.06 ± 6.23	13.84 ± 15.27		64.87 ± 16.49	8.21 ± 6.34	15.74 ± 18.63
Ours	78.35 ± 8.74*	0.91 ± 1.09	8.15 ± 12.06*		68.52 ± 15.60*	1.73 ± 3.24*	12.26 ± 19.37		69.51 ± 12.63*	1.60 ± 2.17*	7.30 ± 7.75*

Figs. 6 and 7 show qualitative segmentation results for the ATLAS v2.0 and Shift MS Dataset, respectively. These results demonstrate that our framework consistently provides more accurate predictions compared to other models. Notably, in regions indicated by the yellow arrows, other models have either false positives or fail to segment the foreground. The improvements in our framework are particularly evident in the segmentation of small lesions, which are often challenging for other segmentation models to recognize.

Fig. 6. Visualization of lesion segmentation in the ATLAS v2.0 dataset for different models.

Here ‘Att UNet’ is short for the Attention UNet.

Fig. 7. Visualization of lesion segmentation on the Shift MS dataset for different models.

Here ‘Att UNet’ is short for the Attention UNet.

4.3.3. Comparisons with different data augmentation methods

We compare the performance of our framework with other data augmentation methods, including voxel-based methods and GAN-based methods. Voxel-based data augmentation methods use a combination of two existing images to create a new synthetic image. Here we adopted six methods: i.e., Mixup (Zhang et al., 2018), CutMix (Yun et al., 2019), Copy–Paste (Ghiasi et al., 2021), TumorCP (Yang et al., 2021), CarveMix (Zhang et al., 2018), and SelfMix (Zhu et al., 2022). For the GAN-based methods (Chaitanya et al., 2021), a conditional GAN was adapted to generate either a deformation field (D), intensity field (I), or a combination of both (D+I) to change the structure and/or the appearance of images. Regardless of the data augmentation method, UNet was the segmentation model used.

Quantitative segmentation model performance on the ATLAS v2.0 dataset and Shift MS dataset are shown in Tables 5 and 6, respectively. For the ATLAS v2.0 dataset, our framework achieves the best segmentation metric, improving DSC by 3.65, ASD by 5.99 mm, and HD95 by 3.27 mm compared to the second-best model. A similar trend is seen for the Shift MS dataset. One potential reason for the improvements seen in our framework is that the comparison voxel-based data augmentation methods create unrealistic foreground areas which may shift the decision boundary of the segmentation model and reduce their ability to generalize. For the GAN-based data augmentation methods, generated deformation and intensity fields do not change the intrinsic characteristics of the original images, which results in models that may overfit the training data.

Table 5. UNet Segmentation model performance for our method and other data augmentation techniques on the ATLAS v2.0 dataset.

For the GAN method, we consider deformation augmentation (D), intensity augmentation (I), and both (D+I). * indicates that our framework significantly improves segmentation performance (p-value < 0.05) compared to the second-best model using a paired Student’s t-test.

Methods	Type	DSC↑	ASD↓	HD95↓
Mixup (Zhang et al., 2018)	Voxel	43.75 ± 35.41	23.45 ± 20.48	43.76 ± 32.54
CutMix (Yun et al., 2019)	Voxel	47.32 ± 32.46	22.71 ± 20.03	40.35 ± 31.56
Copy-Paste (Ghiasi et al., 2021)	Voxel	55.24 ± 31.42	15.57 ± 15.42	25.78 ± 27.44
TumorCP (Yang et al., 2021)	Voxel	55.64 ± 32.59	15.43 ± 15.08	24.86 ± 26.94
CarveMix (Zhang et al., 2018)	Voxel	56.58 ± 31.05	12.31 ± 16.54	23.53 ± 25.98
SelfMix (Zhu et al., 2022)	Voxel	54.13 ± 31.24	16.78 ± 17.55	25.97 ± 26.88
cGAN (D) (Chaitanya et al., 2021)	GAN	52.79 ± 32.23	20.43 ± 22.48	36.72 ± 30.23
cGAN (I) (Chaitanya et al., 2021)	GAN	50.66 ± 30.54	21.14 ± 23.84	39.46 ± 32.36
cGAN (D+I) (Chaitanya et al., 2021)	GAN	51.48 ± 31.26	22.58 ± 28.56	37.22 ± 31.28
Ours	Mixed	60.23 ± 29.48*	6.32 ± 13.68*	20.26 ± 25.81

Table 6. UNet Segmentation model performance for our framework and other data augmentation techniques on the Shift MS dataset.

For the GAN method, we consider deformation augmentation (D), intensity augmentation (I), and both (D+I). * indicates that our framework significantly improves segmentation performance (p-value < 0.05) compared to the second-best model using a paired Student’s t-test.

Methods	Type	Shifts MS Devin				Shifts MS Devout				Shifts MS Evlin
		DSC↑	ASD↓	HD95↓		DSC↑	ASD↓	HD95↓		DSC↑	ASD↓	HD95↓
Mixup (Zhang et al., 2018)	Voxel	52.64 ± 32.68	12.68 ± 13.45	26.75 ± 28.46		48.76 ± 23.45	13.54 ± 14.23	23.45 ± 24.35		38.45 ± 21.56	14.69 ± 15.46	38.58 ± 39.43
CutMix (Yun et al., 2019)	Voxel	54.23 ± 30.69	11.54 ± 12.53	24.54 ± 26.58		50.43 ± 22.59	12.59 ± 13.56	21.76 ± 22.69		40.23 ± 22.63	13.42 ± 12.53	36.44 ± 37.22
Copy-Paste (Ghiasi et al., 2021)	Voxel	60.36 ± 15.48	8.42 ± 9.64	22.43 ± 25.46		58.46 ± 19.53	11.23 ± 10.77	20.69 ± 21.34		43.56 ± 23.28	10.85 ± 11.49	25.12 ± 26.41
TumorCP (Yang et al., 2021)	Voxel	61.55 ± 13.55	6.46 ± 7.69	20.15 ± 27.33		62.44 ± 18.46	9.53 ± 6.75	18.43 ± 20.67		51.12 ± 20.51	9.34 ± 8.32	18.32 ± 15.46
CarveMix (Zhang et al., 2018)	Voxel	65.28 ± 12.47	5.24 ± 5.53	15.63 ± 18.42		63.86 ± 17.63	8.75 ± 7.53	15.23 ± 20.49		54.53 ± 18.24	6.43 ± 7.22	15.44 ± 12.75
SelfMix (Zhu et al., 2022)	Voxel	62.45 ± 14.62	7.53 ± 6.45	19.44 ± 20.47		60.54 ± 18.45	10.84 ± 7.53	16.29 ± 21.54		52.36 ± 19.42	7.52 ± 6.31	16.32 ± 11.42
cGAN (D) (Chaitanya et al., 2021)	GAN	62.87 ± 22.54	9.53 ± 7.65	19.78 ± 22.23		62.12 ± 20.68	6.52 ± 7.98	17.32 ± 21.34		52.65 ± 18.23	12.89 ± 10.25	28.36 ± 24.35
cGAN (I) (Chaitanya et al., 2021)	GAN	60.23 ± 25.65	10.54 ± 8.33	23.45 ± 24.25		60.98 ± 22.45	8.51 ± 8.21	18.55 ± 20.35		51.25 ± 18.78	13.45 ± 11.54	30.48 ± 26.35
cGAN (D+I) (Chaitanya et al., 2021)	GAN	61.27 ± 24.47	9.23 ± 8.11	22.36 ± 21.69		61.73 ± 21.59	6.78 ± 8.07	17.21 ± 19.24		52.73 ± 17.63	12.45 ± 10.48	26.45 ± 24.51
Ours	Mixed	78.35 ± 8.74*	0.97 ± 1.09*	8.15 ± 12.06*		68.52 ± 15.60*	1.73 ± 3.24*	12.26 ± 19.37		69.51 ± 12.63*	1.60 ± 2.17*	7.30 ± 7.75*

4.4. Ablation studies

We validate the effectiveness of the three modules in our framework i.e., Lesion synthesis (SL), Soft Poisson Blending (SPB), and Prototype Consistency (PC). All ablation studies were conducted on the ATLAS v2.0 dataset and Shift MS dataset.

4.4.1. Synthetic lesion

We validated the performance of the segmentation model when using only synthetic lesions for training. The results are shown in the first row of Tables 7 and 8 for the ATLAS v2.0 and Shift MS datasets, respectively. For the ATLAS v2.0 dataset, training with only synthetic lesions improves segmentation performance compared to using only the real data for model training. In contrast, segmentation performance on the Shift MS dataset is worse. The reason for this discrepancy we believe is dataset size, the ATLAS v2.0 dataset contains 655 images but the Shift MS dataset only has 33 images. While foreground mismatch and boundary artifacts caused by directly inserting the synthetic lesions can increase the model’s generalization, they are catastrophic for a small dataset like the MS Shift dataset since the synthetic lesions with boundary artifacts will shift the segmentation model feature distribution.

Table 7. Ablation study on the components of our framework: synthetic lesions (SL), Soft Poisson Blending (SPB), and prototype consistency loss (PC) for the ATLAS v2.0 dataset.

UNet is the segmentation model.

SL	SPB	PC	DSC↑	ASD↓	HD95↓
✓			54.86 ± 29.84	16.96 ± 24.59	34.98 ± 36.01
✓	✓		58.71 ± 27.20	11.30 ± 17.57	27.17 ± 30.28
✓		✓	59.38 ± 30.80	18.93 ± 72.62	30.00 ± 73.03
✓	✓	✓	60.23 ± 29.48	6.32 ± 13.68	20.26 ± 25.81

Table 8. Ablation study on the components of our framework: synthetic lesions (SL), Soft Poisson Blending (SPB), and prototype consistency loss (PC) for the Shift MS dataset.

UNet is the segmentation model.

SL	SPB	PC	Shifts MS Devin
			DSC↑	ASD↓	HD95↓
✓			59.54 ± 24.61	9.68 ± 12.64	22.36 ± 17.31
✓	✓		67.16 ± 14.26	5.98 ± 6.74	20.06 ± 16.99
✓		✓	70.05 ± 20.63	3.11 ± 4.39	13.49 ± 17.06
✓	✓	✓	78.35 ± 8.74	0.97 ± 1.09	8.15 ± 12.06
SL	SPB	PC	Shifts MS Devout
			DSC↑	ASD↓	HD95↓
✓			58.48 ± 25.06	6.89 ± 9.77	17.75 ± 15.71
✓	✓		59.23 ± 26.13	7.11 ± 10.79	16.36 ± 16.76
✓		✓	59.99 ± 24.63	3.30 ± 7.16	15.88 ± 17.47
✓	✓	✓	68.52 ± 15.60	1.73 ± 3.24	12.26 ± 19.37
SL	SPB	PC	Shifts MS Evlin
			DSC↑	ASD↓	HD95↓
✓			42.48 ± 23.76	18.42 ± 19.71	33.52 ± 23.88
✓	✓		53.89 ± 22.22	18.20 ± 28.73	41.92 ± 54.82
✓		✓	54.55 ± 25.35	6.05 ± 7.33	17.30 ± 15.15
✓	✓	✓	69.51 ± 12.63	1.60 ± 2.17	7.30 ± 7.75

4.4.2. Soft Poisson blending

Applying SPB to achieve a consistent appearance with real lesions and a seamless boundary improves the segmentation model performance for both datasets (the second row of Tables 7 and 8). Our results demonstrate that resolving the inconsistent appearance of synthetic lesions improves the model performance. We also observe variations in the standard deviation (STD) of ASD and HD95 across different configurations in Table 7. Notably, applying PC without SPB increases the variability of boundary-related metrics. This is likely due to subtle differences in appearance between real and synthetic lesions. Without SPB, synthetic lesions may contain unnatural boundaries, making it more difficult for the PC loss to guide feature alignment. Consequently, this can lead to inconsistent boundary predictions. In contrast, when SPB is used in conjunction with PC, the synthetic lesions exhibit smoother transitions and better visual coherence with real lesions. This allows the PC regularization to learn more consistent features, resulting in improved segmentation stability and significantly reduced STD, especially in boundary metrics which may be more sensitive to outliers.

4.4.3. Prototype consistency

To address the potential feature gap caused by synthetic and real images, we introduced prototype consistency regularization. This penalty, applied to both real and synthetic lesions, ensures the segmentation model learns similar features for lesions regardless of origin. Results shown in the third row of Tables 7 and 8 demonstrate applying prototype consistency regularization solely to synthetic lesions yields improved segmentation model performance. Moreover, integrating this regularization with a consistent lesion appearance further enhances segmentation performance, as evidenced in the fourth row of Table 7. The Shift MS dataset (Table 8) demonstrates a substantial improvement in segmentation performance compared to models where the consistency penalty was not employed highlighting that feature alignment is most important for small datasets where even a small shift in the synthetic lesion distribution can affect segmentation performance.

5. Conclusion and discussion

We presented a comprehensive framework to augment existing training samples for brain lesion segmentation via a two-stage adversarial autoencoder (AAE) to generate new lesion images. The AAE is trained in a self-supervised manner, but generates synthetic lesions with the same latent space distribution as real lesions. We then augment the synthetic images by using Soft Poisson Blending (SPB) to create a seamless boundary between foreground and background, eliminating boundary artifacts. Finally we introduce a prototype consistency regularization term during segmentation model training to ensure similar features across synthetic and real lesions are learnt. The synthetic lesion samples boost segmentation model performance under the supervision of the prototype consistency penalty. Experiments on two public datasets demonstrate that our framework outperforms other data augmentation approaches and methods that only adapt augmented samples for model training. We do not compare our approach to models based on pre-trained datasets such as SAM (Kirillov et al., 2023) because they are pre-trained on a large-scale dataset, making direct comparisons unfair. Besides, SAM-based methods largely depend on accurate user prompts to achieve good segmentation results, which differs from our fully automatic setting requiring no prompt.

Despite its strong performance on brain abnormality MRI datasets, our current framework has several limitations. First, it has only been validated on MRI data of brain lesions; its generalizability to other imaging modalities (e.g., CT, ultrasound) and lesion types remains to be explored. Second, although our framework implicitly learns lesion variability from real data distributions, it does not explicitly incorporate anatomical priors or region-aware controls. As a result, lesion synthesis may not fully capture the spatial dependencies between brain anatomy and lesion distribution. The current prototype consistency regularization does not model intra-class variation among lesions. While PC loss preserves inter-class structure, further improvements could be achieved by introducing multiple prototypes per class or cluster-aware constraints. In future work, we plan to extend our framework to diverse anatomical structures and modalities, and incorporate disease- or region-specific conditions to enhance the biological plausibility and controllability of lesion synthesis.