Package: segment-anything-hq
Segment Anything in High QualitySegment Anything in High QualityLei Ke, Mingqiao Ye, Martin Danelljan, Yifan Liu, Yu-Wing Tai, Chi-Keung Tang, Fisher YuETH Zurich & HKUSTWe propose HQ-SAM to upgrade SAM for high-quality zero-shot segmentation. Refer to ourpaperfor more details.Updates2023/07/25: Light HQ-SAM is inEfficientSAM seriescombining withGrounded SAM, thanks the support!2023/07/21: HQ-SAM is also in OpenXLab apps, thanks their support!:rocket::rocket: 2023/07/17: We releasedLight HQ-SAMusing TinyViT as backbone, for both fast and high-quality zero-shot segmentation, which reaches41.2 FPS. Refer toLight HQ-SAM vs. MobileSAMfor more details.:trophy::1st_place_medal: 2023/07/14: GroundedHQ-SAMobtains thefirst place:1st_place_medal: in theSegmentation in the Wildcompetition on zero-shot track (hosted inCVPR 2023 workshop), outperforming Grounded SAM. Refer to ourSGinW evaluationfor more details.:fire::fire: 2023/07/05: We releasedSAM tuning instuctionsandHQSeg-44K data.2023/07/04: HQ-SAM is adopted inSAM-PTto improve the SAM-based zero-shot video segmentation performance. Also, HQ-SAM is used inGrounded-SAM,Inpaint AnythingandHQTrack(2nd in VOTS 2023).2023/06/28: We released theONNX export scriptandcolab notebookfor exporting and using ONNX model.2023/06/23: Play with HQ-SAM demo at, which supports point, box and text prompts.2023/06/14: We released thecolab demoandautomatic mask generator notebook.2023/06/13: We released themodel checkpointsanddemo visualization codes.Visual comparison between SAM and HQ-SAMSAM vs. HQ-SAMIntroductionThe 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 9 diverse segmentation datasets across different downstream tasks, where 7 out of them are evaluated in a zero-shot transfer protocol.Quantitative comparison between SAM and HQ-SAMNote: For box-prompting-based evaluation, we feed SAM, MobileSAM and our HQ-SAM with the same image/video bounding boxes and adopt the single mask output mode of SAM.We provide comprehensive performance, model size and speed comparison on SAM variants:Various ViT backbones on COCO:Note: For the COCO dataset, we use a SOTA detector FocalNet-DINO trained on the COCO dataset as our box prompt generator.YTVIS and HQ-YTVISNote:Using ViT-L backbone. We adopt the SOTA detector Mask2Former trained on the YouTubeVIS 2019 dataset as our video boxes prompt generator while reusing its object association prediction.DAVISNote: Using ViT-L backbone. We adopt the SOTA model XMem as our video boxes prompt generator while reusing its object association prediction.Interactive segmentation comparison using various pointsNote:Using ViT-L backbone. On the high-quality COIFT (zero-shot) and DIS val set.InstallationThe code requirespython>=3.8, as well aspytorch>=1.7andtorchvision>=0.8. Please follow the instructionshereto install both PyTorch and TorchVision dependencies. Installing both PyTorch and TorchVision with CUDA support is strongly recommended.Clone the repository locally and install withgit clone https://github.com/SysCV/sam-hq.git
cd sam-hq; pip install -e .The following optional dependencies are necessary for mask post-processing, saving masks in COCO format, the example notebooks, and exporting the model in ONNX format.jupyteris also required to run the example notebooks.pip install opencv-python pycocotools matplotlib onnxruntime onnx timmExample conda environment setupcondacreate--namesam_hqpython=3.8-y
condaactivatesam_hq
condainstallpytorch==1.10.0torchvision==0.11.0cudatoolkit=11.1-cpytorch-cnvidia
pipinstallopencv-pythonpycocotoolsmatplotlibonnxruntimeonnxtimm# under your working directorygitclonehttps://github.com/SysCV/sam-hq.gitcdsam-hq
pipinstall-e.exportPYTHONPATH=$(pwd)Model CheckpointsThree HQ-SAM model versions of the model are available with different backbone sizes. These models can be instantiated by runningfrom segment_anything import sam_model_registry
sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")Download the provided trained model below and put them into the pretrained_checkpoint folder:mkdir pretrained_checkpointClick the links below to download the checkpoint for the corresponding model type. We also providealternative model downloading linkshereor athugging face.vit_b:ViT-B HQ-SAM model.vit_l:ViT-L HQ-SAM model.vit_h:ViT-H HQ-SAM model.vit_tiny(Light HQ-SAMfor real-time need):ViT-Tiny HQ-SAM model.Getting StartedFirst download amodel checkpoint. Then the model can be used in just a few lines to get masks from a given prompt:from segment_anything import SamPredictor, sam_model_registry
sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
predictor = SamPredictor(sam)
predictor.set_image(<your_image>)
masks, _, _ = predictor.predict(<input_prompts>)Additionally, see the usage examples in ourdemo,colab notebookandautomatic mask generator notebook.To obtain HQ-SAM's visual result:python demo/demo_hqsam.pyTo obtain baseline SAM's visual result. Note that you need to download original SAM checkpoint frombaseline-SAM-L modeland put it into the pretrained_checkpoint folder.python demo/demo_sam.pyTo obtain Light HQ-SAM's visual result:python demo/demo_hqsam_light.pyHQ-SAM Tuning and HQ-Seg44k DataWe provide detailed training, evaluation, visualization and data downloading instructions inHQ-SAM training. You can also replace our training data to obtain your own SAM in specific application domain (like medical, OCR and remote sensing).Please change the current folder path to:cd trainand then refer to detailedreadme instruction.Grounded HQ-SAM vs Grounded SAM onSegInWGrounded HQ-SAM wins thefirst place:1st_place_medal: on SegInW benchmark (consist of 25 public zero-shot in the wild segmentation datasets), and outpuerforming Grounded SAM using the same grounding-dino detector.Model NameEncoderGroundingDINOMean APEvaluation ScriptLogOutput JsonGrounded SAMvit-hswin-b48.7scriptlogresultGrounded HQ-SAMvit-hswin-b49.6scriptlogresultPlease change the current folder path to:cd seginwWe provide detailed evaluation instructions and metrics on SegInW inGrounded-HQ-SAM evaluation.Light HQ-SAM vs MobileSAM on COCOWe proposeLight HQ-SAMbased on the tiny vit image encoder provided by MobileSAM. We provide quantitative comparison on zero-shot COCO performance, speed and memory below. Try Light HQ-SAM athere.ModelEncoderAPAP@LAP@MAP@SModel Params (MB)FPSMemory (GB)MobileSAMTinyViT44.361.848.128.838.644.83.7Light HQ-SAMTinyViT45.062.848.829.240.341.23.7Note: For the COCO dataset, we use the same SOTA detector FocalNet-DINO trained on the COCO dataset as our and Mobile sam's box prompt generator.ONNX exportHQ-SAM's lightweight mask decoder can be exported to ONNX format so that it can be run in any environment that supports ONNX runtime. Export the model withpython scripts/export_onnx_model.py --checkpoint <path/to/checkpoint> --model-type <model_type> --output <path/to/output>See theexample notebookfor details on how to combine image preprocessing via HQ-SAM's backbone with mask prediction using the ONNX model. It is recommended to use the latest stable version of PyTorch for ONNX export.CitationIf you find HQ-SAM useful in your research or refer to the provided baseline results, please star :star: this repository and consider citing :pencil::@article{sam_hq,
    title={Segment Anything in High Quality},
    author={Ke, Lei and Ye, Mingqiao and Danelljan, Martin and Liu, Yifan and Tai, Yu-Wing and Tang, Chi-Keung and Yu, Fisher},
    journal = {arXiv:2306.01567},
    year = {2023}
}AcknowledgmentsThanksSAM,Grounded SAMandMobileSAMfor their public code and released models.