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

TL;DR

The paper tackles data scarcity in brain lesion segmentation by introducing a three-stage framework that combines self-supervised lesion synthesis, Soft Poisson Blending for seamless lesion insertion, and prototype-consistency regularization to align real and synthetic lesion features. It leverages a two-stage adversarial autoencoder to generate plausible lesion masks and textures, uses SPB to produce realistic composites, and enforces feature alignment with prototype losses during segmentation training. Across ATLAS v2.0 and Shift MS datasets, the approach achieves state-of-the-art segmentation performance and excels on small datasets where conventional augmentation struggles, including a Dice improvement from 50.36% to 60.23% on ATLAS v2.0. The work demonstrates that high-fidelity synthetic data, when properly blended and aligned in feature space, can substantially boost brain lesion segmentation and generalization, with potential for extension to other modalities and organs.

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 lesion generator, based on an adversarial autoencoder, in a self-supervised manner. 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 U-Net with conventional data augmentation techniques for the ATLAS v2.0 dataset.

Figures (7)

• Figure 1: (a) Real lesions (green arrow) and synthetic lesions (orange arrow) in images from the ATLAS v 2.0 (top) and MS Shift dataset (bottom) demonstrating the synthetic lesions have a similar appearance to real lesions. (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.
• Figure 2: Overview of our framework containing three stages. First, the lesion generator is trained via a self-supervised learning strategy and used to generate synthetic lesions based on constrained latent space sampling in Stage I. In Stage II, we seamlessly compose synthetic lesions into full brain images using the proposed Soft Poisson Blending (SPB) to increase the number of training samples. In Stage III, we train the downstream segmentation model with the prototype consistency regularization to align real and synthetic features.
• Figure 3: The detailed structure of MEB block.
• Figure 4: Lesion mask in three views for (a) real masks and (b) pred-generated masks for the shape AAE model training. Lesion images in three views for (c) real images and (d) pred-generated images for the intensity AAE model training.
• Figure 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.