A Generative Model of Lung CT Conditioned on Radiomics Features

Abstract

Deep learning image generation has been an active area of research in a number of applications. However, traditional generative models are not able to control specific properties of the image outputs. In this work, we propose a deep learning model that produces images according to user-specified texture feature values. We adopted a diffusion transformer architecture and used texture features to condition the reverse process. The model was trained on lung patches from a public lung CT database. Two texture features, autocorrelation and inverse difference derived from the Gray-Level Co-Occurrence Matrix were used as conditional inputs. We evaluated the ability of the model to produce samples with similar feature values as the conditional inputs. Both in-distribution and out-of-distribution conditions were evaluated. Results indicate that the model is able to generate image patches resembling lung parenchyma. The autocorrelation and inverse difference of generated images have good agreement with and exhibit low variability around the conditional inputs. The concordance correlation coefficient between real and generated samples is 0.9962 for autocorrelation and 0.9402 for inverse difference. Visual assessment of image samples reveals that real and generated images share similar features, consistent with their radiomic properties. Findings from this work indicate that the diffusion transformer model is able to generate images with texture features closely aligning with the conditional inputs, supporting its utility for highly controlled data generation for a variety of applications.

1. INTRODUCTION

Image generation has been an active research area in medical imaging. Generated images serve a variety of purposes, including image quality evaluation through virtual clinical trials,^1,2 as anatomically realistic priors for image reconstruction or restoration,^3,4 as well as augmentation of small and/or imbalanced datasets to improve the performance of downstream machine/deep learning models.^5,6 Image generation has traditionally relied on analytical approaches.,^7,8 Recently, deep learning generative models have gained popularity due to their ability to learn high-dimensional distributions and the potential to produce highly realistic-looking images.^9–11

One drawback of deep learning generative models is the lack of control of the properties of image outputs. Frequently, image data from the same anatomical site are treated as samples from the same distribution and used to train a generative model. Such data, however, can cover broad distributions of image features (e.g., textures). A randomly generated sample typically has image features that are randomly sampled from their distributions in the training data. For some applications, it may be desirable to generate outputs with specific properties, for example, to produce a balanced dataset for algorithm training.^12,13

In previous work, Wang et al.¹⁴ developed a conditional generative diffusion model for trabecular bone images conditioned on structural properties including bone volume fraction, trabecular thickness, and spacing. While the method achieved excellent performance, it is unclear if similar concepts can be applied to generate non-binary anatomical images conditioned on higher-order statistical features like textures. In this work, we extend this approach to generate lung patches conditioned on features derived from the gray-level co-occurrence matrix (GLCM).

2. METHODS

2.1. Diffusion Transformer Model

The generative model adopts a diffusion transformer (DiT) architecture¹⁵ shown in Fig.1 where the typical U-Net used in diffusion models¹⁶ is replaced with a vision transformer.¹⁷ The model consists of 12 transformer blocks, each with 6 attention heads with a hidden dimension of 384. Distinct from the architectures in Refs.15 and 14 where image inputs are first converted to low-dimensional latent representations via an autoencoder, we directly provided image data as inputs to the diffusion model. Radiomics features values, r, and time step, t, are provided to each block as conditional inputs. First, r and t are each encoded via a linear layer into 1 × 384 vectors and summed. The resulting vector, z_cond, is used to regress the dimension-wise scaling and shift parameters γ and β in two adaptive layer norm (adaLN) layers using a multi-layer preceptron (MLP).

Figure 1:

Architecture of the diffusion transformer model used to conditionally generate lung patches with user-specific radiomics features.

We optimized the network architecture in terms of: (1) latent representation vs. direct image input, (2) size of the transformer, and (3) patch size.

2.2. Model Training

We extracted lung patches from thoracic CT scans in the LIDC¹⁸ database for training. A deep learning segmentation¹⁹ was first performed to extract the lung region. According to annotations provided in the database, we excluded volumes-of-interest containing lung nodules to focus only on nodule-free lung regions. Axial regions-of-interest (ROIs) of 64×64 pixels were randomly sampled from each scan. A total of 20,000 patches was extracted for training and 2500 patches for testing. All voxel values were clipped between [−1000, 500] Hounsfield units (HU) and normazlied to [−1, 1] before inputted to the model. For each patch, we computed the texture features autocorrelation and inverse difference (ID) from its GLCM and used the values as conditions during training. For GLCM computation, the images were discretized into 64 gray levels between −1000 and 500 HU. Radiomics definitions follow those in PyRadiomics.²⁰ The two features were chosen based on their low correlation in order to exercise the model’s ability to simultaneously generate independent texture features. The AdamW optimizer was used with a learning rate of 1e-4. The batch size was set to 8. The model was found to converge after 75 epochs.

2.3. Evaluation

We evaluated the ability of the model to generate lung patches with user-specified radiomics feature values. The aurocorrelation and ID of the 2500 testing patches were computed and used to condition the reverse sampling process. The corresponding radiomics features of the generated images were then computed. The Pearson correlation coefficient (r) and concordance correlation coefficients (ρ_c) between generated feature values and conditional inputs were calculated for each feature to assess correlation and agreement, respectively. To avoid bias due to the innate distributions of the feature values and to evaluate the spread of generated values around the conditional inputs, we also evenly sampled the 2D space of autocorrelation and ID throughout their respective range and produced 25 samples at each sample point. The coefficient of variation (CV) for all sample points are reported for each radiomics feature.

3. RESULTS

Similar to the architecture optimization performed in Ref.,¹⁵ we experimented with the size of the transformer and patch sizes. For our imaging task, the smallest transformer with a patch size 2 achieved comparable performance with larger networks and was therefore selected. Contrary to previous works^14,15 where a variational autoencoder was adopted to reduce the dimensions of the input patches, we found that this approach yielded worse performance (in terms of the correlation between radiomics features of the outputs and the conditions) compared to using the input patches directly.

The distributions of autocorrelation and ID of the 2500 images in test dataset (blue) and corresponding generated outputs (orange) are plotted in Fig.2(a). The distributions have significant overlap and similar coverage, indicating the model’s ability to produce samples that span the range of feature values in the testing data. The generated feature values vs. conditional inputs for each radiomics feature is plotted in Fig.2(b) and 2(c) with a linear fit, identity line, Pearson correlation coefficent (r), and concordance correlation coefficient (ρ_c) overlaid. For a network with perfect performance, the linear fit should overlap with the identity line, the Pearson and concordance correlations coefficients should both be 1. For both features, strong correlation and agreement were observed. The linear fit of autocorrelation almost completely overlaps with the identity line; that of the ID has similar slope but has a slight positive bias.

Figure 2:

(a) Distributions of conditional inputs (blue) and generated feature values (orange) show high degrees of overlap, indicating that the model is capable of producing samples that cover the full range of feature values. (b) Generated autocorrelation and (c) inverse difference are plotted vs the conditional inputs. The Pearson correlation coefficient (r) indicates high correlation for both features while the concordance correlation coefficient (ρ_c) indicates good agreement. A slight positive bias is observed for inverse difference.

The spread of generated values (dots) around their respective conditional inputs (round markers with black edges) are shown in Fig.3. The normalized distance between the conditional inputs to features within each training sample is plotted as the gray-scale background. Conditional inputs that are “in-distribution” have smaller distances (darker gray regions), while those that are “out-of-distribution” have larger distances (lighter gray regions). For in-distribution conditional inputs, feature values from generated samples exhibit small variability around the conditions with a slight positive bias observed for ID, consistent with the offset seen in Fig.2(c). The model demonstrates some ability to extrapolate for conditional inputs that are less represented in the training data, especially for samples with small autocorrelation. However, the generated clusters become increasingly dispersed at high autocorrelation values where the conditional inputs deviate further from those seen during training. The CV for autocorrelation over all sampled conditional inputs range from 2.752% to 60.273% with a mean of 5.186%; those for ID range from 1.371% to 12.509% with a mean of 3.666%.

Figure 3:

We evenly sampled conditional inputs throughout the autocorrelation and ID range (large, round markers with black edges) and generated 25 samples at each condition. The feature values from the generated samples are plotted as small, dotted markers that are observed to cluster around their corresponding conditional inputs (indicated by the same color). The normalized distance to features from all training data are plotted in the gray-scale overlay to visualize performance for data that in-distribution (smaller distance, darker gray values) vs out-of-distribution (larger distance, lighter gray values).

Example images are shown in Fig.4. Five real lung ROIs at different autocorrelation and ID values are shown in the left column; their respective feature values were used to conditionally generate four ROIs displayed to the right. Feature values of all images are shown on top. As inverse difference decreases (top three rows), both real and generated samples exhibit less local homogeneity seen from the decreasing smoothness in the background, less long-range structures, and more local variations. As autocorrelation increases (bottom three rows), we observe an increasing amount of vessel-like structures with locally uniform regions. The feature values of all generated samples are close to the conditions obtained from the real images.

Figure 4:

Example real lung patches and four samples generated while conditioned on their correlation and ID values (displayed above each figure). The generated samples display textures comparable to the real images and feature values that closely align with those in the real data.

4. CONCLUSION

We presented a diffusion transformer neural network capable of generating lung images with user-specified radiomics features. The feature values of the generated images closely align with the conditional inputs and show low variability. In ongoing work, we will extend the model to generate 3D lung parenchyma and include a wider range of radiomics features that capture a range of conditions from normal to diseased.