高质量图像分割模型 SAM-HQ 实战指南
SAM-HQ（Segment Anything in High Quality）是由 ETH VIS 团队推出的零样本分割模型，在 NeurIPS 2023 上发布。该模型在原始 SAM 基础上优化了掩码生成质量，特别擅长处理边缘细节丰富的物体，支持全图自动分割和交互式点选分割两种模式。
环境配置与模型部署
获取源码并配置运行环境：
git clone https://github.com/SysCV/sam-hq.git
cd sam-hq
conda create -n samhq python=3.9
conda activate samhq
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt下载预训练权重文件（vit_h 版本精度最高）：
wget https://huggingface.co/lkeab/hq-sam/resolve/main/sam_hq_vit_h.pth -P ./pretrained/基础推理实现
以下是完整的图像分割推理代码，包含模型初始化与可视化输出：
import numpy as np
import torch
from PIL import Image
import matplotlib.pyplot as plt
from segment_anything import sam_model_registry
from segment_anything.utils.transforms import ResizeLongestSide
# 配置设备与模型路径
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
ckpt_path = "./pretrained/sam_hq_vit_h.pth"
# 加载 HQ-SAM 模型
sam = sam_model_registry["vit_h"](checkpoint=ckpt_path)
sam.to(device)
sam.eval()
# 预处理函数
def prepare_image(img_path):
raw_img = np.array(Image.open(img_path).convert("RGB"))
transform = ResizeLongestSide(sam.image_encoder.img_size)
input_img = transform.apply_image(raw_img)
input_tensor = torch.as_tensor(input_img, device=device).permute(2, 0, 1).contiguous()
return raw_img, sam.preprocess(input_tensor.unsqueeze(0))
# 执行推理
src_img, img_tensor = prepare_image("demo.jpg")
with torch.no_grad():
features = sam.image_encoder(img_tensor)
sparse_emb, dense_emb = sam.prompt_encoder(points=None, boxes=None, masks=None)
mask_pred, hq_token = sam.mask_decoder(
image_embeddings=features,
image_pe=sam.prompt_encoder.get_dense_pe(),
sparse_prompt_embeddings=sparse_emb,
dense_prompt_embeddings=dense_emb,
multimask_output=True,
hq_token_only=True
)
# 提取高质量掩码并可视化
high_quality_mask = torch.sigmoid(mask_pred[:, -1, :, :]).squeeze().cpu().numpy()
plt.imshow(src_img)
plt.imshow(high_quality_mask, alpha=0.6, cmap='jet')
plt.axis('off')
plt.savefig("output.png", bbox_inches='tight', pad_inches=0)交互式点选分割
通过指定前景/背景点实现精细化分割控制：
def segment_with_points(image_path, fg_coords, bg_coords=None):
"""
fg_coords: 前景点坐标列表 [(x1, y1), (x2, y2), ...]
bg_coords: 背景点坐标列表，可选
"""
img_np, img_prep = prepare_image(image_path)
h, w = img_np.shape[:2]
# 坐标归一化
def normalize(coords):
if coords is None:
return None
arr = np.array(coords)
return arr / np.array([w, h])
fg_norm = normalize(fg_coords)
bg_norm = normalize(bg_coords) if bg_coords else None
# 构建点标签（1=前景, 0=背景）
point_coords = np.vstack([fg_norm, bg_norm]) if bg_norm else fg_norm
point_labels = np.array([1]*len(fg_coords) + [0]*len(bg_coords)) if bg_norm else np.ones(len(fg_coords))
# 编码提示
coords_torch = torch.as_tensor(point_coords[None, ...], device=device)
labels_torch = torch.as_tensor(point_labels[None, ...], device=device)
with torch.no_grad():
img_feats = sam.image_encoder(img_prep)
sparse_emb, dense_emb = sam.prompt_encoder(
points=(coords_torch, labels_torch),
boxes=None,
masks=None
)
masks, scores, hq_feat = sam.mask_decoder(
image_embeddings=img_feats,
image_pe=sam.prompt_encoder.get_dense_pe(),
sparse_prompt_embeddings=sparse_emb,
dense_prompt_embeddings=dense_emb,
multimask_output=True,
hq_token_only=True
)
# 选择置信度最高的掩码
best_idx = torch.argmax(scores)
final_mask = masks[0, best_idx].cpu().numpy()
return final_mask
# 示例：分割花朵的特定花瓣
target_mask = segment_with_points(
"flower.jpg",
fg_coords=[(450, 320), (480, 350)], # 花瓣上的两个点
bg_coords=[(300, 200)] # 背景排除点
)批量处理与性能优化
针对视频序列或大批量图像的高效处理方案：
from tqdm import tqdm
import os
class HQSegmentor:
def __init__(self, model_type="vit_h", checkpoint=None):
self.sam = sam_model_registry[model_type](checkpoint=checkpoint)
self.sam.to(device)
self.sam.eval()
self.input_size = self.sam.image_encoder.img_size
def encode_batch(self, image_batch):
"""批量编码图像特征"""
with torch.cuda.amp.autocast(): # 混合精度加速
return self.sam.image_encoder(image_batch)
def batch_segment(self, img_dir, output_dir, grid_size=32):
"""自动网格采样分割"""
os.makedirs(output_dir, exist_ok=True)
img_files = [f for f in os.listdir(img_dir) if f.endswith(('.jpg', '.png'))]
for fname in tqdm(img_files):
img_path = os.path.join(img_dir, fname)
img_raw, img_tensor = self.prepare(img_path)
# 生成网格点提示
h, w = img_raw.shape[:2]
grid_x = np.linspace(0, w-1, grid_size)
grid_y = np.linspace(0, h-1, grid_size)
grid_points = np.array([[x, y] for y in grid_y for x in grid_x])
# 分块处理避免显存溢出
batch_points = np.array_split(grid_points, max(1, len(grid_points)//64))
all_masks = []
for pts in batch_points:
mask = self.predict_from_points(img_raw, pts)
all_masks.append(mask)
# 合并掩码结果
merged = self.merge_masks(all_masks)
self.save_overlay(img_raw, merged, os.path.join(output_dir, fname))
# 初始化并运行
segmentor = HQSegmentor(checkpoint="./pretrained/sam_hq_vit_l.pth")
segmentor.batch_segment("./input_images/", "./segmented_output/")相关生态工具
Autodistill-SAM-HQ：自动化标注流水线，支持将 SAM-HQ 作为教师模型生成伪标签，用于训练下游检测模型。
Grounded-SAM-HQ：融合 Grounding DINO 的文本检测能力与 SAM-HQ 的分割精度，实现"文本描述→目标定位→高质量分割"的全流程。
