Learning to Prompt Segment Anything Models

Jiaxing Huang,Kai Jiang,Jingyi Zhang,Han Qiu,Lewei Lu,Shijian Lu,Eric Xing

2024-01-10

Abstract:Segment Anything Models (SAMs) like SEEM and SAM have demonstrated great potential in learning to segment anything. The core design of SAMs lies with Promptable Segmentation, which takes a handcrafted prompt as input and returns the expected segmentation mask. SAMs work with two types of prompts including spatial prompts (e.g., points) and semantic prompts (e.g., texts), which work together to prompt SAMs to segment anything on downstream datasets. Despite the important role of prompts, how to acquire suitable prompts for SAMs is largely under-explored. In this work, we examine the architecture of SAMs and identify two challenges for learning effective prompts for SAMs. To this end, we propose spatial-semantic prompt learning (SSPrompt) that learns effective semantic and spatial prompts for better SAMs. Specifically, SSPrompt introduces spatial prompt learning and semantic prompt learning, which optimize spatial prompts and semantic prompts directly over the embedding space and selectively leverage the knowledge encoded in pre-trained prompt encoders. Extensive experiments show that SSPrompt achieves superior image segmentation performance consistently across multiple widely adopted datasets.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper attempts to address the issue of how to obtain suitable prompts in Segment Anything Models (SAMs) to improve the performance of image segmentation tasks. Specifically, the paper identifies two main challenges:

1. **Limited search space in spatial prompt learning**: SAMs are constrained by the 2D coordinate system when optimizing spatial prompts, leading to a limited search space and potentially suboptimal spatial prompts.
2. **Side effects of text prompt encoders**: The text prompt encoders in SAMs (such as CLIP, UniCL/Florence) are primarily pre-trained on image-text data centered around objects, which results in good performance for foreground object descriptions but relatively weaker performance for background descriptions. Therefore, when using these encoders for semantic prompt learning, they may benefit from the knowledge they have already learned but may also be negatively affected by the knowledge they have not fully learned.

To address these issues, the paper proposes the **Spatial-Semantic Prompt Learning (SSPrompt)** method, which directly optimizes spatial and semantic prompts in the embedding space and selectively leverages the knowledge from pre-trained prompt encoders. This approach allows for effective prompt learning with a small amount of labeled data, improving the segmentation performance of SAMs on downstream datasets. Extensive experimental validation shows that SSPrompt performs excellently on multiple widely-used datasets, achieving state-of-the-art performance.