[논문 리뷰] Diff-UNet: A Diffusion Embedded Network for Volumetric Segmentation

Zhaohu Xing, Liang Wan외 3명|arXiv (Cornell University)|2023. 03. 18.

Radiomics and Machine Learning in Medical Imaging인용 수 41

한 줄 요약

Diff-UNet은 확산 모델을 3D U-Net에 통합하여 체부 영상 의학 이미지 분할(volumetric medical image segmentation)을 수행하고, Step-Uncertainty based Fusion(SUF)을 사용해 다단계 예측을 결합하여 MRI와 CT 데이터셋 전반에서 견고한 결과를 얻는다.

ABSTRACT

In recent years, Denoising Diffusion Models have demonstrated remarkable success in generating semantically valuable pixel-wise representations for image generative modeling. In this study, we propose a novel end-to-end framework, called Diff-UNet, for medical volumetric segmentation. Our approach integrates the diffusion model into a standard U-shaped architecture to extract semantic information from the input volume effectively, resulting in excellent pixel-level representations for medical volumetric segmentation. To enhance the robustness of the diffusion model's prediction results, we also introduce a Step-Uncertainty based Fusion (SUF) module during inference to combine the outputs of the diffusion models at each step. We evaluate our method on three datasets, including multimodal brain tumors in MRI, liver tumors, and multi-organ CT volumes, and demonstrate that Diff-UNet outperforms other state-of-the-art methods significantly. Our experimental results also indicate the universality and effectiveness of the proposed model. The proposed framework has the potential to facilitate the accurate diagnosis and treatment of medical conditions by enabling more precise segmentation of anatomical structures. The codes of Diff-UNet are available at https://github.com/ge-xing/Diff-UNet

연구 동기 및 목표

• 노이즈 제거 확산 모델을 활용한 강건한 3D 의료 부피 분할의 동기 부여.
• 레이블 임베딩을 통해 다중 라벨 분할을 처리하는 확산 기반 엔드 투 엔드 프레임워크 제안.
• 시맨틱 정보를 추출하기 위한 별도의 피처 인코더와 디노이징 모듈로 구성된 디노이징 모듈 도입.
• 테스트 중 다중 확산 스텝에서 예측을 융합하기 위한 Step-Uncertainty 기반 Fusion(SUF) 모듈 개발.
• Diff-UNet은 BraTS2020, MSD Liver, BTCV 데이터셋에서 최첨단 방법보다 우월함을 입증합니다.

제안 방법

• 확산 기반 예측을 가능하게 하기 위해 다중 클래스 분할을 원-핫 인코딩 다채널 라벨 맵으로 변환.
• X0 예측을 위해 노이즈가 있는 원-핫 라벨과 볼륨 특징을 융합하는 피처 인코더와 디노이징-유넷으로 구성된 디노이징 모듈.
• Dice, BCE, MSE 손실을 결합한 합성 손실로 학습.
• 테스트 중 다수의 DDIM 스텝을 생성하고 스텝 인덱스와 예측 불확실성을 고려하는 SUF 가중 방식으로 예측을 융합.
• MC와 유사 샘플링으로 각 스텟당 다중 순방향 패스를 통해 불확실성을 추정하고 이를 최종 분할 출력에서 스텝 예측의 가중치로 사용.

실험 결과

연구 질문

• RQ1 확산 모델을 3D U-Net 아키텍처에 효과적으로 삽입해 다중 라벨 의료 부피 분할을 수행할 수 있는가?
• RQ2 Step-Uncertainty 기반 Fusion 메커니즘이 도전적인 데이터셋에서 확산 기반 분할의 견고성을 개선하는가?
• RQ3 Diff-UNet은 다기관 CT 및 다모달 MRI 분할 벤치마크에서 최첨단 방법과 비교하여 어떤 성능을 보이는가?

주요 결과

• Diff-UNet은 BraTS2020에서 평균 Dice 점수 85.35%로 우수한 성능을 달성하고 여러 SOTA 방법과 비교하여 TC/WT/ET 영역이 개선되었습니다.
• MSD Liver에서 Diff-UNet은 Liver에서 95.72% Dice를 달성하고 HD95는 0.222이며, Liver Tumor에서 51.65% Dice와 17.280 HD95를 달성하고, 영역 전반에서 73.69 Dice와 8.751 HD95를 평균.
• BTCV에서 Diff-UNet은 8개 기관에서 HD95 8.115로 최고의 평균 Dice(83.75)를 달성.
• 충실도 연구에서 별도의 Feature Encoder(FE)와 SUF 모듈의 도입이 성능을 향상시키며, 불확실성 샘플링에서 S=4에서 최적의 결과를 보임.
• 정성적 결과는 Diff-UNet이 경쟁 방법보다 작은 대상에 대해 더 정확한 분할을 생성한다는 것을 시사.