HongyouZhou/ruac-sam2.1-hiera-bplus

RUAC for SAM 2.1 Hiera Base Plus

Official checkpoint for Segment Anything with Robust Uncertainty-Accuracy Correlation (ICML 2026). RUAC adds a Weibull Bayesian mask decoder and an adversarial uncertainty-estimation training branch to SAM 2.1. It produces prompted segmentation masks together with per-pixel uncertainty maps.

Model details

• Base model: SAM 2.1 Hiera Base Plus
• Training data: MOSE training split
• Evaluation: zero-shot across 23 out-of-domain datasets
• Default uncertainty estimator: Bernoulli entropy from 20 Weibull Monte Carlo samples
• Weight format: safetensors
• Original training checkpoint SHA-256: 4f43fd1bd6ed421dbdb6a4d57b68dabddb0ccb69f8c95de089b859922e299561
• Released model.safetensors: 330,844,812 bytes, SHA-256 678dae09c6834504d20e825793d8270f83dbc1c3c2d542ce843247d7154b46e4

The value 20 is the released default and main setting. Other Monte Carlo sample counts belong to the parameter-sensitivity study and are not release defaults.

Installation

git clone --branch version/019 https://github.com/HongyouZhou/sam2.git
git clone https://github.com/HongyouZhou/BNDL.git
git clone https://github.com/HongyouZhou/ruac.git

python -m pip install -e sam2
python -m pip install -e './ruac[hub]'
export PYTHONPATH=$PWD:$PWD/sam2:$PWD/BNDL:$PWD/BNDL/BNDL_upload:$PYTHONPATH

Load the model

from ruac.hub import load_ruac_predictor

predictor = load_ruac_predictor(
    repo_id="HongyouZhou/ruac-sam2.1-hiera-bplus",
    device="cuda",
    mc_samples=20,
)

RUAC is not registered with Transformers AutoModel; the loader above builds the custom Bayesian SAM2 decoder and validates the released state dict.

Point-prompted image inference

import numpy as np
from PIL import Image

image = np.array(Image.open("image.jpg").convert("RGB"), copy=True)
predictor.set_image(image)
masks, predicted_iou, _ = predictor.predict(
    point_coords=np.array([[640, 360]], dtype=np.float32),
    point_labels=np.array([1], dtype=np.int32),
    multimask_output=True,
    return_logits=True,
)

selected = int(np.argmax(predicted_iou))
mask = masks[selected] > 0
aux = predictor.get_last_aux_outputs()
uncertainty_candidates = aux["bndl"]["pixel_uncertainty_sampling"][0]
assert uncertainty_candidates.shape[-1] == masks.shape[0]
uncertainty = uncertainty_candidates[:, :, selected]

The loader aligns uncertainty channels with the masks returned by SAM2. The four internal mask-token channels are retained under aux["bndl"]["all_mask_tokens"] for reproducing the original evaluation's all-hypothesis aggregation.

See the complete model loading guide and runnable inference example.

Limitations

RUAC follows SAM2's prompt-based segmentation interface. It does not perform semantic classification, and uncertainty quality can vary with the domain, prompt choice, and Monte Carlo sample count. The released model was trained on MOSE and should be evaluated before deployment in safety-critical settings.

Licenses

The released weights are derived from the Apache-2.0-licensed SAM 2.1 checkpoint and are distributed under Apache 2.0. The RUAC implementation is MIT licensed. Users must also follow the licenses of the SAM2 and BNDL dependencies.

Citation

@inproceedings{ruac2026,
  title         = {Segment Anything with Robust Uncertainty-Accuracy Correlation},
  author        = {Zhou, Hongyou and Toussaint, Marc and Shao, Ling and Ye, Zihan},
  booktitle     = {Proceedings of the 43rd International Conference on Machine Learning},
  year          = {2026},
  eprint        = {2605.10603},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CV}
}