[Paper Review] Learnable Ophthalmology SAM

Zhongxi Qiu, Yan Hu+2 more|arXiv (Cornell University)|Apr 26, 2023

Retinal Imaging and Analysis12 citations

TL;DR

This paper introduces a learnable prompt layer for Segment Anything (SAM) to enable one-shot, multi-modal ophthalmology image segmentation by freezing the backbone and training a task-specific prompt and head. It demonstrates improved segmentation across multiple ophthalmology tasks and datasets with strong generalization.

ABSTRACT

Segmentation is vital for ophthalmology image analysis. But its various modal images hinder most of the existing segmentation algorithms applications, as they rely on training based on a large number of labels or hold weak generalization ability. Based on Segment Anything (SAM), we propose a simple but effective learnable prompt layer suitable for multiple target segmentation in ophthalmology multi-modal images, named Learnable Ophthalmology Segment Anything (SAM). The learnable prompt layer learns medical prior knowledge from each transformer layer. During training, we only train the prompt layer and task head based on a one-shot mechanism. We demonstrate the effectiveness of our thought based on four medical segmentation tasks based on nine publicly available datasets. Moreover, we only provide a new improvement thought for applying the existing fundamental CV models in the medical field. Our codes are available at \href{https://github.com/Qsingle/LearnablePromptSAM}{website}.

Motivation & Objective

• Motivate segmentation in ophthalmology due to multi-modal imaging and limited labels.
• Propose a learnable prompt layer inserted between SAM's transformer blocks to inject medical prior knowledge.
• Enable one-shot fine-tuning by training only the prompt layers and task head.
• Demonstrate effectiveness across multiple tasks (vessel, lesion, retinal layer) and datasets.
• Assess generalization via zero-shot experiments across datasets.

Proposed method

• Freeze the SAM backbone while inserting a learnable prompt layer between each transformer block.
• Use two 1x1 convolutions with LayerNorm and GELU, plus a depthwise 3x3 convolution to form each Prompt_i.
• Compute prompts from the i-th transformer's input features as Prompt_i = delta(LN(W1(delta(LN(DW3(delta(LN(W1(fi))))))))).
• Attach a task-specific segmentation head with upsampling and multi-scale convolution to produce segmentation maps.
• Train only the prompt layers and the task head with a one-shot mechanism.
• Evaluate on nine public datasets across three segmentation tasks (vessel, lesion, retinal layer) using metrics such as Dice, Precision, Recall, BM, and IoU.

Experimental results

Research questions

• RQ1 Can a learnable prompt layer enable SAM to segment ophthalmology-specific targets across multi-modal images without full fine-tuning?
• RQ2 Does the proposed prompt learn medical priors that generalize across datasets and modalities?
• RQ3 How does one-shot training of prompts and head compare to full fine-tuning or head-only tuning in medical segmentation?
• RQ4 What is the method’s performance and generalization on vessel, lesion, and retinal layer segmentation tasks?

Key findings

• The learnable ophthalmology prompt yields competitive or improved segmentation performance on several tasks with one-shot fine-tuning.
• In zero-shot tests, prompts trained on one dataset generalize well to other two-modal vessel segmentation datasets and to OCTA datasets from different manufacturers.
• On OCT retinal layer segmentation (ARoI), Dice improves by about 25% over the baseline SAM.
• The method can segment large vessels in OCTA and perform retinal layer segmentation more effectively than SAM alone in several cases.
• Small or tiny targets remain challenging for SAM-based prompts, indicating a need for further prompt enhancement for tiny objects.