[논문 리뷰] Generative AI for Medical Imaging: extending the MONAI Framework

Walter Hugo Lopez Pinaya, Mark S. Graham외 22명|arXiv (Cornell University)|2023. 07. 27.

Generative Adversarial Networks and Image Synthesis인용 수 39

한 줄 요약

본 논문은 MONAI Generative Models를 제시한다. 이는 MONAI를 확장하여 2D/3D 모달리티에서 확산, 트랜스포머, GAN 등 다양한 의학 영상 생성 모델을 학습, 평가, 배포할 수 있는 오픈 소스 플랫폼이며, 사전 학습 모델과 모듈형 구성요소를 포함한다.

ABSTRACT

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.

연구 동기 및 목표

• 프라이버시 보호 데이터 공유와 다양한 의료 영상 작업을 촉진하기 위해 생성형 AI의 활용을 고무한다.
• 의료 영상에서 생성 모델을 학습, 평가, 배치하기 위한 표준화되고 모듈식인 플랫폼을 제공한다.
• CT, MRI, X-ray 등 다양한 모달리티와 데이터셋 간의 적용 가능성(2D/3D)을 시연한다.
• 재현성과 최첨단 방법과의 비교를 간소화하기 위해 사전 학습된 모델과 구성 요소를 제공한다.

제안 방법

• 타임스텝 조건화와 조건 정보용 공간 변환기를 갖춘 DiffusionModelUNet를 구현한다.
• DDIMScheduler 및 PNDMScheduler와 다중 노이즈 프로필(선형, 스케일드 선형, 코사인)을 갖는 Scheduler 추상화를 도입한다.
• 잠재 공간 생성을 위한 AutoencoderKL 및 VQVAE를 사용한 Latent Diffusion Models를 지원한다.
• 조건화 및 적대적 학습을 위한 ControlNets 및 PatchDiscriminator 변형(PatchDiscriminator, MultiScalePatchDiscriminator)을 포함한다.
• 고차원 데이터를 1D 트랜스포머에 맞추는 Ordering 클래스와 자기회귀 트랜스포머를 제공한다.
• 이미지 합성을 위한 SPADE 유사 적응 정규화 통합 및 적대적 구성요소와 결합한다.
• 스펙트럴 손실, 패치 기반 적대 손실, 인지적 손실 등 손실 함수와 MONAI-사전학습 네트워크(2D RadImageNet, 3D MedicalNet)를 포함한다.
• 계산 자원을 균형 있게 사용하기 위한 2.5D 접근법을 통한 3D 지각 손실과 함께 FID, MMD, MS-SSIM 등의 평가 지표를 지원한다.

Figure 1: Synthetic images from Latent Diffusion Models. A) Mammogram, B) Chest X-Ray, C) Retinal OCT, D) 2D slice from T1-weighted brain image, and E) axial, coronal, and sagittal view of a 3D brain image.

실험 결과

연구 질문

• RQ1 MONAI Generative Models가 단일 MONAI 프레임워크 내에서 의학 영상에서 최첨단 생성 방법들(확산, 트랜스포머, GAN)을 재현하고 확장할 수 있는가?
• RQ2 이 모델들이 2D/3D 의학 데이터를 넘나들며 여러 모달리티(CT, MRI, X-ray)와 해부학적 영역에 걸쳐 일반화되는가?
• RQ3 조건 신호(텍스트 프롬프트, 의미적 레이아웃)가 생성된 영상의 품질 및 참조와의 정합성에 어떤 영향을 주는가?
• RQ4 플랫폼 내에서 압축(VQ-VAE, AutoencoderKL)과 생성 구성요소(확산, 트랜스포머)의 모듈성 및 교환 가능성은 어느 정도인가?
• RQ5 이 플랫폼이 의학 영상에서 Out-of-Distribution 탐지, 이미지 변환, 초해상도와 같은 다운스트림 작업을 지원할 수 있는가?

주요 결과

Dataset	Data Type	Dimensions	FID	MS-SSIM	MS-SSIM Recons.
MIMIC-CXR	2D Chest X-ray	512×512	8.8325	0.4276	0.9834
CSAW-M	2D Mammography	640×512	1.9061	0.5356	0.9789
Retinal OCT	2D OCT	512×512	2.2501	0.3593	0.8966
UK Biobank	2D Brain Slice MRI	160×224	2.1986	0.5368	0.9876
UK Biobank	3D Brain MRI	160×224×160	0.0051	0.9217	0.9820

• 다양한 데이터셋(MIMIC-CXR, CSAW-M, UK Biobank, retinal OCT)에 대해 학습된 잠재 확산 모델은 2D 및 3D 데이터 전반에서 고품질의 다양한 샘플을 생성한다.
• FID 점수는 데이터셋에 따라 낮음에서 보통 범위이며(예: 3D UK Biobank brain MRI 0.0051; CSAW-M 1.9061), MS-SSIM 재구성은 충실한 재구성을 나타낸다(예: 3D Brain MRI 0.9820).
• 텍스트 프롬프트를 통한 조건부 샘플링은 BiomedCLIP의 높은 CLIP 정합성을 달성하면서도 허용 가능한 FID를 유지하여 제어 가능한 생성을 보여준다.
• 모듈성: VQ-VAE + Transformer 및 VQ-VAE + Diffusion Model 변형이 MIMIC-CXR에서 경쟁력 있는 FID(9.1995 대 8.0457)를 보이며, 잠재 생성 백본의 교환 가능성을 보여준다.
• 시연에는 Out-of-Distribution 탐지(AUC가 여러 Decathlon 데이터세트에서 Transformer 기반 방법으로 1.0으로 보고)와 이미지 변환(FLAIR to T1-weighted MRI)으로 높은 재현성 지표(PSNR, MAE, MS-SSIM)를 포함한다.
• 연쇄 확산 접근법을 통한 이미지 초해상도는 3D 데이터에서 FID가 거의 0에 가깝고 PSNR 약 29.8, MS-SSIM 약 0.98의 강한 정합성과 다양성을 달성한다.

Figure 2: CLIP Score vs FID Pareto curve for the LDM trained on chest X-Ray data. We sweep over guidance values of [1, 1.5, 1.75, 2, 3, 4, 5, 6, 7, 8, 9, 10]