SAM-OCTA: A Fine-Tuning Strategy for Applying Foundation Model to OCTA Image Segmentation Tasks

As a typical architecture for deep image processing, the vision transformer (ViT) is frequently used for segmentation tasks in OCTA [7]. In OCTA images, the distribution of RV is extensive, and it requires the models to effectively utilize the global information in the images. The TCU-Net, OCT2Former, and StruNet methods have improved ViT, achieving continuous RV segmentation and addressing issues such as vessel discontinuities or missing segments [8, 9, 10]. Other methods, from the perspectives of efficiency, denoising, and the utilization of three-dimensional data, have designed a series of techniques and strategies, achieving promising segmentation results on OCTA datasets [2, 11, 12, 13, 14]. The above-mentioned methods have demonstrated that existing deep networks are capable of achieving precise segmentation of RV and FAZ.

The SAM is a foundational vision model for general image segmentation. With the ability to segment diverse objects, parts, and visual structures in various scenarios, SAM takes prompts in the form of points, bounding boxes, or coarse masks as input. Its remarkable zero-shot segmentation capabilities enable its easy transfer to numerous applications through simple prompting [5]. Although SAM has established an efficient data engine for model training, there are relatively few cases collected for medical applications or other rare image scenarios. Therefore, some fine-tuning methods have been applied to SAM to improve its performance in certain segmentation failure cases [15, 16]. The common characteristic of these fine-tuning methods is that they introduce additional network layers on top of the pre-trained SAM. By adding a small number of trainable parameters, fine-tuning becomes feasible through training on the new dataset. The advantage of fine-tuning methods lies in their ability to preserve SAM’s strong zero-shot capabilities and flexibility.

The image encoder utilizes a ViT pre-trained with the masked auto-encode method. The ViT model comes in three variants: vit-b, vit-l, and vit-h which can only process fixed-size inputs (e.g. $1024*1024*3$). To support input images of different resolutions, scaling and padding operations are employed. In this study, we used the image encoder from the ”vit-h” model for the fine-tuning process.

As shown in Figure 2, OCTA data is inherently in 3D format, but most datasets provide en-face 2D projection forms. En-face projection is obtained through layer-wise segmentation based on vascular anatomical structures. As SAM requires three-channel images as input, in this work, we stack projection layers in different depths of OCTA images to adapt to this input format. The benefit of this approach is that it preserves the vascular structure information in the OCTA images while fully utilizing SAM’s feature-extracting capabilities. Fine-tuning aims to retain SAM’s powerful image-understanding capabilities while enhancing its performance on OCTA. The approach used in this paper involves utilizing the LoRA technique [17], which introduces additional linear network layers in each transformer block of the image encoder, similar in form to a ResNet block. During the training process, the weights of the SAM are frozen, and only the newly introduced parameters are updated.

The prompt encoder is divided into two types of prompts: sparse prompts (points, boxes, text) and dense prompts (masks). In our work, we chose points as the prompt for OCTA segmentation. For each sample’s prompt points input, assuming there are $n$ points in the prompt point input, it can be represented as ${(x_{1},y_{1},1),(x_{2},y_{2},1),...,(x_{n},y_{n},0)}$, where $x$ and $y$ denote the coordinates of prompt points in the image. The values ”1” and ”0” indicate positive (foreground) and negative (background) points, respectively. The prompt encoder of SAM will perform embedding on this input, and due to its pre-training, it can appropriately integrate with the information from the input image.

The prompt points generation strategy has two types: the global mode and the local mode. The global mode is applied to all OCTA segmentation tasks, while the local mode is specific to artery/vein segmentation. The study of segmenting individual vessels, as a local segmentation task, has not been attempted in previous works. By the prompt encoder, more accurate regional vessel segmentation can be achieved in OCTA datasets. For this task, the first step is to identify and label all connected components in the segmentation masks with unique identifiers. Due to the weak connectivity at the endpoints of some vessels in OCTA labels, we adopt the eight-connectivity criterion. Then positive points are randomly selected within each connected component. Due to varying numbers of vessels in different samples, to standardize their data format, negative points are added from the background adjacent to the connected components. The prompt points generation process can be described as Figure 3.

The role of the mask decoder is to efficiently map the image embeddings, prompt embeddings, and output tokens to a segmentation mask. A modified version of the transformer decoder block is employed, followed by a dynamic mask prediction head. For an image input and corresponding prompt input, the mask decoder outputs multiple segmentation masks to represent objects at different semantic levels. In this work, the loss function (in the fine-tuning process) is computed based on the segmentation output with the highest confidence.

The publicly available dataset used in this paper is OCTA-500 [18]. The OCTA-500 dataset contains 500 samples, classified based on the field of view (FoV): $3mm*3mm$ (3M) and $6mm*6mm$ (6M). The corresponding image resolutions are $304*304$ and $400*400$, with 200 and 300 samples respectively. The OCTA-500 dataset provides annotations for RV, FAZ, capillary, artery, and vein. The adopted data augmentation tool is Albumentations [19]. The data augmentation strategies include horizontal flipping, brightness and contrast adjustment, and random slight rotation.

The SAM is deployed on A100 graphic cards with 80 GB memory. The 10-fold cross-validation is adopted to evaluate the training results. The optimizer used is AdamW, and the learning rate adopts a warm-up strategy, starting from $10^{-5}$ and gradually increasing to $10^{-3}$.

The loss functions used for fine-tuning vary depending on the segmentation tasks. For FAZ and capillary, the Dice loss is employed. However, for RV, artery, and vein, the clDice loss is utilized which is more feasible for tubular segmentation [20]. These two loss functions can be represented as:

$\displaystyle L_{clDice}=0.2*L_{Dice}+0.8*L_{clDice}^{\prime},$

where $L_{Dice}=1-\frac{2*|\hat{Y}\cap Y|}{|\hat{Y}|+|Y|}$,

$L_{clDice}^{\prime}=1-2*\frac{T{prec}(\hat{Y_{s}},Y)*T{sens}(Y_{s},\hat{Y})}{T{prec}(\hat{Y_{s}},Y)+T{sens}(Y_{s},\hat{Y})}$,

$Y$ → the ground-truth,

$\hat{Y}$ → the predicted value,

$Y_{s},\hat{Y_{s}}$ → soft-skeleton($Y$, $\hat{Y}$), and

$T{prec},T{sens}$ → precision and sensitivity.

We conducted extensive experiments with various cases on the OCTA datasets. The segmentation results using metrics Dice, and Jaccard, which are calculated as follows:

$\displaystyle Dice(\hat{Y},Y)=\frac{2|\hat{Y}\cap Y|}{|\hat{Y}|+|Y|},\ \ Jaccard(\hat{Y},Y)=\frac{|\hat{Y}\cap Y|}{|\hat{Y}\cup Y|}.$