[논문 리뷰] Diffusion Models for Implicit Image Segmentation Ensembles

Julia Wolleb, Robin Sandkühler외 3명|arXiv (Cornell University)|2021. 12. 06.

Medical Image Segmentation Techniques인용 수 71

한 줄 요약

논문은 뇌종양 분할 마스크를 생성하기 위해 확산 확률 모델을 훈련시키고 확률적 샘플링을 이용해 암묵적 앙상블을 형성하여 불확실성 맵을 생성하고 BRATS2020에서 분할 성능을 향상시킵니다.

ABSTRACT

Diffusion models have shown impressive performance for generative modelling of images. In this paper, we present a novel semantic segmentation method based on diffusion models. By modifying the training and sampling scheme, we show that diffusion models can perform lesion segmentation of medical images. To generate an image specific segmentation, we train the model on the ground truth segmentation, and use the image as a prior during training and in every step during the sampling process. With the given stochastic sampling process, we can generate a distribution of segmentation masks. This property allows us to compute pixel-wise uncertainty maps of the segmentation, and allows an implicit ensemble of segmentations that increases the segmentation performance. We evaluate our method on the BRATS2020 dataset for brain tumor segmentation. Compared to state-of-the-art segmentation models, our approach yields good segmentation results and, additionally, detailed uncertainty maps.

연구 동기 및 목표

• 의료 영상 분할에서 불확실성 정량화를 동기화하고 병변 분할 성능을 향상시킨다.
• 입력 이미지를 훈련 및 샘플링 중 아틀라스 프라이어로 사용하는 확산 모델 기반 방법을 제안한다.
• 확률적 확산 샘플링을 통해 분할 마스크의 암묵적 앙상블 생성을 가능하게 한다.
• 여러 분할 샘플의 픽셀 단위 분산을 통해 불확실성 맵을 제공한다.

제안 방법

• 그라운드 트룻 분할 마스크에 노이즈를 추가하면서 입력 MR 이미지를 concat으로 프라이어로 도입하여 분할 마스크를 생성하기 위해 DDPM을 훈련시킨다.
• 이미지 프라이어 X_t = b 〮irst(x_b,t)로 조건을 주고 평균/분산 분할 앙상블을 생성하도록 DDPM 샘플링 방정식을 수정한다.
• 선형 노이즈 스케줄을 사용하여 T=1000 스텝, U-Net 아키텍처로 디노이징 스텝 ε_θ를 예측한다.
• 여러 개의 분할 샘플을 샘플링하고 그 분산을 계산하여 픽셀 단위 불확실성 맵을 계산한다.
• 여러 개의 샘플링 마스크를 평균화하고 임계값을 적용하여 최종 분할을 얻는 앙상블을 생산한다.
• BRATS2020에서 nnU-Net 및 SegNet과 벤치마킹하고 Dice, Jaccard, HD95 지표를 보고한다.

실험 결과

연구 질문

• RQ1 디퓨전 모델을 의료 영상의 의미적 분할에 이미지와 같은 프라이어로 조건을 주어 적응시킬 수 있는가?
• RQ2 확률적 확산 샘플링으로 여러 분할 샘플을 생성하는 것이 단일 샘플 예측에 비해 분할 지표를 향상시키는가?
• RQ3 확산 기반 분할이 BRATS2020에서 정확도 및 불확실성 시각화 측면에서 최첨단 방법(nnU-Net, SegNet)과 어떻게 비교되는가?

주요 결과

• 단일 샘플 DDPM 분할은 Dice 0.866 및 Jaccard 0.795를 달성하고 HD95는 6.052이다.
• 5개의 확산 샘플로 앙상블하면 Dice가 0.881, Jaccard가 0.819로 향상되며 HD95는 5.178이다.
• nnU-Net은 Dice 0.891, Jaccard 0.831, HD95 5.004를 달성한다.
• SegNet은 Dice 0.839, Jaccard 0.761, HD95 7.190을 보인다.
• 베이지안 SegNet(MC 드롭아웃, 5 샘플)에서 Dice 0.838, Jaccard 0.747, HD95 13.707을 보인다.
• 앙상블 평균 맵은 향상된 분할 성능과 상세한 불확실성 시각화를 가능하게 한다.