[논문 리뷰] Anatomically-Controllable Medical Image Generation with Segmentation-Guided Diffusion Models

Nicholas Konz, Yuwen Chen외 2명|arXiv (Cornell University)|2024. 02. 07.

Medical Image Segmentation Techniques인용 수 4

한 줄 요약

이 논문은 샘플링 중에 부분적 또는 완전한 해부학적 분할 마스크를 조건으로 삼아 해부학적으로 현실적인 의료 영상 생성을 위한 분할 마스크 유도 확산 모델을 제안한다. 새로운 마스크 제거 학습 전략을 사용함으로써, 일부 구조가 마스크로 가려져 있어도 고해상도이자 해부학적으로 타당한 영상을 학습 과정에서 생성할 수 있으며, 제어넷(ControlNet)과 같은 최신 기술들보다도 뛰어난 해부학적 현실성과 마스크 충실도를 달성한다. 특히 유방 MRI 및 CT 데이터셋에서 성능이 뛰어나다.

ABSTRACT

Diffusion models have enabled remarkably high-quality medical image generation, yet it is challenging to enforce anatomical constraints in generated images. To this end, we propose a diffusion model-based method that supports anatomically-controllable medical image generation, by following a multi-class anatomical segmentation mask at each sampling step. We additionally introduce a random mask ablation training algorithm to enable conditioning on a selected combination of anatomical constraints while allowing flexibility in other anatomical areas. We compare our method ("SegGuidedDiff") to existing methods on breast MRI and abdominal/neck-to-pelvis CT datasets with a wide range of anatomical objects. Results show that our method reaches a new state-of-the-art in the faithfulness of generated images to input anatomical masks on both datasets, and is on par for general anatomical realism. Finally, our model also enjoys the extra benefit of being able to adjust the anatomical similarity of generated images to real images of choice through interpolation in its latent space. SegGuidedDiff has many applications, including cross-modality translation, and the generation of paired or counterfactual data. Our code is available at https://github.com/mazurowski-lab/segmentation-guided-diffusion.

연구 동기 및 목표

• 확산 기반 의료 영상 생성에서 낮은 전반적 해부학적 현실성 문제를 해결하기 위해.
• 이미지 생성 중에 부분적 또는 불완전한 해부학적 분할 마스크를 탄력적으로 조건으로 사용할 수 있도록 하기 위해.
• 학습 중 마스크 제거를 통한 자기지도 학습을 활용하여 무조건적 이미지 생성 성능을 향상시키기 위해.
• 사전 정렬된 이미지 합성 및 반사적 분석과 같은 응용 분야를 촉진하기 위해.
• 기존의 마스크 조건 기반 및 무조건적 생성 모델보다 해부학적 정확도와 입력 마스크 준수 능력에서 뛰어난 성능을 내기 위해.

제안 방법

• 모델은 모든 노이즈 제거 단계에서 분할 마스크를 조건으로 삼는 이미지 공간 확산 과정을 사용한다.
• 입력 마스크 내에서 결손이 있는 해부학적 클래스의 모든 가능한 조합을 노출시키는 새로운 마스크 제거 학습 전략을 적용한다.
• 이 학습 방식은 모델이 마스크가 없는 영역에서 현실적인 해부학을 추론하고 생성할 수 있도록 하며, 제공된 제약 조건을 충실하게 따르도록 한다.
• 모델은 실제 의료 영상과 그에 해당하는 분할 마스크를 기반으로 종단 간(end-to-end)으로 훈련된다.
• 생성된 이미지의 마스크 충실도를 평가하기 위해 보조 분할 네트워크를 사용한다.
• 조건 기반 생성(부분 마스크 사용 시)과 무조건적 생성(빈 마스크 사용 시)을 모두 지원하며, 자기지도 학습 덕분에 후자의 경우 더 높은 현실성 확보가 가능하다.

Figure 1: Standard diffusion models (right) can fail to create realistic tissue, even if the overall image appears high-quality. From left to right: (1) randomly sampled real breast MRI images, (2) synthetic images created by our segmentation-guided diffusion model, and (3) by the same model but wit

실험 결과

연구 질문

• RQ1 모든 노이즈 제거 단계에서 분할 마스크 유도를 통합함으로써 확산 모델이 더 해부학적으로 현실적인 의료 영상 생성이 가능한가?
• RQ2 부분적 또는 제거된 마스크가 제공될 경우, 모델이 마스크가 없는 영역에서 타당한 해부학을 생성할 수 있는가?
• RQ3 마스크 제거 학습이 무조건적(빈 마스크) 설정에서도 현실적인 이미지 생성 능력을 향상시키는가?
• RQ4 ControlNet과 같은 최신 기술들과 비교해 모델의 해부학적 현실성과 마스크 충실도 성능은 어떠한가?
• RQ5 모델을 사용해 후속 응용 분야를 위한 사전 정렬된 쌍체 의료 영상 합성을 수행할 수 있는가?

주요 결과

• 제안된 모델은 유방 MRI 데이터셋에서 마스크 충실도의 덱스 점수(Dice score)가 0.9027을 기록하였고, CT 장기 데이터셋에서는 0.8980을 기록하여 ControlNet의 각각 0.3636과 0.1132를 크게 앞서나갔다.
• 동일한 데이터셋에서 생성된 이미지의 분할 겹침과 실제 이미지 예측 간의 덱스 점수는 각각 0.8593과 0.8797을 기록하여 높은 해부학적 정확도를 나타냈다.
• 빈 마스크(무조건적 생성) 조건에서도, 마스크 제거 학습을 통해 자기지도 학습을 수행한 덕분에 표준 무조건적 확산 모델보다 더 해부학적으로 현실적인 이미지를 생성하였다.
• 모델의 합성 데이터 기반으로 훈련된 분할 네트워크는 실제 데이터 기반으로 훈련된 네트워크와 거의 유사한 성능을 보였으며, 덱스 차이가 0.04 이내였고, 이는 생성된 이미지의 높은 현실성을 확인하는 데 기여했다.
• 혈관과 섬유선모세관 조직 등 일부 해부학적 구조만 제약을 둘 경우에도 다양한 해부학적으로 타당한 이미지를 성공적으로 생성하였고, 동시에 유방 형태 등 다른 요소들에 대해서는 다양성을 허용하였다.
• 저자들은 코드베이스와 공동으로 정렬된 유방 MRI 생성 데이터셋을 공개하여 공동 연구를 지원하였다.

Figure 3: Breast MRI (top) and CT organ (bottom) images generated with different segmentation-conditional and unconditional models. Samples from conditional models ( blue ) are below their input mask, and unconditional model ( red ) samples are random. For breast MRI, the breast, blood vessel, and f