MedSegDiff: Medical Image Segmentation with Diffusion Probabilistic Model
Diffusion Probabilistic Model (DPM) has recently become one of the hottest topics in computer vision. Its image generation applications, such as Imagen, Latent Diffusion Models, and Stable Diffusio...
proceedings.mlr.press
medical segmentation + diffusion model의 최신 논문
Abstract
Diffusion Probabilistic Model(DPM)은 비전분야에서 떠오르는 분야 중 하나이다. 최근의 연구들은 DPM을 이용해 image deblurring, super-resolution, anomaly detection등에서도 사용되기에 우리는 medical image segmentation에 적용해보려 한다. DPM의 단계적인 regional attention을 향상시키기 위해 단계별로 조건을 적용하는 Dynamic Conditional Encoding을 제안한다. 또한 이 과정에서 발생하는 high-frequency 잡음을 제거하기 위해 Feature Frequency Parser(FF-Parser)ㅇ르 제안한다.
Method
우리의 모델은 DDPM을 기반으로 하며 diffusion, reverse diffusion의 2단계로 진행된다. 순전파에서 segmentation label \(x_0\)은 \(T\)단계에 걸쳐 가우시안 노이즈가 추가되며 복원단계에서 원본 데이터로 회복한다. 우리는 기본적인 DPM의 구현을 따르며 U-Net형태의 네트워크를 적용하였다.
Dynamic Conditional Encoding
- 대부분의 conditional DPM의 conditional 사전은 unique feature embedding이다. 하지만 medical segmentation은 모호한 객체들을 가지고 있어 배경과 이들을 구별하기 어렵다. 또한 current step의 segmentation map은 대상 영역이 표시되어 있지만 완전히 정확하지는 않다. 그렇기에, current step \(x_t\)에 conditional raw image encoding을 통합하려고 한다. 이는 feature level에서 구현된다. 인코더에 대해 각 단계에서의 특징맵 \(m_I^k\)와 인코딩된 \(m_x^k\)이 같은 크기로 융합되어진다. 이는 Resnet34의 convolutional stage에 해당하며 중간의 2단계에 걸쳐 진행된다. Dynamic Conditional Encoding은 segmentation을 동적으로 보정하는데 도움을 주지만 \(x_t\)단계에서의 noise가 더해지기 때문에 FF-Parset을 통해 노이즈를 제한하여야 한다.
FF-Parser
- 주된 아이디어는 Fourier-space feature에 적용되어 있는 attentive map을 학습하는 것이다. 특징맵 \(m\in R^{H\times W\times C}\)을 만족하는 디코더에 대해 2D FFT(Fast Fourier Transform)을 적용한다. 이후 attentive map을 곱한 후 reverse FFT를 하여 이를 학습하는 것이라고 한다.(잘 이해가 안됨)
 |
 |
 |
수식으로 보면 feature map \(m\)에 대해 학습가능한 attentive map\(A\)를 추가하여 이를 복원하도록 훈련시키는 것 같다.
모델은 DPM의 훈련과정을 따르며 ResUNet의 모델 구조를 수정한 형태이다.
Implementation Details
샘플링에 따라 Large, Base, Small 모델이 존재하는데 이는 각 x6, x5, x4의 크기로 샘플링하여 훈련한 모델이다. DPM의 원 논문보다는 적은 100번의 diffusion step을 적용하였고 256x256크기의 이미지를 사용하였다.
Optimizer : AdamW
Batch size : 32
