Self-supervised Brain Lesion Generation for Effective Data Augmentation of Medical Images

We follow the model architecture proposed in  [30] 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) [12] 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.

For training the AAE models we use a three-term loss: reconstruction loss $\mathcal{L}_{rec}$, latent adversarial loss $\mathcal{L}_{adv\_z}$, and image adversarial loss $\mathcal{L}_{adv\_x}$. $\mathcal{L}_{rec}$ computed as the mean absolute error (MAE) between the input image $\mathcal{I}$ and the reconstructed image $\mathcal{\hat{I}}$:

$\mathcal{L}_{rec}$ guarantees $\mathcal{\hat{I}}$ and $\mathcal{I}$ look similar in general.

The latent adversarial loss $\mathcal{L}_{adv\_z}$ is formulated to ensure that the latent space of the lesions has a normal distribution:

where $\mathcal{N}(0,1)$ is a normal distribution of mean $0$ and standard deviation $1$. Similarly, the image adversarial loss $\mathcal{L}_{adv\_x}$ is designed to ensure the reconstructed image is realistic in appearance. It is defined as:

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 (grey 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 [10], we first generate $n|n\sim\mathcal{U}(1,5)$ ellipsoids with overlap to simulate a general lesion shape. The half-axis lengths of three directions follow the uniform distribution $\mathcal{U}(5,15)$. Elastic deformations controlled by $\sigma|\sigma\sim\mathcal{U}(3,6)$ are applied to the ellipsoids to make the shape more natural and irregular. Finally, we add Perlin noise [28] 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.

To generate the 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 [11] 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.

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.

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 FastFurfer [8] 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.

We developed Soft Poisson Blending (SPB), which is a modified implementation of Poisson Image Editing [27]. 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 $\Omega$. The boundary of $\Omega$ is defined as $\partial\Omega$. $f$, the value function defined on $\Omega$, has a value of $f^{*}$ at $\partial\Omega$. We solve the optimization problem defined as:

where $V$ is the guidance vector field and $\nabla$ 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:

where $\Delta$ is the Laplacian operator, and $\operatorname{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 $(\nabla s,\nabla 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 $\nabla g$ on $\partial\Omega$ being much higher than $\nabla 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:

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

The ATLAS v2.0 dataset [19] 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 (1mm³ 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.

The Shift MS dataset [23] is a multi-center white matter multiple sclerosis segmentation dataset comprised of i.e., MSSEG-1 [5], ISBI[3], PubMRI [18] 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), the out-domain development set (Dev_out), and the in-domain evaluation set (Eva_in) are used for evaluating model segmentation performance only.

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 1mm³ isotropic spacing. The ground-truth segmentation mask is determined by the consensus of annotations acquired from clinical experts.

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 the 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.

We compare the segmentation model performance using our framework 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. These enhancements not only elevate the base model’s ability to accurately segment images but also emphasize a crucial insight: the quality of the dataset plays a more critical role than the network architecture in achieving model generalization for segmentation tasks. 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.

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.

Fig. 6 and Fig. 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.

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 [43], CutMix [42], Copy-Paste [6], TumorCP [41], CarveMix [43], and SelfMix [45]. For the GAN-based methods [4], 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 Table 5 and Table 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.

We validated the performance of the segmentation model when using only synthetic lesions for training. The results are shown in the first row of Table 7 and Table 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.

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 Table 7 and Table 8). Our results demonstrate that resolving the inconsistent appearance of synthetic lesions improves the model performance.

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.