[논문 리뷰] A Diffusion Model Predicts 3D Shapes from 2D Microscopy Images

Dominik Waibel, Ernst Röell외 3명|arXiv (Cornell University)|2022. 08. 30.

Cell Image Analysis Techniques인용 수 4

한 줄 요약

이 논문은 단일 2D 형광 이미지를 조건으로 하여 2D 미세망원 영상에서 현실적인 3D 세포 형상을 예측하는 확산 모델인 DISPR를 소개한다. 이 모델은 합성 3D 재구성 영상을 활용한 데이터 증강을 통해 적은 수의 클래스 분류 성능을 크게 향상시켜, 적은 수의 클래스에 대한 매크로 F1 스코어를 55.2%에서 72.2%로 상승시켰다.

ABSTRACT

Diffusion models are a special type of generative model, capable of synthesising new data from a learnt distribution. We introduce DISPR, a diffusion-based model for solving the inverse problem of three-dimensional (3D) cell shape prediction from two-dimensional (2D) single cell microscopy images. Using the 2D microscopy image as a prior, DISPR is conditioned to predict realistic 3D shape reconstructions. To showcase the applicability of DISPR as a data augmentation tool in a feature-based single cell classification task, we extract morphological features from the red blood cells grouped into six highly imbalanced classes. Adding features from the DISPR predictions to the three minority classes improved the macro F1 score from $F1_ ext{macro} = 55.2 \pm 4.6\%$ to $F1_ ext{macro} = 72.2 \pm 4.9\%$. We thus demonstrate that diffusion models can be successfully applied to inverse biomedical problems, and that they learn to reconstruct 3D shapes with realistic morphological features from 2D microscopy images.

연구 동기 및 목표

• 3D 공형 현미경의 느린 속도와 높은 비용 문제를 해결하기 위해 2D 영상에서 빠르고 현실적인 3D 형상 예측을 가능하게 한다.
• 단일 세포 형태 분류에서의 데이터 부족과 클래스 불균형 문제를 해결하기 위해 합성 3D 형상을 생성하여 데이터 증강을 수행한다.
• 확산 모델의 분포적 성질을 활용하여 단일 2D 입력에서 다양한 현실적인 3D 재구성 영상을 생성한다.
• 불균형 생물학적 데이터셋에서의 하류 분류 성능 향상에 있어 확산 기반 3D 형상 예측의 유용성을 평가한다.

제안 방법

• 모델는 T 단계 동안 점진적으로 가우시안 노이즈를 3D 세포 볼륨에 추가하는 전방 노이징 과정을 뒤집는 덴즈잉 확산 확률 모델로 훈련된다.
• 모델은 2D 형광 현미경 영상과 세그멘테이션 마스크로 조건화되며, 이는 덴즈잉 과정 중 사전 지식으로 기능한다.
• 추론 과정에서 모델은 순수한 노이즈에서 시작하여 T번의 덴즈잉 단계를 수행하며, t=0 단계에서 3D 형상 예측을 생성한다.
• 노이즈 스케줄을 최적화하여 추론 효율성을 향상시키며, 변분 하한 손실을 사용해 모델을 훈련시킨다.
• 3D 볼륨 내 장거리 의존성을 모델링하기 위해 U-Net 기반 아키텍처에 스위프 커넥션과 어텐션 메커니즘을 통합한다.
• 소수 클래스에 대해 각 2D 영상당 다섯 개의 3D 재구성 영상을 예측하여 합성 3D 형상을 생성함으로써 데이터 증강을 수행한다.

Fig. 1 : Diffusion steps of one forward pass through DISPR visualized during inference between step $t=999$ with the largest amount of Gaussian noise and the DISPR prediction of $t=0$ , for a stromatocyte cell (see section 2 ). Note that the zoom level in the visualizations is not uniform.

실험 결과

연구 질문

• RQ1 확산 모델은 단일 2D 현미경 영상에서 현실적인 3D 세포 형상을 생성할 수 있는가?
• RQ2 확산 모델의 분포적 출력은 불균형 단일 세포 분류 작업의 데이터 증강에 도움이 되는가?
• RQ3 DISPR가 생성한 3D 형상의 형태적 특징은 참값과 비교해 현실성과 정확도 면에서 어떻게 다른가?
• RQ4 DISPR가 생성한 데이터를 사용할 경우 소수 클래스 분류의 매크로 F1 스코어는 어느 정도 향상되는가?

주요 결과

• 소수 클래스 분류 성능 향상을 위해 DISPR가 생성한 3D 형상의 특징을 사용한 데이터 증강을 통해 적혈구 형태 분류의 매크로 F1 스코어가 55.2% ± 4.6%에서 72.2% ± 4.9%로 향상되었다.
• 가중치 F1 스코어는 85.7% ± 2.1%에서 88.6% ± 1.1%로 상승하여 전체 분류 성능 향상을 나타냈다.
• 데이터 증강 후, 케라토시트의 정확히 분류된 수는 14개에서 24개로, 다중엽세포는 1개에서 8개로 증가했다.
• 최신의 결정론적 모델 대비 DISPR가 생성한 형상은 상대 표면 거칠기와 곡률 오차가 낮아 형태적 충실도가 높음을 보였다.
• 모델는 생성된 형상 분포에서 참값 데이터셋의 클래스 특화 형태적 특징 분포를 성공적으로 반영했다.
• 추론 시간이 더 길었음에도 불구하고(예측당 4:45분), DISPR는 SHAPR와 같은 결정론적 오토에인코더보다 다양한 현실적인 3D 재구성 영상을 더 잘 생성했다.

Fig. 2 : DISPR is trained to denoise the 3D volume $\mathbf{x}_{b,t}$ containing added stochastic Gaussian noise, to obtain $\mathbf{x}_{b,t-1}$ , thus reversing the noising step $q$ . In each forward pass $p_{\theta}$ , we constrain our 3D model with one 2D image $b$ containing a fluorecent image a