PaperlayerView exact revision on arXiv

Source-linked AI summary

Segment Anything in High Quality

Lei Ke, Mingqiao Ye, Martin Danelljan, Yifan Liu, Yu-Wing Tai, Chi-Keung Tang, Fisher Yu

arXiv:2306.01567v2cs.CV

TL;DR

SAM’s zero-shot, promptable segmentation remains unreliable for thin structures and intricate object boundaries. HQ-SAM adds a lightweight output token and global-local feature fusion, trained on 44K fine-grained masks while reusing SAM’s pretrained model. Across 10 diverse datasets, including 8 zero-shot evaluations, it improves mask quality while preserving zero-shot capability and efficiency.

  • Problem

    SAM often produces coarse boundaries, misses thin structures, and generates incorrect or broken masks despite its broad zero-shot segmentation ability.

  • Method

    HQ-SAM minimally adapts SAM with a High-Quality Output Token and global-local feature fusion, training only introduced parameters on HQSeg-44K fine-grained masks.

  • Results

    HQ-SAM produces higher-quality masks while maintaining zero-shot capability across 10 diverse segmentation datasets, including 8 zero-shot transfer evaluations.

  • Takeaways & Limitations

    HQ-SAM extends SAM toward high-quality zero-shot segmentation using 44K accurate masks and negligible additional model overhead.

  • Takeaways & Limitations

    HQ-SAM shares SAM’s heavy ViT encoder and therefore cannot achieve real-time video-processing speed without the Light HQ-SAM variant.

Abstract

from arXiv · show

The recent Segment Anything Model (SAM) represents a big leap in scaling up segmentation models, allowing for powerful zero-shot capabilities and flexible prompting. Despite being trained with 1.1 billion masks, SAM's mask prediction quality falls short in many cases, particularly when dealing with objects that have intricate structures. We propose HQ-SAM, equipping SAM with the ability to accurately segment any object, while maintaining SAM's original promptable design, efficiency, and zero-shot generalizability. Our careful design reuses and preserves the pre-trained model weights of SAM, while only introducing minimal additional parameters and computation. We design a learnable High-Quality Output Token, which is injected into SAM's mask decoder and is responsible for predicting the high-quality mask. Instead of only applying it on mask-decoder features, we first fuse them with early and final ViT features for improved mask details. To train our introduced learnable parameters, we compose a dataset of 44K fine-grained masks from several sources. HQ-SAM is only trained on the introduced detaset of 44k masks, which takes only 4 hours on 8 GPUs. We show the efficacy of HQ-SAM in a suite of 10 diverse segmentation datasets across different downstream tasks, where 8 out of them are evaluated in a zero-shot transfer protocol. Our code and pretrained models are at https://github.com/SysCV/SAM-HQ.

1 Introduction

SAM offers flexible, zero-shot segmentation but often produces coarse boundaries and incorrect masks for intricate structures. HQ-SAM addresses these failures with a minimal adaptation, and evaluation across diverse datasets shows improved mask quality while preserving zero-shot capability.

  • SAM supports zero-shot segmentation from points, bounding boxes, or coarse masks across diverse objects and scenarios.
  • SAM often produces coarse boundaries, misses thin structures, and generates broken or substantially erroneous masks in challenging cases.
  • HQ-SAM targets these failures while preserving SAM’s zero-shot capabilities and flexible prompting.
  • Less than 0.5% additional parameters enable HQ-SAM to extend SAM toward high-quality segmentation with minimal adaptation.
  • HQ-SAM is evaluated against SAM on 10 diverse segmentation datasets, including 8 zero-shot transfer evaluations.
  • HQ-SAM produces higher-quality masks while maintaining zero-shot capability compared with SAM.

2 Related Work

Prior high-quality segmentation methods largely target specific tasks or rely on refinement techniques, while foundation-model research emphasizes prompt-based adaptation for transfer.

  • Existing high-quality segmentation methods are mostly trained for specific tasks in close-world settings.
  • Some prior methods refine segmentation outputs with graphical models such as CRF or region growing.
  • Prompt-based learning adapts pre-trained foundation models to downstream tasks without fine-tuning their internal parameters.

3 Method

HQ-SAM minimally adapts SAM by adding a High-Quality Output Token and fusing global and local features, while preserving the pretrained model and promptable inference pipeline. It trains only the introduced components on fine-grained masks and adds little computational overhead.

  • 3 Method: HQ-SAM reuses SAM’s pretrained weights and adds two key components: a High-Quality Output Token and Global-local Feature Fusion.
  • 3.2 HQ-SAM: The HQ-Output Token enters SAM’s mask decoder alongside prompt and output tokens, then predicts high-quality masks through associated MLP layers.
  • 3.2.2 Global-local Fusion for High-quality Features: HQ-Features fuse early-layer local details, final-layer global context, and SAM mask-decoder shape information.
  • 3.3 Training and Inference of HQ-SAM: HQSeg-44K contains 44,320 fine-grained image-mask annotations collected from six existing datasets without overlap with zero-shot evaluation datasets.
  • 3.3 Training and Inference of HQ-SAM: During training, SAM’s parameters remain fixed while only the HQ-Output Token, its three-layer MLP, and fusion convolutions are learned.
  • 3.3 Training and Inference of HQ-SAM: Inference sums SAM and HQ-SAM mask logits for correction before upsampling the result to the original resolution.
  • 3.3 Training and Inference of HQ-SAM: Less than 0.5% more parameters and 96% of SAM’s original speed provide substantially better segmentation quality with negligible overhead.

4 Experiments

Experiments show that HQ-SAM improves fine-grained mask quality and zero-shot transfer across diverse image and video segmentation tasks, while retaining efficient adaptation and mobile deployment options.

  • Evaluation Setup: HQ-SAM improves segmentation quality across 10 diverse datasets, including eight evaluated under zero-shot transfer.The evaluation spans fine-grained image, open-world, instance, and video segmentation benchmarks.
  • Ablation Experiments: HQ-Output Token ablations raise mBIoU on DIS from 52.8 to 70.4 and improve HRSOD by 3.8 points over SAM.The HQ-Output Token also produces larger gains than context-token alternatives on four high-quality datasets.
  • Ablation Experiments: Global-local HQ-Features improve averaged mBIoU by 2.6 points, with final-layer and early-layer features increasing it from 80.1 to 81.8.The final-layer feature contributes global context, while the early-layer feature adds local details.
  • Ablation Experiments: HQ-SAM outperforms direct decoder fine-tuning and heavy post-refinement on COCO while avoiding the poor zero-shot performance associated with overfitting.HQ-SAM improves COCO APB by 1.1, compared with 0.4 from output-token-only fine-tuning.
  • Zero-shot Comparison with SAM: Compared with SAM, HQ-SAM consistently improves boundary-sensitive metrics, including 1.1 APB on COCO, 0.7 APstrict B75 on LVIS, and 3.8 Tube Boundary APB on HQ-YTVIS.HQ-YTVIS also gains 2.9 Tube Mask APM.
  • Mobile Efficiency: Light HQ-SAM reaches 41.2 FPS and raises MobileSAM’s zero-shot COCO AP from 44.3 to 45.0 with a 1.7MB parameter increase.This variant targets efficient mobile deployment using a tiny ViT image encoder.

5 Conclusion

HQ-SAM adds high-quality zero-shot segmentation to SAM with negligible overhead, using a lightweight output-token design trained on 44K accurate masks. The method is evaluated across eight zero-shot benchmarks spanning image and video tasks.

  • HQ-SAM introduces a lightweight High-quality Output Token that replaces SAM’s original output token for high-quality mask prediction.
  • HQ-SAM is trained on 44K highly accurate masks while retaining the original SAM model’s zero-shot segmentation capabilities.
  • HQ-SAM evaluates zero-shot transfer on 8 segmentation benchmarks covering both image and video tasks.
  • The work targets data-efficient and computation-affordable ways to extend SAM-like foundational segmentation models.

6 Supplementary experiments

Supplementary experiments show that HQ-SAM consistently improves SAM across backbones and zero-shot video settings, including higher-quality contour and tube-mask measurements. The method also remains more robust when box prompts contain noise.

  • SAM vs. HQ-SAM on Various Backbones: HQ-SAM consistently outperforms SAM across ViT-B, ViT-L, ViT-H, and TinyViT backbones, with over 10 points higher mBIoU on four HQ datasets.
  • SAM vs. HQ-SAM on Various Backbones: On COCO with ViT-B, HQ-SAM raises APB from 28.2 to 31.3 and AP from 44.4 to 46.7 with a 1.1% parameter increase.
  • Zero-shot Video Instance Segmentation Comparison: HQ-SAM increases Tube Mask AP on YTVIS 2019 from 51.8 to 53.2, a gain of 1.4 points.
  • Zero-shot Video Object Segmentation Comparison: On DAVIS validation data, HQ-SAM improves J &F from 82.0 to 83.2 and F from 84.9 to 86.1.
  • Robustness to Input Box Prompts: HQ-SAM is more robust than SAM to noisy input box prompts, with its relative mBIoU advantage increasing across noise scales.

7 Additional Implementation details

Training freezes SAM’s pretrained parameters and learns only the added HQ-SAM components. Inference follows SAM’s pipeline while using the HQ-Output Token for the high-quality mask.

  • Training Details: Training freezes SAM’s pretrained parameters and learns the HQ-Output Token, its three-layer MLP, and three HQ-Feature-fusion convolutions.
  • Training Details: Two transposed convolutions with 2×2 kernels and stride 2 upscale encoder embeddings from 64×64 to 256×256.
  • Training Details: HQ-SAM treats the new HQ-Output Token as a fifth mask token alongside SAM’s original four mask tokens.
  • Implementation Details: At inference, HQ-SAM follows SAM’s pipeline and uses the mask predicted by the HQ-Output Token as the high-quality output.
  • Implementation Details: Interactive evaluation uses the ground-truth mask’s center point for single-point prompts and randomly sampled ground-truth points for multi-point evaluation.

8 More Details of HQSeg-44K

HQSeg-44K provides 44,320 fine-grained masks from six datasets, without overlap with the zero-shot evaluation data. Experiments attribute substantial gains to this curated training set and show improvements remain when related training splits are removed.

  • Data compostion of HQSeg-44K: HQSeg-44K contains 44,320 extremely accurate image-mask annotations collected from six existing datasets.
  • Data compostion of HQSeg-44K: The six source datasets contribute approximately 7.4K mask labels each, covering DIS, ThinObject-5K, FSS, ECSSD, MSRA-10K, and DUT-OMRON.
  • Data compostion of HQSeg-44K: HQSeg-44K has no image or annotation overlap with the paper’s zero-shot evaluation datasets.
  • Effect of HQSeg-44K: Training with 44K SA-1B samples lowers averaged mBIoU from 71.1 to 70.1, whereas HQSeg-44K training raises it from 71.1 to 81.8 across four datasets.
  • Visual comparisons: Visual comparisons on DIS show more accurate boundaries for HQ-SAM than SAM under the same red-box prompt.
  • Zero-shot results on DIS and ThinObject-5K: Removing DIS or ThinObject-5K training splits still leaves HQ-SAM over 10.0 points higher on DIS-mIoU and 9.0 points higher on ThinObject-mIoU.

9 More Visual Results Comparison

HQ-SAM is compared visually with SAM across image and video segmentation settings, including zero-shot prompts and noisy boxes. The comparisons show improved mask quality, detail preservation, robustness, and fewer artifacts, while retaining some failures in extreme cases.

  • Overall comparison: Across the reported visual comparisons, HQ-SAM improves over SAM while preserving zero-shot segmentation power, but does not always produce fully correct masks.
  • Image segmentation: HQ-SAM produces higher-quality masks with fewer artifacts than SAM on COCO validation images in a zero-shot setting.The comparison uses box prompts generated by a SOTA FocalNet-DINO detector trained on COCO.
  • Image segmentation: HQ-SAM delivers more consistent and robust segmentation than SAM when the input boxes are ground-truth or noisy.
  • Image segmentation: HQ-SAM preserves more details and fixes broken-hole mask errors compared with SAM across zero-shot HRSOD, NDD20, and web images.Two failure cases remain in extremely dark environments and for very tiny metal rods.
  • Video segmentation: The visual comparisons also evaluate zero-shot video object segmentation on DAVIS 2017 using identical video-box prompts from pre-trained XMem.
  • Video segmentation: The visual comparisons evaluate zero-shot video instance segmentation on YTVIS 2019 using identical video-box prompts from pre-trained Mask2Former.
Loading 2306.01567v2…