[논문 리뷰] LeFusion: Controllable Pathology Synthesis via Lesion-Focused Diffusion Models

Hantao Zhang, Liu, Yuhe외 5명|arXiv (Cornell University)|2024. 03. 21.

Advanced MRI Techniques and Applications인용 수 4

한 줄 요약

LeFusion는 병변 영역에 집중된 확산 모델을 도입하여 심장 MRI 병변 합성 성능을 향상시키며, 정상 해부학의 구조적 일관성을 유지하기 위해 전방 확산 배경 정보 통합을 통해 배경의 무결성을 보존한다. 이는 DE-MRI에서 심근경색 및 노리플로우 병변의 nnUNet 분할 정확도를 햖을 수 있는 최신 기술 수준의 합성 데이터 증강을 달성한다.

ABSTRACT

Patient data from real-world clinical practice often suffers from data scarcity and long-tail imbalances, leading to biased outcomes or algorithmic unfairness. This study addresses these challenges by generating lesion-containing image-segmentation pairs from lesion-free images. Previous efforts in medical imaging synthesis have struggled with separating lesion information from background, resulting in low-quality backgrounds and limited control over the synthetic output. Inspired by diffusion-based image inpainting, we propose LeFusion, a lesion-focused diffusion model. By redesigning the diffusion learning objectives to focus on lesion areas, we simplify the learning process and improve control over the output while preserving high-fidelity backgrounds by integrating forward-diffused background contexts into the reverse diffusion process. Additionally, we tackle two major challenges in lesion texture synthesis: 1) multi-peak and 2) multi-class lesions. We introduce two effective strategies: histogram-based texture control and multi-channel decomposition, enabling the controlled generation of high-quality lesions in difficult scenarios. Furthermore, we incorporate lesion mask diffusion, allowing control over lesion size, location, and boundary, thus increasing lesion diversity. Validated on 3D cardiac lesion MRI and lung nodule CT datasets, LeFusion-generated data significantly improves the performance of state-of-the-art segmentation models, including nnUNet and SwinUNETR. Code and model are available at https://github.com/M3DV/LeFusion.

연구 동기 및 목표

• 임상 심장 MRI에서의 희소 데이터 분포와 알고리즘 편향 문제를 해결하기 위해 다양하고 현실적인 합성 병변을 생성함으로써 데이터 균형 문제를 해결한다.
• 이전의 생성 모델들이 병변을 합성할 때 정상 해부학의 품질을 유지하지 못하는 한계를 극복한다.
• DE-MRI에서 심근경색과 지속성 미세혈관 폐색을 데이터 기반으로 제어 가능한 방법으로 합성하기 위한 방법을 개발한다.
• 여러 병변 유형을 동시에 모델링하고 다양한 병변 마스크를 생성하여 데이터 증강 성능을 향상시킨다.
• nnUNet와 같은 최신 기술 수준의 모델을 사용하여 하류 분할 성능 향상 효과를 검증한다.

제안 방법

• 확산 학습 목표를 병변 영역에만 집중시켜 복잡도를 감소시키고 제어 가능성을 향상시킨다.
• 역확산 과정에 전방 확산 배경 정보를 통합하여 정상 해부학의 구조적 정밀도를 유지한다.
• 다중 채널 분해를 활용한 공동 모델링 전략을 제안하여 병렬적으로 여러 병변 유형을 처리한다.
• 합성 데이터의 형태와 질감 다양성을 높이기 위해 병변 마스크에 기반한 확산 기반 생성 모델을 개발한다.
• RePaint에 영감을 받은 리패ainting 스타일 프레임워크를 도입하지만, 병변 전용 학습과 배경 보존을 우선시하도록 적응시켰다.
• Emidec DE-MRI 심장 병변 데이터셋에서 학습 및 평가를 수행하였으며, 하류 분할 벤치마크로 nnUNet를 사용하였다.

실험 결과

연구 질문

• RQ1 병변 중심 확산 모델은 배경 해부학을 유지하면서도 심장 MRI에서 고해상도 심장 병변을 합성할 수 있는가?
• RQ2 병변 영역에 집중된 확산 학습이 배경-병변을 함께 학습하는 것에 비해 제어 가능성과 합성 품질을 어떻게 향상시키는가?
• RQ3 다양한 병변 유형(예: 심근경색 및 지속성 미세혈관 폐색)의 공동 모델링이 데이터 다양성과 하류 분할 성능에 얼마나 기여하는가?
• RQ4 학습된 확산 모델을 통해 생성된 병변 마스크는 수작업 마스크에 비해 더 다양하고 현실적인 형태를 가지는가?
• RQ5 LeFusion가 생성한 합성 데이터는 nnUNet와 같은 최신 기술 수준의 분할 네트워크 성능을 실제로 향상시키는가?

주요 결과

• LeFusion-J는 PMO 병변 클래스에서 가장 높은 Dice 점수(88.7)를 기록하여 LeFusion-S(85.2)와 RePaint(83.1)를 앞서며 공동 병변 모델링 성능이 뛰어남을 입증했다.
• 33×5개의 합성 정상 케이스(N”)를 활용한 결과, MI의 Dice 점수가 92.1로 상승하여 데이터 스케일 증가에 따른 일관된 성능 향상을 보였다.
• LeFusion-J는 기준 방법 대비 평균 PSNR(28.4)와 SSIM(0.91)이 가장 높아 이미지 품질과 실제 병변과의 구조 유사도가 뛰어남을 나타냈다.
• 확산 모델이 생성한 병변 마스크는 실제 마스크에 더 가까이 보이며, 수작업 마스크에 비해 더 높은 형태 다양성을 보였다.
• Cond-Diffusion 및 RePaint 방법은 배경의 구조를 손상시켜, 특히 PMO 병변에서 낮은 Dice 점수를 기록했다.
• LeFusion의 마스크 생성 기반 합성 데이터 증강은 분할 성능을 크게 향상시켜, 데이터 증강에 효과적임을 입증했다.