LeFusion: Controllable Pathology Synthesis via Lesion-Focused Diffusion Models

Inspired by diffusion-based inpainting (Lugmayr et al., 2022; Avrahami et al., 2022), we aim to decouple the generation of lesions and background. As shown in Fig. 2, inpainting predicts the missing parts of an image, particularly lesions. For standard diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020), the inference process starts by sampling a noise vector $x_{T}\sim\mathcal{N}(0,1)$ and gradually denoising it to produce the output image $x_{0}$. We focus on 3D images only in this study, while the method can be easily extended to 2D. The original grayscale image is denoted as $\hat{x}_{0}\in\mathbb{R}^{D\times H\times W\times 1}$, where $D$, $H$, and $W$ represent the 3D image size. The reverse diffusion step from $x_{t}$ to $x_{t-1}$ is decoupled into lesion and background components, as shown below. Here, $M_{f}$ and $M_{b}$ represent the lesion foreground mask and background mask. The lesion $o_{t-1}$ is predicted through the reverse diffusion process $p_{\theta}$ using a 3D U-Net, while $\hat{x}_{t-1}$ is derived from $\hat{x}_{0}$ by adding noise through forward diffusion $q$.

$${x_{t-1}}=o_{t-1}\odot M_{f}+\hat{x}_{t-1}\odot M_{b}\\ ,o_{t-1}\sim p_{\theta}\left(\hat{x}_{t},t\right),\hat{x}_{t-1}\sim q\left(\hat{x}_{0},t\right)\\ .$$

(1)

This approach preserves the background accurately without requiring prediction.

The method above, while suitable for all diffusion-based inpainting, cannot guarantee that the denoised output $o_{t-1}$ is focused on the lesion area. Unlike standard diffusion training, where the target region may vary, we know that the part to be inpainted is specifically the lesion in our application. Therefore, we can design the model to predict only the lesion and ignore other regions. To achieve this, we introduce a lesion-focused loss during training. The general diffusion process starts with the original image $\hat{x}_{0}$, adding Gaussian noise over $T$ time steps (forward diffusion). The neural network is trained to predict the noise distribution at time step $t$ (reverse diffusion), conditioned on the noised image $\hat{x}_{t}$ and the time step. To ensure the model focuses only on the lesion, a mask $M_{f}$ is applied, calculating the loss exclusively within the lesion region. The training objective is defined as follows, where $\bm{\varepsilon}\in\mathbb{R}^{D\times H\times W\times 1}$ represents the noise sampled from a Gaussian distribution:

$$\mathbb{E}_{\hat{x}_{0},\epsilon\sim\mathcal{N}(0,1),t}\left[M_{f}\|\epsilon-p_{\theta}\left(\hat{x}_{t},t\right)\|_{2}\right]\\ .$$

(2)

Despite the change in the training objective, inpainting inference remains unaffected. As shown in Eq. 1, outside $M_{f}$, the predicted lesion $o_{t-1}$ is replaced by the real noised background $\hat{x}_{t-1}$.

In the section above, the model relies on the noised background $\hat{x}_{t}$ to infer the texture of lesion within the mask. While this works for lesions with minimal texture differences (as with cardiac lesions), it becomes problematic with multi-peak data. As shown in Fig. 3, lung nodules exhibit distinct texture types. We empirically show that relying solely on the background can lead to mode collapse.

To address this, we propose a simple yet effective approach, histogram-based texture control. The lesion texture histogram $h$ is used as a condition via cross attention Rombach et al. (2022), i.e.,

$$o_{t-1}\sim p_{\theta}\left(\hat{x}_{t},h,t\right)\\ .$$

(3)

During training, the histogram is computed from the ground truth, and during inference, texture types can be controlled by adjusting the histogram. Notably, this approach requires no additional lesion type annotations, such as nodule attenuation.

In Sec. 4, we demonstrate that the proposed histogram-based texture control is crucial for generating lesions on normal scans. Without it, models tend to fail, biasing towards healthy appearances and producing overly subtle lesions, which degrades the performance of downstream segmentation tasks using these synthetic data.

The method above focuses on single-class lesions, but many medical applications, such as cardiac MRI, require modeling multiple lesion types, like myocardial infarction (MI) and persistent microvascular obstruction (PMO). To capture correlations between lesion types and generate textures for multiple lesions simultaneously, we use a joint modeling strategy called multi-channel decomposition, where each channel corresponds to a different lesion type. The diffusion model generates each lesion in its respective channel, and they are combined through lesion masks.

We expand the input image $\hat{x}_{0}\in\mathbb{R}^{D\times H\times W\times 1}$ to $\hat{x}_{t}\in\mathbb{R}^{D\times H\times W\times n}$, based on the number of lesion classes $n$, where $n=2$ in the cardiac MRI experiments. Similarly, $\bm{\varepsilon}\in\mathbb{R}^{D\times H\times W\times n}$ and $o_{t-1}\sim p_{\theta}\left(\hat{x}_{t},t\right)\in\mathbb{R}^{D\times H\times W\times n}$ are extended to $n$ channels. The training objective is:

$$\mathbb{E}_{\hat{x}_{0},\epsilon\sim\mathcal{N}(0,1),t}\sum_{i=1}^{n}\left[M_{f}^{(i)}\left\|\epsilon-p_{\theta}\left(\hat{x}_{t},t\right)\right\|_{2}^{(i)}\right]\\ ,$$

(4)

where $*^{(i)}$ refers to the channel $i$ of $*$. To combine these channels, we compute $\sum_{i=1}^{n}\left[M_{f}^{(i)}o_{t-1}^{(i)}\right]$.

LIDC: Multi-Peak Lung Nodule CT.

We use LIDC dataset (Armato III et al., 2011), which contains 1,010 chest CT scans, from which 2,624 regions of pathological (P) interest (ROIs) corresponding to lung nodules were extracted, along with 135 cases of healthy patients. The dataset was divided into an 808-case training set, comprising 2,104 lung nodule ROIs, and a 202-case test set, containing 520 lung nodule ROIs. Additionally, 3,076 normal ( N) ROIs were cropped from the 135 healthy patients, representing regions where lung nodules typically appear. These normal ROIs were used for data augmentation in the experiments.

Emidec: Multi-Class Cardiac Lesion MRI. The Emidec dataset (Lalande et al., 2022) consists of examinations featuring DE-MRI in a short-axis orientation. This dataset offers access to 100 labeled cases, including 33 normal (N) and 67 pathological (P). The annotations cover 5 classes: background, left ventricle (LV), myocardium (Myo), myocardial infarction (MI), and persistent microvascular obstruction (PMO). We split the 67 P cases into 57 for training and 10 for testing. The 57 P cases are used to train the data synthesis model. In the downstream evaluation (Sec. 4.2), we use those models to synthesize P cases based on both 57 P and 33 N as the training set.

The following synthesis algorithms are compared with the LeFusion.

Copy-Paste. We used the masks from real lesion data and matched them with normal data, copying the original lesion textures onto normal cases to generate new synthetic data.

Hand-Crafted (Hu et al., 2023). The lesion mask is represented by the overlapping of multiple ellipsoidal lesion masks, followed by several random morphological operations. The texture is approximated using Gaussian noise and softened through interpolation and Gaussian filtering.

RePaint (Lugmayr et al., 2022) or Blended-Diffusion Avrahami et al. (2022). This inference-stage method removes the lesion mask from the image and then fills in the lesion texture by combining forward-diffused backgrounds with reverse-diffused foregrounds. The model employs global training loss, which lacks the capability to focus on lesion information. Due to the absence of guidance from lesion category information, it is unable to specify corresponding lesions and is confined to simulating the generation of single-class lesions;

Cond-Diffusion (Ho et al., 2020; Rombach et al., 2022).

These methods use the lesion mask and background image information as conditional inputs (Rombach et al., 2022) to a diffusion model (Ho et al., 2020). However, a downside of this approach is that it disrupts the background information. Furthermore, directly using multiple masks as conditional inputs fails to control the corresponding categories, limiting the method to modeling the generation of single-class lesions.

Cond-Diffusion (L) (Chen et al., 2024). Cond-Diffusion (L) is conceptually a latent diffusion (Rombach et al., 2022) version of Cond-Diffusion but adds VQGAN (Esser et al., 2021) to map image into latent space for diffusion. For a fair comparison, we fine-tuned open-source code and pre-trained weights by Chen et al. (2024) and used the model outputs directly.

LeFusion and the Variants (Ours). Apart from standard LeFusion, there are two variants for fine control of lesion textures. Histogram-Based Texture Control (*-H): A variant of LeFusion that incorporates histogram control information, using the input histogram to guide the generation of multi-peak lesion textures. Multi-Channel Decomposition (*-J): When handling multi-class lesions, standard LeFusion trains individual models separately, lacking of modeling for the correlation between different lesions. LeFusion-J is a generalized version to model multi-class lesions jointly.

We show that LeFusion can effectively benefit downstream application of training nnUNet (Isensee et al., 2021) and SwinUNETR (Hatamizadeh et al., 2021) to perform lung nodule segmentation. We use the following synthetic subset settings: P’: $2104\times 1$ synthetic ROIs from the 808 real pathological cases P; N’: $3076\times 1$ synthetic samples from the 135 normal subjects N; N”: $3076\times 2$ synthetic cases from the 135 normal subjects N.

Tab. 1 show the Dice and normalized surface distance (NSD). For the first group (Texture Synthesis with Real Masks), the texture of Hand-Crafted (Hu et al., 2023) and RePaint (Lugmayr et al., 2022) differs significantly from the real texture, making it challenging to achieve satisfactory results. Cond-Diffusion (Ho et al., 2020; Rombach et al., 2022) and Cond-Diffusion (L) (Chen et al., 2024), on the other hand, disrupt the background structure of the generated images. Our baseline model, LeFusion, is impacted by the pixel distribution of the background due to the lack of histogram control information, resulting in only a slight improvement over the baseline. We also synthesized lesion data on normal data N using Hand-Crafted Synthetic Masks and masks matched from lesion data, arriving at similar conclusions. In the final set of experiments, we validated the effectiveness of the DiffMask we designed for lesion synthesis and further explored accuracy improvements with increasing amounts of synthetic data. Compared to the baseline, in terms of Dice, we achieved improvements of 5.18% and 4.75% for nnUNet (Isensee et al., 2021) and SwinUNETR (Hatamizadeh et al., 2021), respectively. NSD improvements were 4.4% and 4.53%, respectively.

We use the following synthetic subset settings :

P’: $57\times 1$ synthetic cases from the 57 real pathological cases P; N’: $33\times 2$ synthetic cases from the 33 normal cases N; N”: $33\times 5$ synthetic cases from the 33 normal cases N.

We use combinations of these subsets to train nnUNet (Isensee et al., 2021) and SwinUNETR (Hatamizadeh et al., 2021). The results of the two types of lesions are reported in 2 in terms of the Dice (0-100, higher is better). Due to page limitations, the corresponding NSD table is provided in Appendix Tab. A1.

In the first group of Tab. 2, we apply texture synthesis with real masks. The Hand-Crafted (Hu et al., 2023) produces textures that differ from real textures, leading to a decrease in baseline performance. Cond-Diffusion (Ho et al., 2020; Rombach et al., 2022) and Cond-Diffusion (L) (Chen et al., 2024) disrupts background structure, blurring lesion categories, decreasing MI’s performance. RePaint (Lugmayr et al., 2022), focusing on global information, struggles to generate textures that conform to lesion characteristics, resulting in a significant decrease in the Dice for PMO lesions. LeFusion models the two lesions separately, ignoring the correlation between lesions, which leads to improved accuracy for MI but decreased accuracy for PMO lesions. In contrast, our proposed LeFusion-J achieves superior results; For the second group, we expanded the normal data utilizing texture synthesis with hand-crafted synthetic masks. Due to RePaint (Lugmayr et al., 2022)’s inability to distinguish between multiple lesion categories, we did not repeat experiments for it. From the experiments, we observed results similar to those of the first group; For the last, we evaluated our proposed lesion mask synthesis. Our method significantly improved the performance for both MI and PMO. Additionally, as data volume expanded, segmentation Dice consistently improved.

Fig. 5 shows the generation results of different methods for lung nodule CT and cardiac MRI based on lesion images and normal images. We have also quantitatively calculated and compared the similarity between our generated images and real images; more details can be found in Appendix D. Fig. 5(a) displays the synthesized visualization of lung nodules (red). The Cond-Diffusion method (Ho et al., 2020; Rombach et al., 2022) and Cond-Diffusion (L) (Chen et al., 2024) disrupt the background structure. The lesions generated by Hand-Crafted (Hu et al., 2023) and RePaint (Lugmayr et al., 2022) fail to capture texture information, such as the grayscale variations characteristic of the lesions. Our baseline LeFusion, without histogram control information, is easily influenced by background features, resulting in the generation of relatively shallow lesions. In contrast, our LeFusion-H can better utilize histogram information to control the generation of lesions. Fig. 5(b) displays the synthesized visualization of two lesions for the heart: MI (blue) and PMO (red). The generation results for the heart show similar conditions. It is important to note that since the heart has two types of lesions, LeFusion-J (joint modeling of both lesions) can more accurately reflect the textures of both lesion types and, compared to LeFusion, results in smoother transitions at the boundaries between the two lesions and the background. However, due to the small contrast difference between heart lesions, LeFusion-J-H with histogram control achieves similar effects. More visualizations can be found in the supplementary materials.

We studied the effect of histogram control on the lung nodule dataset. The visualization results are shown in Fig. 6. Under different histogram controls, the attenuation (“transparency”) of the generated lesions changes from shallow to deep. Without histogram control information, the generated lung nodules tend to match the pixel distribution of the normal lung background, resulting in overly light lesions. We selected 100 subsets and used LeFusion-H and LeFusion to generate each sample twice, calculating the similarity between pairs using Peak Signal-to-Noise Ratio (PSNR) and Structural Similarity Index Measure (SSIM), lower means more diverse. As shown in the figure, with histogram control, there is greater diversity between samples.

We also conducted a quantitative analysis of the impact of histogram control on the lesion area. The addition of the control histogram influences the distribution of the original data’s histogram, effectively shifting or weighting it based on the control input. Empirically, this interpretation aligns well with the observed distributions. Further details can be found in Appendix C.

We also visualize synthetic lesion masks, as shown in Fig. A1 and Fig. A2. Compared to the hand-crafted masks (Hu et al., 2023), our diffusion-generated masks are closer to the real masks and exhibit a more diverse range of shape patterns. More details can be found in Appendix A.