Task-Specific Adaptation of Segmentation Foundation Model via Prompt Learning

TL;DR

This paper introduces a prompt learning-based adaptation method for the Segment Anything Model to improve its performance on task-specific instance segmentation, especially for out-of-distribution objects, with more efficient training.

Contribution

It proposes a prompt learning module and a point matching module to customize SAM for specific segmentation tasks, addressing ambiguity and training efficiency issues.

Findings

Enhanced segmentation accuracy on customized scenarios

More efficient training process for task-specific adaptation

Improved boundary detail alignment in segmentation results

Abstract

Recently, foundation models trained on massive datasets to adapt to a wide range of tasks have attracted considerable attention and are actively being explored within the computer vision community. Among these, the Segment Anything Model (SAM) stands out for its remarkable progress in generalizability and flexibility for image segmentation tasks, achieved through prompt-based object mask generation. However, despite its strength, SAM faces two key limitations when applied to instance segmentation that segments specific objects or those in unique environments (e.g., task-specific adaptation for out-of-distribution objects) not typically present in the training data: 1) the ambiguity inherent in input prompts and 2) the necessity for extensive additional training to achieve optimal segmentation. To address these challenges, we propose a task-specific adaptation (i.e., customization) of the segmentation foundation model via prompt learning tailored to SAM. Our method involves a prompt learning module (PLM), which adjusts input prompts into the embedding space to better align with peculiarities of the target task, thereby enabling more efficient training. Furthermore, we introduce a point matching module (PMM) to enhance the feature representation for finer segmentation by ensuring detailed alignment with ground truth boundaries. Experimental results on various customized segmentation scenarios demonstrate the effectiveness of the proposed method.