SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lei Ke\*1,2 Mingqiao Ye\*1 Martin Danelljan Yifan Liu Yu-Wing Tai3 Chi-Keung Tang2 Fisher Yu1 ETH Zürich 2HKUST 3Dartmouth College
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.
Accurate segmentation of diverse objects is fundamental for a wide range of scene understanding applications, including image/video editing, robotic perception, and AR/VR. Trained with billion-scale mask labels, the Segment Anything Model (SAM) [21] was recently released as a foundational vision model for general image segmentation. SAM is capable of segmenting a wide range of objects, parts, and visual structures in diverse scenarios, by taking a prompt consisting of points, a bounding box, or a coarse mask as input. Its zero-shot segmentation abilities have led to a rapid paradigm shift, as it can be transferred to numerous applications through simple prompting.
While SAM has achieved impressive performance, its segmentation results are still unsatisfactory in many cases. In particular, SAM suffers from two key problems: 1) Coarse mask boundaries, often even neglecting the segmentation of thin object structures, as shown in Figure 1. 2) Incorrect predictions, broken masks, or large errors in challenging cases. This is often related to SAM misinterpreting thin structures, such as the kite lines in the rightmost column of Figure 1. These types of failures severely limit the applicability and effectiveness of foundational segmentation models, such as SAM, in particular for automated annotation and image/video editing tasks, where highly accurate image masks are crucial.
We propose HQ-SAM, which can predict highly accurate segmentation masks, even in very challenging cases (see Figure 1), without compromising the strong zero-shot capabilities and flexibility of the
SAM Prediction
HQ-SAM Prediction
Figure 1: The predicted masks of SAM vs. our HQ-SAM, given the same red box or several points on the object as input prompts. HQ-SAM produces significantly more detailed results with very accurate boundaries. In the rightmost column, SAM misinterprets the thin structure of the kite lines, and produces a large portion of errors with broken holes for the input box prompt.
original SAM. To preserve the efficiency and zero-shot performance, we propose a minimal adaptation of SAM, adding less than  $0.5\%$  parameters, to extend its capability to high-quality segmentation.
Directly fine-tuning the SAM decoder or introducing a new decoder module severely degrades the general zero-shot segmentation performance. We therefore propose the HQ-SAM architecture, which tightly integrates with and re-uses the existing learned SAM structure, in order to fully preserve the zero-shot performance. First, we design a learnable HQ-Output Token that is input to SAM's mask decoder, alongside the original prompt and output tokens. Unlike the original output tokens, our HQ-Output Token and its associated MLP layers are trained to predict a high-quality segmentation mask. Second, instead of only re-using the SAM's mask decoder features, our HQ-Output Token operates on a refined feature set to achieve accurate mask details. In particular, we use both global semantic context and local fine-grained features by fusing SAM's mask decoder features with early and late feature maps from its ViT encoder. During training, we freeze the entire pre-trained SAM parameters, while only updating our HQ-Output Token, its associated three-layer MLPs, and a small feature fusion block.
Learning accurate segmentation requires a dataset with accurate mask annotations of diverse objects with complex and detailed geometries. SAM is trained on the SA-1B dataset, which contains 11M images with 1.1 billion masks automatically generated by a SAM-like model. However, using this extensive dataset presents significant cost implications and falls short of achieving the desired high-quality mask generations pursued in our work, as evident by SAM's performance in Figure 1. Consequently, we compose a new dataset, called HQSeg-44K, which contains 44K extremely fine-grained image mask annotations. HQSeg-44K is constructed by merging six existing image datasets [35, 29, 26, 38, 8, 46] with highly accurate mask labels, covering over 1,000 diverse semantic classes. Thanks to the smaller-scale dataset and our minimal integrated architecture, HQ-SAM can be trained in only 4 hours on 8 RTX 3090 GPUs.
Figure 2: Performance vs. speed vs. model size for an array of SAM variants [21, 52].
To validate the effectiveness of HQ-SAM, we perform extensive quantitative and qualitative experimental analysis. We provide a comprehensive performance-speed-model size comparison on SAM variants [21, 52] in Figure 2. We compare HQ-SAM with SAM on a suite of 10 diverse segmentation datasets across different downstream tasks, where 8 out of them are under a zero-shot transfer protocol, including COCO [31], UVO [42], SGinW [58], LVIS [14], HQ-YTVIS [20], BIG [6], COIFT [29]
and HR-SOD [51]. This rigorous evaluation demonstrates that the proposed HQ-SAM can produce higher-quality masks while maintaining the zero-shot capability compared with SAM.
High-quality Segmentation Existing works for high-quality segmentation are mostly trained for a specific segmentation task, like image and video instance segmentation [22, 19, 20, 40, 44], semantic segmentation [30, 54, 39, 50] or panoptic segmentation [9], in a close-world paradigm. Some of them focus on post-segmentation refinement using with graphical models such as CRF [23] or region growing [10]. However, the CRF-based refinement is adhere to low-level color boundaries without fully utilizing high-level semantic context and cannot fix large segmentation errors. While some refinement-based works adopt separate deep networks for cascade iterative refinement [6, 37], they are prone to overfitting as shown by our experiment. Compared to these high-quality segmentation [19, 22, 33] or segmentation refinement methods, we focus on accurately segmenting diverse objects on new data with flexible prompting, and build a high-quality zero-shot segmentation model that generalizes to various segmentation tasks and domains. Unlike the post segmentation refinement works [6, 37], to preserve the zero-shot segmentation capability of SAM, HQ-SAM predicts the new high-quality mask directly by reusing the image encoder and mask decoder of SAM, instead of taking the coarse mask and images as the input and feeding it into a separate refinement network. The model architecture of HQ-SAM builds upon SAM with negligible overhead, where we propose efficient token learning for accurate mask predictions. This is completely different from previous high-quality segmentation works, and we show its effectiveness across a wide range of zero-shot experiments.
Fine-tuning and Prompt Tuning for Foundation Models Foundation models [2, 1] first appear in the NLP community, where large language models such as GPT series [2] show strong zero-shot generalization to unseen tasks and data. Then, some prompt-based learning works [16, 27, 17] are proposed to help these pre-trained models generalize to the downstream tasks instead of fine-tuning the internal model parameters [15] for better transfer learning. For vision-based foundation models [21, 43, 59], prompt engineering [56, 45, 49, 57] that freezes the pre-trained model is first explored in vision-language models, such as CLIP [36]. These prompts with learnable parameters are designed to help downstream tasks with better context optimization. Different from the existing prompt-based or finetuning works, we focus on the minimal adaptation of SAM toward high-quality segmentation. We directly use the proposed HQ-Output Token output for accurate mask prediction, instead of only leveraging some learnable parameters [56] to help context learning and better generalization.
We propose HQ-SAM to upgrade SAM for high-quality zero-shot segmentation. HQ-SAM is lightweight and only introduces two important adaptations to the SAM model. In Sec 3.1, we first briefly review the architecture of SAM on which HQ-SAM is built. Then, in Sec 3.2, we introduce our HQ-SAM with High-Quality Token (HQ-Output Token) and Global-local Feature Fusion, which are the key components to achieve better segmentation quality for SAM while preserving its zero-shot capability. Finally, in Sec 3.3, we describe the training and inference process of HQ-SAM, which is both data and computationally efficient.
3.1 Preliminaries: SAM
SAM [21] is composed of three modules: (a) Image encoder: a heavy ViT-based backbone for image feature extraction, resulting in image embedding in spatial size  $64 \times 64$ . (b) Prompt encoder: encoding the interactive positional information from the input points/boxes/masks to provide for the mask decoder. (c) Mask decoder: a two-layer transformer-based decoder takes both the extracted image embedding with the concatenated output and prompt tokens for final mask prediction. The released SAM model is trained on the large-scale SA-1B dataset, which contains over 1 billion automatically generated masks ( $400 \times$  more masks than any existing segmentation datasets [14, 24]) and 11 million images. Thus, SAM shows valuable strong zero-shot generalization to new data without the necessity for additional training. However, we also note that SAM training is very expensive, where distributively training ViT-H-based SAM for 2 epochs on SA-1B requires 256 GPUs with a large batch size of 256 images. For more SAM method details, we refer readers to [21].
Figure 3: HQ-SAM introduces HQ-Output Token and Global-local Feature Fusion to SAM for high-quality mask prediction. To keep the zero-shot capability of SAM, the lightweight HQ-Output Token reuses SAM's mask decoder, and generates new MLP layers for performing point-wise product with fused HQ-Features. During training, only a few learnable parameters in HQ-SAM are trainable while we fix the model parameters of the pre-trained SAM. The prompt encoder is omitted here for clarity. Error correction is simply used as a direct element-wise sum between the predicted logits of the SAM's Output Token and the HQ-Output Token during inference.
3.2 Ours: HQ-SAM
In this section, we describe the architecture of the HQ-SAM network. To preserve the zero-shot transfer capability of SAM, while preventing model overfitting or catastrophic forgetting, instead of directly finetuning SAM or adding a new heavy decoder network, we take a minimal adaptation approach as much as possible. To this end, HQ-SAM reuses the pre-trained model weights of SAM as much as possible with only two new key components, namely, High-Quality Output Token and Global-local Feature Fusion, as illustrated in Figure 3. HQ-SAM can thus be regarded as a high-quality zero-shot segmentation model evolved from SAM with negligible extra model parameters and computation cost.
3.2.1 High-Quality Output Token
We propose efficient token learning for improving the mask quality of SAM. As shown in Figure 3, in SAM's original mask decoder design, the output token (similar to object query in DETR [3]) is adopted for mask prediction, which predicts dynamic MLP weights and then performs point-wise product with the mask features. To promote SAM's mask quality in HQ-SAM, instead of directly taking SAM's coarse masks as input, we introduce the HQ-Output token and a new mask prediction layer for high-quality mask prediction.
In Figure 3, by reusing and fixing SAM's mask decoder, a new learnable HQ-Output Token (size of  $1 \times 256$ ) is concatenated with SAM's output tokens (size of  $4 \times 256$ ) and prompt tokens (size of  $N_{prompt} \times 256$ ) as the input to the SAM's mask decoder. Similar to the original output token, in each attention layer, HQ-Output Token first performs self-attention with other tokens and then conducts both token-to-image and the reverse image-to-token attention for its feature updating. Note that HQ-Output Token uses the point-wise MLP shared by the other tokens in each decoder layer. After passing through two decoder layers, the updated HQ-Output Token has access to the global image context, the critical geometric/type information of prompt tokens as well as hidden mask information of the other output tokens. Finally, we add a new three-layer MLP to generate dynamic convolutional kernels from the updated HQ-Output Token, which then performs spatially point-wise product with the fused HQ-feature for high-quality mask generation.
Instead of directly finetuning SAM or further adding a heavy post-refinement network, we only allow the HQ-Output Token and its associated three-layer MLPs to be trained for correcting the mask errors of SAM's output token. This is completely different from existing high-quality segmentation models [19, 6, 20, 22]. We identify two main advantages of our efficient token learning through extensive experiments: 1) This strategy significantly improves SAM's mask quality while only
introducing negligible parameters compared to original SAM, making HQ-SAM training extremely time and data-efficient; 2) The learned token and MLP layers do not overfit to mask the annotation bias of a specific dataset, thus keeping SAM's strong zero-shot segmentation capability on new images without catastrophic knowledge forgetting.
3.2.2 Global-local Fusion for High-quality Features
Very accurate segmentation also requires input image feature with both rich global semantic context and local boundary details. To further promote mask quality, we enrich both the high-level object context and low-level boundary/edge information in the mask decoder features of SAM. Instead of directly using SAM's mask decoder feature, we compose the new high-quality features (HQ-Features) by extracting and fusing features from different stages of the SAM model: 1) The early layer local feature of SAM's ViT encoder with spatial shape  $64 \times 64$ , which captures more general image edge/boundary details [12]. Concretely, we extract the feature after the first global attention block of the ViT encoder, and for ViT-Large based SAM, this is the 6th block output for the 24 blocks in total; 2) The final layer global feature of SAM's ViT encoder with shape  $64 \times 64$ , which has more global image context information; 3) The mask feature in SAM's mask decoder with size  $256 \times 256$ , which is also shared by the output tokens, contains strong mask shape information.
As shown in Figure 3, to obtain the input HQ-Features, we first upsample the early-layer and final-layer encoder features to the spatial size  $256 \times 256$  by transposed convolution. Then, we sum up these three types of features in an element-wise manner after simple convolutional processing. We show that this global-local feature fusion is simple while effective, yielding detail-preserving segmentation results with a small memory footprint and computation burden. We also perform detailed ablation on the effect of each feature source in the experimental section (Table 3).
3.3 Training and Inference of HQ-SAM
Training Data Construction To train HQ-SAM in a data-efficient manner, instead of further training on SA-1B [21], we compose a new training dataset HQSeg-44K which contains 44,320 extremely accurate image mask annotations. We note that the released SA-1B dataset only contains automatically generated mask labels, missing very accurate manual annotation on objects with complex structures. Due to the annotation difficulty, HQSeg-44K leverages a collection of six existing image datasets including DIS [35] (train set), ThinObject-5K [29] (train set), FSS-1000 [26], ECSSD [38], MSRA-10K [8], DUT-OMRON [46] with extremely fine-grained mask labeling, where each of them contains 7.4K mask labels on average. To make HQ-SAM robust and generalizable to new data, HQSeg-44K contains diverse semantic classes of more than 1,000. We show the advantage of using HQSeg-44K by comparing HQ-SAM training with 44K randomly sampled images and masks from SA-1B [21] in our supplemental analysis.
HQ-SAM Training During training, we fix the model parameters of the pre-trained SAM model while only making the proposed HQ-SAM learnable. The learnable parameters thus only include the HQ-Output Token, its associated three-layer MLP and three simple convolutions for HQ-Features fusion. Since SAM is designed for flexible segmentation prompts, we train HQ-SAM by sampling mixed types of prompts including bounding boxes, randomly sampled points, and coarse masks input. We generate these degraded masks by adding random Gaussian noise in the boundary regions of the GT masks. For generalizability to different object scales, we use large-scale jittering [13]. We use a learning rate of 0.001 and train our HQ-SAM for 12 epochs, with a learning rate drop after 10 epochs. We train on 8 Nvidia GeForce RTX 3090 GPUs with a total batch size of 32, which takes 4 hours to train for 16.6K iterations. Please refer to our supplemental file for more details.
HQ-SAM Inference We follow the same inference pipeline of SAM but use the mask prediction from HQ-Output token as high-quality mask prediction. During inference, we sum the predicted logits of the SAM mask (by Output Token) and our predicted mask (by HQ-Output Token) for mask correction on spatial resolution  $256 \times 256$ . Then we up-sample the corrected mask to the original resolution  $1024 \times 1024$  as our output.
SAM vs. HQ-SAM on Training and Inference In Table 1, we report detailed training and inference comparisons between our HQ-SAM and SAM. While HQ-SAM produces substantially better segmentation quality, its training is very quick and affordable, which only takes 4 hours with 8 RTX3090 GPUs. HQ-SAM is also lightweight and efficient, introducing negligible increases in model parameters, GPU memory usage, and inference time per image.
Table 1: Training and inference comparison between ViT-L [11] based SAM and HQ-SAM. HQ-SAM brings negligible extra computation burden to SAM, with less than  $0.5\%$  increase in model parameters and reaching  $96\%$  of its original speed. SAM-L is trained on 128 A100 GPUs for 180k iterations. Based on SAM-L, we only need to train our HQ-SAM on 8 RTX3090 GPUs for 4 hours.
Method	Learnable Params (M)	Training # GPU	Batch Size	Time (h)	Inference
FPS	Mem.
SAM [21]	1191	128	128	N/A	5.0	7.6G
HQ-SAM	5.1	8	32	4	4.8	7.6G
Datasets For training we use the compiled HQSeg-44K, described in Section 3.3. For a comprehensive evaluation of the segmentation performance of HQ-SAM, we perform experiments on a wide range of datasets, including four extremely fine-grained segmentation datasets: DIS [35] (validation set), ThinObject-5K [29] (test set), COIFT [29] and HR-SOD [51]. Besides, we experiment on popular and challenging benchmarks across various image/video-based segmentation tasks in zero-shot settings, such as COCO [31], SGinW [58], UVO [42], LVIS [14], HQ-YTVIS [20] and BIG [6].
Evaluation Metrics To accurately quantify improvements in mask quality, instead of only employing the standard mask AP or mask mIoU, we also adopt boundary metrics mBIoU and boundary  $\mathrm{AP}_B$  [5]. We also evaluate on stricter  $\mathrm{AP}_B^{\mathrm{strict}}$  by adjusting the default dilation ratio from 0.02 to 0.01 on UVO [42] and LVIS [14]. For evaluation on the four fine-grained segmentation datasets [35, 29, 51], we also report the averaged boundary and mask IoU among them. For video instance segmentation evaluation on HQ-YTVIS [20], we use both Tube Boundary  $\mathrm{AP}^B$  and Tube Mask  $\mathrm{AP}^M$ .
4.2 Ablation Experiments
We conduct detailed ablation studies on the proposed HQ-SAM using ViT-Large as the backbone, analyzing the impact of the proposed HQ-Output Token and HQ-Features on segmentation quality especially in zero-shot cases. For ablation experiments, we use the four aforementioned extremely accurate segmentation datasets, namely, DIS (val) [35], ThinObject-5K (test) [29], COIFT [29] and HR-SOD [51] as well as the COCO validation set.
Effect of the High-Quality Output Token. HQ-SAM employs HQ-Output Token for high-quality mask prediction. Table 2 compares our HQ-Output Token to the baseline SAM and other existing prompt-token learning strategies, such as adding an additional three context tokens [56] as learnable vectors into the SAM's mask decoder for better context learning. Compared to using context tokens, the HQ-Output token consistently brings larger performance gains on four high-quality datasets, with  $13.2\mathrm{mBIoU}$  on DIS and  $2.7\mathrm{mBIoU}$  on COIFT datasets. We also perform other ablation experiment variants, such as computing the scaled dot product [18] between the original SAM's output token and our HQ-Output token or restricting the mask loss to only inside the boundary regions, and find they slightly decrease the averaged performance on the four evaluation datasets. Compared to SAM, HQ-SAM significantly improves the mBIoU on DIS benchmark from 52.8 to 70.4 and also promotes the mBIoU on the HRSOD dataset for 3.8 points.
Ablation on the Global-local Fusion for HQ-Features Table 3 tabulates the effect of global-local fusion, where the importance of each feature component is analyzed in HQ-Features during the fusion process. Compared to directly using the mask decoder feature of SAM, the entire HQ-Features bring an obvious advantage of  $2.6\mathrm{mBiOU}$  on four highly accurate segmentation datasets. The final-layer ViT encoder feature with global context increases the mBiOU from 80.1 to 81.3 while the early-layer feature with local details further promotes the mBiOU to 81.8. We also replace the proposed global-local fusion with the conventional FPN to build a feature pyramid for fusion, and found this brought an inferior performance, decreasing from 89.1 to 87.4 mIoU.
Comparison to SAM finetuning or post-refinement. In Table 4, we compare our efficient token adaptation strategy to adding an extra post-refinement network [6] and model finetuning, including directly finetuning SAM's mask decoder or only finetuning its output token for mask prediction. Adding an extra heavy post-refinement network brings limited averaged performance increase on