[논문 리뷰] SwinVFTR: A Novel Volumetric Feature-learning Transformer for 3D OCT Fluid Segmentation

Khondker Fariha Hossain, Sharif Amit Kamran외 3명|arXiv (Cornell University)|2023. 03. 16.

Retinal Imaging and AnalysisMedicine인용 수 3

한 줄 요약

SwinVFTR는 광학 간섭 단층촬영(OCT) 볼륨에서 정밀한 유체 세분화를 위해 설계된 새로운 3차원 볼륨형 트랜스포머 아키텍처이다. 채널 기반 볼륨 샘플링, 이동 윈도우 주의 메커니즘과 다중 수용장치 필드 잔류 블록을 갖춘 수정된 Swin 트랜스포머, 그리고 볼륨형 주의 스킵 연결을 활용하여, 각각 Spectralis, Cirrus, Topcon OCT 데이터셋에서 평균 Dice 점수 0.72, 0.59, 0.68를 기록하며 최신 기술 수준의 성능을 달성하였다.

ABSTRACT

Accurately segmenting fluid in 3D optical coherence tomography (OCT) images is critical for detecting eye diseases but remains challenging. Traditional autoencoder-based methods struggle with resolution loss and information recovery. While transformer-based models improve segmentation, they arent optimized for 3D OCT volumes, which vary by vendor and extraction technique. To address this, we propose SwinVFTR, a transformer architecture for precise fluid segmentation in 3D OCT images. SwinVFTR employs channel-wise volumetric sampling and a shifted window transformer block to improve fluid localization. Moreover, a novel volumetric attention block enhances spatial and depth-wise attention. Trained using multi-class dice loss, SwinVFTR outperforms existing models on Spectralis, Cirrus, and Topcon OCT datasets, achieving mean dice scores of 0.72, 0.59, and 0.68, respectively, along with superior performance in mean intersection-over-union (IOU) and structural similarity (SSIM) metrics.

연구 동기 및 목표

• 3D OCT 유체 세분화에서 자동코드어 모델의 해상도 손실 및 특징 복구 부족 등의 한계를 해결하기 위해.
• 3D 볼륨형 OCT 데이터에서 깊이 정보를 고려하지 못하는 기존 2D 트랜스포머 모델의 단점을 극복하기 위해.
• 다양한 OCT 장치 출력에 일반화 가능한 3D 기반, 제조사에 관계없는 세분화 모델을 개발하기 위해.
• 망막층 내 미세한 유체 경계 및 소규모 유체 침착물의 국소화를 향상시키기 위해.
• Dice, IOU, SSIM 지표를 활용한 다중 클래스 유체 세분화에서 뛰어난 성능을 달성하기 위해.

제안 방법

• 깊이 축(B-scan)에 따라 OCT 볼륨을 채널 기반 볼륨 샘플링 기법으로 자르며, 공간 해상도를 유지하면서도 변동하는 깊이 크기에 적응할 수 있도록 한다.
• 이동 윈도우 자기 주의 메커니즘과 다중 수용장치 필드(MRF) 잔류 하위 블록(합성곱과 확장된 합성곱의 조합)을 갖춘 수정된 Swin 트랜스포머 블록을 도입하여 특징 학습 능력을 향상시킨다.
• 에코더와 디코더 간 스킵 연결에 표준 잔류 연결 대신 새로운 볼륨형 주의(VA) 블록을 적용하여 공간적 및 깊이 기반 주의 병합을 가능하게 한다.
• 다중 해상도 특징 융합을 활용한 3D 합성곱 디코더를 사용하여 고해상도 세분화 맵을 재구성한다.
• 내막막 유체(IRF), 망막하 유체(SRF), 색소상피세포 박리(PED) 등 다중 클래스 세분화를 최적화하기 위해 다중 클래스 Dice 손실을 사용해 모델을 훈련시킨다.
• 세 가지 제조사 전용 데이터셋에서 성능을 최적화하기 위해 데이터 증강 및 하이퍼파rameter 튜닝(LR, 배치 크기, 에포크 수)을 적용한다.

실험 결과

연구 질문

• RQ1 3D 트랜스포머 기반 아키텍처가 기존 2D 및 3D CNN/트랜스포머 모델보다 3D OCT 볼륨의 유체 영역 세분화에서 우월한 성능을 보일 수 있는가?
• RQ2 채널 기반 볼륨 샘플링이 다양한 제조사에서의 깊이(B-scan 수)가 다른 OCT 데이터셋에서 효과적인 훈련을 가능하게 하는가?
• RQ3 제안된 볼륨형 주의 블록이 표준 잔류 스킵 연결에 비해 세분화 정확도를 얼마나 향상시키는가?
• RQ4 다중 수용장치 필드 잔류 블록의 통합이 소규모 유체 탐지에 유리한 특징 표현을 어떻게 향상시키는가?
• RQ5 모델이 재학습 없이도 서로 다른 OCT 스캐너 제조사(Spectralis, Cirrus, Topcon) 간에 일반화되는가?

주요 결과

• SwinVFTR는 Spectralis 데이터셋에서 평균 Dice 점수 0.72를 기록하며, 비교된 모든 모델보다 뛰어난 성능을 보였다.
• Cirrus 데이터셋에서 SwinVFTR는 평균 Dice 점수 0.59를 기록했으며, ResUNet-3D 대비 PED 세분화에서 10배 향상된 성능을 보였다.
• Topcon 데이터셋에서 SwinVFTR는 평균 Dice 점수 0.68를 기록했으며, ResUNet-3D 대비 PED 세분화에서 1.2배 향상된 성능을 보였다.
• 모든 세 데이터셋에서 가장 높은 평균 IOU(mIOU)와 구조적 유사도 지수(SSIM)를 기록하여, 뛰어난 경계 국소화 능력과 구조적 유사성 유지 능력을 입증했다.
• 절단 실험 결과 볼륨형 주의(VA) 및 다중 수용장치 필드(MRF) 블록이 모두 필수적임을 확인하였으며, 두 구성 요소를 모두 포함한 전체 모델이 각각의 부재한 변형보다 뛰어난 성능을 보였다.
• 배경 클래스 유무에 관계없이 높은 성능를 유지하여, 클래스 불균형에 대한 강건성과 낮은 거짓 양성률을 입증했다.