gatilin/sam2_segmenter

SAM2 ONNX 独立推理包

模型: sam2-hiera-base-plus (Segment Anything Model 2)

文件清单

文件	大小	说明
vision_encoder.onnx	264 MB	Vision Encoder (Hiera backbone + FPN)
mask_decoder.onnx	16 MB	Mask Decoder (点/bbox prompt → mask)
conv_s0.npz	33 KB	FPN conv 权重 (256→32 ch)
conv_s1.npz	65 KB	FPN conv 权重 (256→64 ch)
prompt_encoder.npz	29 KB	Prompt Encoder 权重 (位置编码 + 类型嵌入)
preprocessor_config.json	—	预处理配置
infer.py	—	独立推理脚本

依赖

pip install numpy opencv-python onnxruntime

无需 torch / CUDA

用法

# 单点分割
python infer.py --image photo.jpg --point 500 400

# 多点分割 (每 2 个点自动分组)
python infer.py --image photo.jpg --point 500 400 --point 600 300

# bbox 中心点分割
python infer.py --image photo.jpg --bbox 100 200 500 600

# 保存结果
python infer.py --image photo.jpg --point 500 400 --save result.png --save-mask mask.png

模型规格

Vision Encoder

• 输入: pixel_values — (1, 3, 1024, 1024) float32
• 输出 (7 个):
  • last_hidden_state, fpn_0 (256ch, 256×256), fpn_1 (256ch, 128×128), fpn_2 (256ch, 64×64)
  • pos_0, pos_1, pos_2 (对应位置编码)

Mask Decoder

• 输入: image_embeddings, image_pos_enc, sparse_prompt_emb (B,1,2,256), dense_prompt_emb, high_res_feat_0/1
• 输出: iou_scores, pred_masks, object_score_logits

推理流程

1. 图像预处理: resize longest edge → 1024, pad → 1024×1024, ImageNet normalize
2. Vision Encoder: 1 次调用 → 图像特征缓存 (可复用多次分割)
3. Prompt Encoder: 点坐标 → sinusoidal 位置编码 + 类型嵌入
4. Mask Decoder: 生成 256×256 mask → resize 回原始尺寸

坐标系

• 输入坐标使用 原始图像像素坐标 (自动缩放到 1024 空间)
• 支持前景点 (label=1) 和背景点 (label=0)

限制

• ONNX mask_decoder 固定 n_points=2 (单点自动填充 not-a-point)
• 仅支持 batch_size=1
• 固定输入尺寸 1024×1024