[논문 리뷰] DiffuseExpand: Expanding dataset for 2D medical image segmentation using diffusion models

Shitong Shao, Xiaohan Yuan외 4명|arXiv (Cornell University)|2023. 04. 26.

Generative Adversarial Networks and Image Synthesis인용 수 8

한 줄 요약

DiffuseExpand는 확산 확률 모델을 사용하여 2D 의학 영상 분할을 위한 고품질의 짝지어진 이미지-마스크 샘플을 생성하고, 데이터 확장을 효율적으로 수행하며 분할 성능을 향상시킵니다. 마스크 우선 합성, 분류기 가이드, 가속 샘플링 및 최종 품질 필터로 구성된 네 단계 파이프라인을 도입합니다.

ABSTRACT

Dataset expansion can effectively alleviate the problem of data scarcity for medical image segmentation, due to privacy concerns and labeling difficulties. However, existing expansion algorithms still face great challenges due to their inability of guaranteeing the diversity of synthesized images with paired segmentation masks. In recent years, Diffusion Probabilistic Models (DPMs) have shown powerful image synthesis performance, even better than Generative Adversarial Networks. Based on this insight, we propose an approach called DiffuseExpand for expanding datasets for 2D medical image segmentation using DPM, which first samples a variety of masks from Gaussian noise to ensure the diversity, and then synthesizes images to ensure the alignment of images and masks. After that, DiffuseExpand chooses high-quality samples to further enhance the effectiveness of data expansion. Our comparison and ablation experiments on COVID-19 and CGMH Pelvis datasets demonstrate the effectiveness of DiffuseExpand. Our code is released at https://github.com/shaoshitong/DiffuseExpand.

연구 동기 및 목표

• 프라이버시 및 라벨링 제약으로 인한 의학 영상 분할의 데이터 부족 문제를 해결한다.
• 높은 다양성과 정합성을 가진 이미지-마스크 쌍 샘플을 생성하기 위한 확산 모델 기반 파이프라인을 개발한다.
• 단계별 컨디셔닝, 분류기 가이드 및 샘플 선별을 통해 데이터 확장 품질을 향상시킨다.
• COVID-19 흉부 X선 및 CGMH Pelvis 데이터셋에서 효과를 입증한다.
• 오픈 소스 코드를 제공하고 ablation 연구 및 비교를 통해 개선점을 검증한다.

제안 방법

• Stage I: 분류기 가이드 없이 사전 학습된 확산 확률 모델(DPM)을 미세조정하여 분포를 의학 영상에 매핑한다.
• Stage II: 분류기 가이드를 가능하게 하기 위해 세그멘터를 학습하고 Dice 및 BCE 손실을 통해 이미지-마스크 정합성을 평가한다.
• Stage III: DPM Solver++를 사용하여 여러 조건하에서 이미지-마스크 쌍을 합성하고, 분류기 가이드를 기반으로 한 조건부 샘플링을 가능하게 한다.
• Stage IV: 학습된 네트워크와 Dice 임계치를 사용하여 저질 샘플을 제거하고 고품질의 확장 데이터를 보장한다.
• Introduce a continuous-time probability flow ODE formulation for conditional sampling and leverage classifier guidance with a temperature parameter tau to control gradient magnitudes.
• Provide ablation studies to compare classifier-free vs. classifier-guided sampling and assess the impact of tau/s on diversity and Dice performance.

실험 결과

연구 질문

• RQ1 확산 모델이 2D 의학 영상 분할을 위한 다양하고 잘 정렬된 이미지-마스크 쌍을 생성할 수 있는가?
• RQ2 단계적 컨디셔닝과 분류기 가이드는 확장된 분할 데이터세트의 품질과 활용도를 향상시키는가?
• RQ3 샘플 선별(Stage IV)이 Dice 점수와 데이터 확장의 샘플 다양성에 미치는 영향은 무엇인가?
• RQ4 의료 영상 분야에서 확산 기반 확장 방법이 GAN 기반이나 증강 기반 접근법과 어떻게 비교되는가?

주요 결과

• DiffuseExpand는 테스트 데이터셋에서 Dice 점수를 향상시키는 고충실도하고 다양한 이미지-마스크 쌍을 산출한다.
• 학습된 세그멘터를 이용한 분류기 가이드와 Stage IV 선별을 결합하면 이미지와 마스크 간 정합성이 향상된다.
• 가이드를 위한 tau 기반 컨디셔닝은 유사한 설정에서 gradient-scaling 값 s보다 더 높은 Dice 점수를 산출한다.
• 확산 기반 확장은 Dice 기반 분할 성능 측면에서 Synthmed와 XLsor를 능가한다.
• Stage IV 선별은 샘플 다양성을 감소시키지만 전반적 품질을 높여 하위 분할의 성능에 기여한다.