SqueezeSAM: User friendly mobile interactive segmentation

TL;DR

SqueezeSAM addresses the need for fast, on-device interactive segmentation by replacing SAM's heavy ViT-H encoder with a compact UNet-based encoder-decoder that retains segmentation quality. It leverages salient object detection to generate initial prompts and uses an early fusion strategy to incorporate user clicks, plus data augmentation to enable whole-object segmentation. The model is 62.5× faster and 31.6× smaller than SAM, achieving mIOU within 1% of the original on benchmarks, and runs in real time on iPhone CPUs with potential for 14MB int8 deployment. The approach enables automatic segmentation for mobile photography workflows and demonstrates improved user experience in whole-object extraction without extensive user input.

Abstract

The Segment Anything Model (SAM) has been a cornerstone in the field of interactive segmentation, propelling significant progress in generative AI, computational photography, and medical imaging. Despite its ability to process arbitrary user input and generate corresponding segmentation masks, SAM's 600 million parameter architecture, based on ViT-H, is not compatible with current mobile hardware due to its high computational demands and large model size. Our research aims to adapt SAM for use in mobile photography applications. To this end, we have developed a fully convolutional SqueezeSAM model architecture, which is 62.5 times faster and 31.6 times smaller than the original SAM, making it a viable solution for mobile applications. Furthermore, our tiny model achieves an mIOU within 1% of the original VIT-H architecture. Automated segmentation holds significant value in the creation flow for photography applications, as evidenced by its adoption by leading industry players like apple and capcut. To facilitate this automation, we employ salient object detection and simulate potential user clicks for foreground object selection, generating an initial segmentation mask that users can subsequently edit interactively. A common user expectation is that a click on a specific part of an object will result in the segmentation of the entire object. For example, a click on a person's t-shirt in a photo should ideally segment the entire person, not just the t-shirt. However, SAM typically only segments the clicked area. We address this limitation through a novel data augmentation scheme. Consequently, if a user clicks on a person holding a basketball, both the person and the basketball are segmented together, aligning with user expectations and enhancing the overall user experience.

Figures (9)

• Figure 1: A comparative analysis of SqueezeSAM and SAM. (Top) Speed comparison between SqueezeSAM and SAM on a single NVIDIA A100 GPU. (Bottom) Performance/Runtime/Parameter comparison of SAM and other efficient models.
• Figure 2: Salient SqueezeSAM pipeline. Salient Object Detection feeds initial points to SqueezeSAM, so the user can get an initial segmentation without clicking on the screen. After that, the user can interactively edit and improve the segmentation. It runs in real-time on an iPhone.
• Figure 3: SqueezeSAM model architecture. Our proposed SqueezeSAM can take point coordinates from user's clicks, or "guess" user's intention from a saliency mask, and generate the corresponding segmentation output. For example, the model can begin by segmenting one object, and the user can interactively edit the segmentation or segment more objects.
• Figure 4: SqueezeSAM encoder-decoder model architecture. DoubleConv modules are described in later figures. iC = input channels from the previous layer, not including concatenated skip-connections. oC = output channels. DS = downsampling factor, relative to the original input image.
• Figure 5: DoubleConvDown module used in the SqueezeUNet encoder.