HQ-SAM (Segment Anything in High Quality)

If you want to skip building the SAM 2 CUDA extension (which provides post-processing to remove small holes and sprinkles in output masks), set the SAM2_BUILD_CUDA environment variable to 0. This will not affect core segmentation results in most cases.

SAM2_BUILD_CUDA=0 pip install -e ".[notebooks]"

If you are working with the seginw folder for Segmentation in the Wild benchmarks, you must install GroundingDINO in editable mode.

cd seginw
python -m pip install -e GroundingDINO

To use the HQ-SAM 2 predictor and run the provided example notebooks, install the [notebooks] extra via pip.

pip install -e ".[notebooks]"

Install Grounding DINO in editable mode using pip. If you have a CUDA environment, ensure the CUDA_HOME environment variable is set to enable GPU acceleration. If no CUDA is available, the package will be compiled in CPU-only mode.

pip install -e .

Run the training process using torch.distributed.launch. You must specify the checkpoint path, the model type, and the output directory.

Arguments:

• --checkpoint: Path to the initialization checkpoint.
• --model-type: The ViT backbone type (e.g., vit_b, vit_l, vit_h).
• --output: Path to the directory where training results will be saved.
• --nproc_per_node: Number of GPUs to use.

python -m torch.distributed.launch --nproc_per_node=<num_gpus> train.py --checkpoint <path/to/checkpoint> --model-type <model_type> --output <path/to/output>

To train HQ-SAM, you must prepare the HQSeg-44K dataset. The dataset can be downloaded from Hugging Face. The directory structure must follow this pattern:

data
|____DIS5K
|____cascade_psp
| |____DUTS-TE
| |____DUTS-TR
| |____ecssd
| |____fss_all
| |____MSRA_10K
|____thin_object_detection
| |____COIFT
| |____HRSOD
| |____ThinObject5K

Download all model checkpoints using the provided shell script in the checkpoints directory, or download the HQ-SAM 2 large checkpoint individually from Hugging Face.

cd checkpoints && \
./download_ckpts.sh && \
cd ..

Download and unzip the SegInW dataset into the data directory.

cd data
wget https://huggingface.co/sam-hq-team/SegInW/resolve/main/seginw.zip
unzip seginw.zip

The expected structure is data/seginw/<category_name>/ containing various subfolders like Airplane-Parts, Bottles, etc.

To ensure the CUDA extension is built and to catch any errors during the process, use the SAM2_BUILD_ALLOW_ERRORS=0 environment variable. This is useful if you want to enable the mask post-processing step and want the installation to fail explicitly if the build fails.

pip uninstall -y SAM-2 && \
rm -f ./sam2/*.so && \
SAM2_BUILD_ALLOW_ERRORS=0 pip install -v -e ".[notebooks]"

Download the required checkpoints for GroundingDINO, SAM, and HQ-SAM into the pretrained_checkpoint directory.

cd pretrained_checkpoint
wget https://huggingface.co/ShilongLiu/GroundingDINO/resolve/main/groundingdino_swinb_cogcoor.pth
wget https://huggingface.co/sam-hq-team/sam-hq-training/resolve/main/pretrained_checkpoint/sam_vit_h_4b8939.pth
wget https://huggingface.co/lkeab/hq-sam/resolve/main/sam_hq_vit_h.pth

For a quick setup, install the segment-anything-hq package using pip. You can then use the sam_model_registry to load models directly in Python.

pip install segment-anything-hq

from segment_anything_hq import sam_model_registry
model_type = "<model_type>" # e.g., "vit_l", "vit_b", "vit_h", or "vit_tiny"
sam_checkpoint = "<path/to/checkpoint>"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)