[Paper Review] SwinVFTR: A Novel Volumetric Feature-learning Transformer for 3D OCT Fluid Segmentation

Khondker Fariha Hossain, Sharif Amit Kamran+3 more|arXiv (Cornell University)|Mar 16, 2023

Retinal Imaging and AnalysisMedicine3 citations

TL;DR

SwinVFTR is a novel 3D volumetric transformer architecture designed for precise fluid segmentation in optical coherence tomography (OCT) volumes. It employs channel-wise volumetric sampling, a modified Swin transformer with shifted-window attention and multi-receptive field residual blocks, and a volumetric attention skip connection, achieving state-of-the-art performance with mean Dice scores of 0.72, 0.59, and 0.68 on Spectralis, Cirrus, and Topcon OCT datasets, respectively.

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.

Motivation & Objective

• To address the limitations of autoencoding models in 3D OCT fluid segmentation, particularly resolution loss and poor feature recovery.
• To overcome the shortcomings of existing 2D transformer models that lack depth context in volumetric OCT data.
• To develop a 3D-capable, vendor-agnostic segmentation model that generalizes across different OCT device outputs.
• To improve localization of fine fluid boundaries and small fluid deposits in retinal layers.
• To achieve superior performance on multi-class fluid segmentation using Dice, IOU, and SSIM metrics.

Proposed method

• Proposes a channel-wise volumetric sampling technique to crop OCT volumes along the depth axis (B-scans), preserving spatial resolution while adapting to variable depth sizes.
• Introduces a modified Swin transformer block with shifted-window self-attention and a multi-receptive field (MRF) residual sub-block combining convolution and dilated convolution for enhanced feature learning.
• Employs a novel volumetric attention (VA) block in skip connections between encoder and decoder to replace standard residual connections, enabling spatial and depth-wise attention fusion.
• Uses a 3D convolutional decoder with multi-scale feature fusion to reconstruct high-resolution segmentation maps.
• Trains the model using multi-class Dice loss to optimize segmentation across intra-retinal fluid (IRF), sub-retinal fluid (SRF), and pigment epithelium detachment (PED).
• Applies data augmentation and hyperparameter tuning (learning rate, batch size, number of epochs) to optimize performance across three vendor-specific datasets.

Experimental results

Research questions

• RQ1 Can a 3D transformer-based architecture outperform existing 2D and 3D CNN/transformer models in segmenting fluid regions in 3D OCT volumes?
• RQ2 Does channel-wise volumetric sampling enable effective training on OCT datasets with varying depth (B-scan count) across different vendors?
• RQ3 To what extent does the proposed volumetric attention block improve segmentation accuracy compared to standard residual skip connections?
• RQ4 How does the integration of multi-receptive field residual blocks enhance feature representation in the Swin transformer encoder for small fluid detection?
• RQ5 Does the model generalize across different OCT scanner vendors (Spectralis, Cirrus, Topcon) without retraining?

Key findings

• SwinVFTR achieves a mean Dice score of 0.72 on the Spectralis dataset, outperforming all compared models.
• On the Cirrus dataset, SwinVFTR achieves a mean Dice score of 0.59, with a 10× improvement in PED segmentation compared to ResUNet-3D.
• On the Topcon dataset, SwinVFTR achieves a mean Dice score of 0.68, showing a 1.2× improvement in PED segmentation over ResUNet-3D.
• The model achieves the highest mean Intersection-over-Union (mIOU) and Structural Similarity Index (SSIM) across all three datasets, indicating superior boundary localization and structural fidelity.
• Abllation studies confirm that both the volumetric attention (VA) and multi-receptive field (MRF) blocks are essential, with the full model outperforming variants lacking either component.
• The model maintains high performance with and without background class, demonstrating robustness to class imbalance and low false-positive rates.