U-SAM: Upgrade Segment Anything Model With Semantic-Aware and Memory-Efficient

Abstract

Segment Anything Model (SAM) has achieved remarkable success in the field of class-agnostic image segmentation by utilizing points or boxes as prompts. However, we identify two significant limitations when compared to traditional image segmentation models: (1) Trained in a category-agnostic interactive segmentation manner, SAM lacks the ability to discern object granularity and semantics, rendering it ineffective for traditional instance, semantic, and panoptic segmentation tasks. (2) SAM’s inefficient use of instance-independent visual features and tokens necessitates maintaining unique features and tokens for each instance, leading to excessive GPU memory consumption and diminished segmentation efficiency. To address these issues, we propose the Universal Segment Anything Model (U-SAM), a semantic-aware and memory-efficient segmentation model designed to perform both promptable and traditional segmentation tasks within a compact and unified framework. Specifically, U-SAM enhances SAM by integrating the Multi-Scale Semantic-Aware Image Encoder (S2IE), thus providing multi-scale semantic features for achieving traditional image segmentation tasks. Additionally, U-SAM is equipped with a Twin Token Mask Decoder (T2MD) which reduces GPU memory overhead by substituting replicated visual features with replicated tokens. Extensive experiments across interactive, instance, semantic, and panoptic segmentation demonstrate U-SAM’s promising results. Notably, U-SAM is 9× smaller and 10× faster than SAM, showing strong performance in zero-shot segmentation. Moreover, U-SAM surpasses the SOTA object-prompter-based model, RSPrompter, by achieving a 6.2% increase in PQ, operating 14× faster, and cutting training memory usage by 61%.