Segment Anything in High Quality

TL;DR

HQ-SAM enhances the original Segment Anything Model by significantly improving mask quality for intricate objects, using minimal additional training and preserving its zero-shot capabilities across diverse datasets.

Contribution

The paper introduces HQ-SAM, a high-quality segmentation extension that maintains SAM's promptability and efficiency while significantly improving mask detail through a learnable output token.

Findings

HQ-SAM outperforms SAM in mask quality on multiple datasets.

HQ-SAM achieves high-quality segmentation with minimal additional training.

The method maintains zero-shot transferability across diverse tasks.

Abstract

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.