Source-linked AI summary

Generative AI for Medical Imaging: extending the MONAI Framework

Walter H. L. Pinaya, Mark S. Graham, Eric Kerfoot, Petru-Daniel Tudosiu, Jessica Dafflon, Virginia Fernandez, Pedro Sanchez, Julia Wolleb, Pedro F. da Costa, Ashay Patel, Hyungjin Chung, Can Zhao, Wei Peng, Zelong Liu, Xueyan Mei, Oeslle Lucena, Jong Chul Ye, Sotirios A. Tsaftaris, Prerna Dogra, Andrew Feng, Marc Modat, Parashkev Nachev, Sebastien Ourselin, M. Jorge Cardoso

arXiv:2307.15208v1eess.IVcs.CV

TL;DR

Sharing healthcare data for medical-imaging research is constrained by privacy concerns, and generative-model implementation and evaluation can be complex. MONAI Generative Models provides reusable components and a standardized platform for developing, evaluating, and deploying models. Experiments demonstrate use across modalities, applications, and 2D and 3D scenarios, while the authors describe the platform as reducing development effort.

  • Problem

    Privacy concerns impede sharing healthcare datasets, while the complexity of generative models and varied evaluation metrics complicate their implementation.

    Source: §1 · pp. 2–3
  • Method

    The authors extend MONAI with an open-source platform of generative models and related components for standardized training, evaluation, and deployment.

    Source: §1 · p. 3, §4 · p. 12
  • Results

    Experiments demonstrate the platform across modalities, applications, and 2D and 3D scenarios, with reported high-quality image generation across evaluated scenarios.

    Source: §1 · p. 3, §3.1 · p. 7, §3.2 · p. 9
  • Takeaways & Limitations

    The authors state that the modular platform can reduce time and effort when applying generative models to image synthesis and other applications.

    Source: §4 · p. 12

Abstract

from arXiv · show

Recent advances in generative AI have brought incredible breakthroughs in several areas, including medical imaging. These generative models have tremendous potential not only to help safely share medical data via synthetic datasets but also to perform an array of diverse applications, such as anomaly detection, image-to-image translation, denoising, and MRI reconstruction. However, due to the complexity of these models, their implementation and reproducibility can be difficult. This complexity can hinder progress, act as a use barrier, and dissuade the comparison of new methods with existing works. In this study, we present MONAI Generative Models, a freely available open-source platform that allows researchers and developers to easily train, evaluate, and deploy generative models and related applications. Our platform reproduces state-of-art studies in a standardised way involving different architectures (such as diffusion models, autoregressive transformers, and GANs), and provides pre-trained models for the community. We have implemented these models in a generalisable fashion, illustrating that their results can be extended to 2D or 3D scenarios, including medical images with different modalities (like CT, MRI, and X-Ray data) and from different anatomical areas. Finally, we adopt a modular and extensible approach, ensuring long-term maintainability and the extension of current applications for future features.

1 Introduction

Privacy concerns impede sharing healthcare data needed for medical imaging research, while complex generative models and varied evaluation practices make implementation difficult. The paper introduces MONAI Generative Models, an open-source platform applied across modalities, anatomical areas, and 2D and 3D scenarios.

1 Introduction

  • Privacy concerns around sensitive healthcare data impede sharing and create a bottleneck for developing and deploying machine-learning methods in medical imaging research.The introduction frames access to shared healthcare datasets as crucial to advancing this research.Source: §1 · p. 2
  • Generative models learn data patterns and can generate plausible new samples, with diffusion models, autoregressive transformers, GANs, and VAEs among the approaches described.The paper frames model training as estimating the underlying data distribution.Source: §1 · p. 2
  • Synthetic data offers a privacy-preserving alternative to directly sharing patient data, and studies report similar or superior downstream performance when synthetic data supplements real data [2] [30].The introduction also cites studies reporting high-quality generated medical images [29] [37] [10].Source: §1 · p. 2
  • Generative models also support anomaly detection, image-to-image translation, image enhancement, and reconstruction from undersampled or noisy data.Examples include CT-to-MRI translation, denoising, artifact removal, and MRI or CT reconstruction for faster imaging.Source: §1 · pp. 2–3
  • MONAI Generative Models addresses implementation complexity and varied evaluation practices with an open-source platform demonstrated in five experiments across CT, MRI, and X-Ray, multiple body parts, and 2D and 3D settings.The experiments span applications including out-of-distribution detection, image translation, and super-resolution.Source: §1 · p. 3

2 Features

MONAI Generative Models packages reusable generative-model components alongside evaluation metrics and training losses. Its features cover diffusion, transformer, and GAN-based methods, with support for 2D and 3D data.

2 Features

  • The Python package and documentation bring together components for generative-model development, evaluation, and training.The feature set includes model components, metrics, and losses following MONAI conventions.Source: §2 · p. 3, §2.1 · p. 4, §2.2 · p. 6

2.1 Models

  • Diffusion support includes a timestep-conditioned UNet, conditioning layers, modular noise schedulers and profiles, latent compression models, an encoder, and ControlNets.AutoencoderKL and VQVAE provide interchangeable compression models for latent diffusion, which supports higher-dimensional 2D or 3D data than traditional diffusion models.Source: §2.1 · p. 4
  • For autoregressive transformers, the Ordering class converts image data into sequences using raster, S-curve, or random ordering, and VQ-VAE codebook indices can serve as tokens.The section also describes autoregressive transformers as modeling likelihood through conditionals over preceding input features.Source: §2.1 · pp. 4–5
  • GAN-related components include SPADE for spatially conditioned generation and patch-based discriminators, including a multi-scale version for assessing image realism at different resolutions.The patch-based discriminator also supports adversarial components in latent compression models.Source: §2.1 · p. 5

2.2 Metrics and Losses

  • The package includes FID and MMD for fidelity evaluation and MS-SSIM for assessing synthetic-data diversity across image-detail scales.The metrics are designed to integrate into existing MONAI workflows.Source: §2.2 · pp. 5–6
  • Training losses include spectral, patch-based adversarial, and perceptual losses, with 2D RadImageNet and 3D MedicalNet features and a 2.5D option for memory-constrained 3D use.The 2.5D approach applies available 2D networks to selected slices in three orientations.Source: §2.2 · p. 6

3 Experiments and Applications

The experiments demonstrate MONAI Generative Models across medical-image synthesis, modular latent models, out-of-distribution detection, image translation, and 3D super-resolution.

3 Experiments and Applications

  • The experiments section presents applications intended to demonstrate the platform’s capabilities, with code for all experiments available online.The cited passage supplies the repository address: https://github.com/Warvito/generative_monai.Source: §3 · p. 6

3.1 Experiment I: Adaptability to different types of medical image

  • Across the evaluated image scenarios, the LDMs generated high-quality images; the study reports autoencoder reconstruction MS-SSIM, sample FID, and pairwise-sample MS-SSIM.The experiments included 2D and 3D brain images, mammograms, chest X-rays, and retinal OCT; Table 1 presents model performance.Source: Table 1 · p. 7, §3.1 · p. 7, Figure 1 · p. 8

3.2 Experiment II: Modularity of latent models

  • The VQ-VAE reconstruction scored MS-SSIM 0.9689, comparable to the KL-regularized autoencoder; the study also presents transformer and diffusion model performance using the same compression model.The comparison pairs a VQ-VAE with an autoregressive transformer or an LDM.Source: §3.2 · p. 9, Table 2 · p. 10

3.3 Experiment III: Application to out-of-distribution detection

  • The OOD experiment uses VQ-VAE plus transformers to obtain likelihoods, training on BRATs as in-distribution data and evaluating other Medical Decathlon classes as OOD data.Table 3 reports detection performance as AUC scores.Source: §3.3 · p. 10, Table 3 · p. 10

3.4 Experiment IV: Application to Image Translation

  • The ControlNet translated 2D FLAIR slices to T1-weighted images with PSNR=26.2458 ± 1.0092, MAE=0.02632 ± 0.0036, and MS-SSIM=0.9526 ± 0.0111.The model was trained using the UK Biobank 2D model and evaluated against true T1-weighted images.Source: §3.4 · p. 10

3.5 Experiment V: Application to image super-resolution

  • FID of 0.0024 and MS-SSIM of 0.9141 indicate higher fidelity and similar diversity versus single-step generation, while cascaded generation takes about 13 minutes rather than 22 seconds.The timing comparison used an NVIDIA TITAN RTX and DDIMScheduler with 200 timesteps; the authors attribute the improvement to more network parameters at the same training compute.Source: §3.5 · pp. 11–12
  • PSNR = 29.8042 ± 0.4173, MAE=0.0181±0.0009, and MS-SSIM=0.9806±0.0017 were reported for super-resolving test images against their ground truth.The authors use these metrics to assess the model’s super-resolution performance.Source: §3.5 · p. 12

4 Conclusions

MONAI Generative Models is presented as a modular MONAI extension intended to reduce the effort of developing and deploying generative models across synthesis and other applications. Future work targets newer models, model comparison, and broader application support including MRI reconstruction.

4 Conclusions

  • The modular platform is intended to reduce the time and effort of applying generative models to image synthesis and diverse applications.The authors describe future support for newer models, model comparison, and applications such as MRI reconstruction.Source: §4 · pp. 12–13
  • Future work will incorporate more recent models and improve model comparison and support for applications such as MRI reconstruction.The authors state these developments are intended to contribute to practical implementations of medical generative models.Source: §4 · pp. 12–13

Appendix A - Additional Samples

The appendix presents random samples from latent diffusion models trained on multiple medical-image datasets, ControlNets trained on UK Biobank, and 3D low-resolution brain images before and after super-resolution.

Appendix A - Additional Samples

  • Random LDM samples are shown for CSAW-M mammograms and MIMIC-CXR images, generated with 200 DDIM steps.The figures identify the two training datasets but do not state a comparative outcome.Source: Figure 5 · p. 17
  • Random LDM samples are shown for retinal OCT, 2D UK Biobank, and 3D UK Biobank datasets, generated with 200 DDIM steps.Source: Figure 6 · p. 18
  • The appendix shows random samples from UK Biobank-trained ControlNets and 3D low-resolution UK Biobank LDM images alongside their super-resolved versions.Source: Figure 7 · p. 19, Figure 8 · p. 20

Appendix B - Hyperparameters and Implementation Details

The appendix reports training configurations for compression and diffusion models, text conditioning, a VQ-GAN and transformer, and a ControlNet.

Appendix B - Hyperparameters and Implementation Details

  • Experiment I compression models omit attention and use L1, adversarial, perceptual, and KL losses, with distinct 2D and 3D learning rates.The 2D models use adversarial weight 0.005, perceptual weight 0.002, and KL weight 1e-8; perceptual loss uses a 2.5D approach on 25% of slices.Source: Appendix B · p. 21, Table 4 · p. 21
  • Experiment I diffusion models use AdamW at learning rate 0.000025 and a scaled-linear 1000-step schedule with beta range [0.0015, 0.0205] and v prediction.Source: Appendix B · p. 21, Table 5 · p. 21
  • Diffusion models use text conditioning from Stable Diffusion 2.1, with prompts tailored to each dataset and 10% empty-string conditioning for classifier-free guidance.Performance samples for FID and MS-SSIM were generated unconditioned; MIMIC-CXR prompts used up to five random report sentences.Source: Appendix B · pp. 21–22
  • Experiment II configures a VQ-GAN with a 3-channel latent space and 512-entry codebook, alongside a transformer trained with cross-entropy loss.The transformer uses 513 tokens, a maximum sequence length of 4096, eight layers, and eight attention heads per layer.Source: Appendix B · p. 22
  • Experiment IV uses the Experiment I 2D brain LDM and trains a ControlNet with a convolutional conditioning network spanning 64, 128, 128, and 256 channels.Other Experiment IV details are stated to follow Experiment I; tables 6 and 7 list Experiment V hyperparameters.Source: Appendix B · p. 22, Table 6 · p. 23, Table 7 · p. 23
Loading 2307.15208v1…