【计算机视觉】Segment Anything 模型部署与实战：从零到一的自动掩码生成

1. Segment Anything模型简介

Segment Anything Model（SAM）是Meta AI推出的通用图像分割模型，它能够根据点、框等简单提示生成高质量的物体掩码。这个模型在包含1100万张图像和10亿个掩码的数据集上进行了训练，具备强大的零样本迁移能力。

我第一次接触SAM时就被它的泛化能力震惊了。随便扔给它一张从未见过的图片，不需要任何微调，它就能准确地分割出图像中的各种物体。这种"开箱即用"的特性，让SAM成为计算机视觉领域的"瑞士军刀"。

SAM的核心优势在于：

• 提示式分割：通过点击或框选等简单交互就能获得精确分割结果
• 全自动分割：无需人工提示即可自动生成图像中所有物体的掩码
• 强大的零样本能力：在未经专门训练的新领域数据上也能表现良好

2. 环境准备与安装

2.1 基础环境配置

在开始之前，我们需要准备好基础运行环境。SAM要求Python≥3.8，以及PyTorch≥1.7和torchvision≥0.8。我强烈建议使用支持CUDA的PyTorch版本，这将大幅提升推理速度。

以下是详细的安装步骤：

# 创建并激活虚拟环境（推荐） conda create -n sam_env python=3.8 -y conda activate sam_env # 安装PyTorch（根据你的CUDA版本选择对应命令） # CUDA 11.3 pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113

如果你不确定该安装哪个版本的PyTorch，可以访问 PyTorch官网 获取适合你系统的安装命令。

2.2 安装Segment Anything

安装好基础环境后，我们可以通过以下命令安装SAM：

pip install git+https://github.com/facebookresearch/segment-anything.git

如果网络连接不稳定，也可以选择克隆仓库后本地安装：

git clone https://github.com/facebookresearch/segment-anything.git cd segment-anything pip install -e .

安装完成后，建议安装一些可选依赖，这些将在后续的示例中使用：

pip install opencv-python pycocotools matplotlib onnxruntime onnx jupyter

3. 模型下载与加载

3.1 下载预训练模型

SAM提供了三种不同规模的预训练模型：

• 默认/ViT-H：性能最好但体积最大
• ViT-L：平衡了性能和速度
• ViT-B：最轻量但精度稍低

对于大多数应用场景，我推荐使用默认的ViT-H模型。你可以从官方提供的链接下载模型权重：

wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth

下载完成后，建议将模型文件放在项目目录下的models文件夹中，方便管理。

3.2 加载模型

加载SAM模型非常简单，以下是完整的代码示例：

import torch from segment_anything import sam_model_registry # 设置模型路径和类型 sam_checkpoint = "models/sam_vit_h_4b8939.pth" model_type = "vit_h" device = "cuda" if torch.cuda.is_available() else "cpu" # 加载模型 sam = sam_model_registry[model_type](checkpoint=sam_checkpoint) sam.to(device=device)

在实际使用中，我发现将模型放到GPU上可以将推理速度提升3-5倍。如果你的显存不足（小于8GB），可以考虑使用ViT-L或ViT-B模型。

4. 自动掩码生成实战

4.1 准备测试图像

让我们先从简单的示例开始。我们将使用SAM自带的示例图像进行测试：

import cv2 import matplotlib.pyplot as plt # 读取并显示图像 image = cv2.imread('images/dog.jpg') image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) plt.figure(figsize=(10,10)) plt.imshow(image) plt.axis('off') plt.show()

4.2 创建掩码生成器

SAM提供了两种主要的使用方式：

1. 提示式分割：通过点或框等提示进行交互式分割
2. 自动掩码生成：自动识别图像中的所有物体并生成掩码

我们先来看自动掩码生成：

from segment_anything import SamAutomaticMaskGenerator # 创建自动掩码生成器 mask_generator = SamAutomaticMaskGenerator( model=sam, points_per_side=32, # 每边采样的点数 pred_iou_thresh=0.86, # 掩码质量的阈值 stability_score_thresh=0.92, # 稳定性得分阈值 crop_n_layers=1, # 裁剪层数 crop_n_points_downscale_factor=2, # 裁剪点下采样因子 min_mask_region_area=100, # 最小掩码区域面积 ) # 生成掩码 masks = mask_generator.generate(image) print(f"检测到 {len(masks)} 个掩码")

4.3 可视化结果

为了直观地查看分割结果，我们需要一个可视化函数：

def show_anns(anns): if len(anns) == 0: return sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True) ax = plt.gca() ax.set_autoscale_on(False) img = np.ones((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1], 4)) img[:,:,3] = 0 for ann in sorted_anns: m = ann['segmentation'] color_mask = np.concatenate([np.random.random(3), [0.35]]) img[m] = color_mask ax.imshow(img) # 显示结果 plt.figure(figsize=(20,20)) plt.imshow(image) show_anns(masks) plt.axis('off') plt.show()

这段代码会为每个检测到的物体分配一个随机的半透明颜色，方便我们观察分割效果。在实际项目中，你可能需要调整SamAutomaticMaskGenerator的参数来获得最佳效果。

5. 提示式分割应用

5.1 基于点的交互式分割

除了自动分割，SAM还支持通过点、框等提示进行交互式分割。这在需要精确控制分割结果的场景中特别有用：

from segment_anything import SamPredictor # 创建预测器 predictor = SamPredictor(sam) predictor.set_image(image) # 定义提示点 (x,y) 和标签 (1表示前景点，0表示背景点) input_point = np.array([[500, 375]]) # 狗的鼻子位置 input_label = np.array([1]) # 前景点 # 预测掩码 masks, scores, logits = predictor.predict( point_coords=input_point, point_labels=input_label, multimask_output=True, # 输出多个可能的结果 ) # 显示最佳结果 best_mask = masks[np.argmax(scores)] plt.figure(figsize=(10,10)) plt.imshow(image) show_mask(best_mask, plt.gca()) show_points(input_point, input_label, plt.gca()) plt.axis('off') plt.show()

5.2 基于框的分割

当物体边界比较清晰时，使用框提示往往能获得更好的效果：

# 定义边界框 [x1,y1,x2,y2] input_box = np.array([425, 300, 700, 500]) # 预测掩码 masks, _, _ = predictor.predict( point_coords=None, point_labels=None, box=input_box[None, :], multimask_output=False, ) # 显示结果 plt.figure(figsize=(10,10)) plt.imshow(image) show_mask(masks[0], plt.gca()) show_box(input_box, plt.gca()) plt.axis('off') plt.show()

在实际应用中，你可以结合点和框提示来获得更精确的分割结果。SAM的灵活性让它能够适应各种复杂场景。

6. 高级应用与优化技巧

6.1 批量处理图像

当需要处理大量图像时，直接使用上述方法效率可能不高。我总结了几个提升性能的技巧：

def batch_process_images(image_paths, output_dir): # 预加载模型 sam = sam_model_registry[model_type](checkpoint=sam_checkpoint) sam.to(device=device) mask_generator = SamAutomaticMaskGenerator(sam) os.makedirs(output_dir, exist_ok=True) for path in image_paths: # 读取图像 image = cv2.imread(path) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 生成掩码 masks = mask_generator.generate(image) # 保存结果 base_name = os.path.basename(path) save_path = os.path.join(output_dir, f"mask_{base_name}") plt.imsave(save_path, visualize_masks(image, masks))

6.2 参数调优指南

SAM的性能很大程度上取决于参数设置。以下是我通过大量实验总结的最佳实践：

1. points_per_side：控制采样密度。值越大找到的小物体越多，但计算量也越大。通常32-64是较好的平衡点。

2. pred_iou_thresh：过滤低质量掩码。建议设置在0.8-0.9之间。

3. stability_score_thresh：控制掩码稳定性。0.9左右通常效果不错。

4. crop_n_layers：启用多尺度处理有助于检测小物体，但会显著增加计算时间。

5. min_mask_region_area：去除小面积噪声。根据应用场景调整，通常100-300比较合适。

# 优化后的参数配置 optimized_generator = SamAutomaticMaskGenerator( model=sam, points_per_side=64, pred_iou_thresh=0.88, stability_score_thresh=0.92, crop_n_layers=1, crop_n_points_downscale_factor=1, min_mask_region_area=200, )

7. 实际应用案例

7.1 医学图像分析

在医疗领域，SAM可以辅助医生快速定位和分割病灶区域。我在一个肺部CT项目中使用SAM取得了不错的效果：

# 加载医学图像 ct_image = load_dicom("lung_scan.dcm") # 调整参数适应医学图像 medical_mask_generator = SamAutomaticMaskGenerator( model=sam, points_per_side=128, # 医学图像需要更高密度采样 pred_iou_thresh=0.92, # 医学应用需要更高精度 stability_score_thresh=0.95, ) # 生成分割结果 medical_masks = medical_mask_generator.generate(ct_image)

7.2 遥感图像处理

对于卫星和航拍图像，SAM可以自动识别建筑物、道路等地物：

# 处理大尺寸遥感图像 def process_large_image(image_path, tile_size=1024): large_image = Image.open(image_path) width, height = large_image.size for i in range(0, width, tile_size): for j in range(0, height, tile_size): tile = large_image.crop((i, j, i+tile_size, j+tile_size)) tile_np = np.array(tile) # 使用较小的points_per_side节省计算资源 masks = mask_generator.generate(tile_np) # 处理并保存结果 process_and_save_masks(masks, i, j)

7.3 工业质检

在生产线上的缺陷检测场景中，SAM可以快速定位产品表面的异常区域：

# 工业质检应用 def detect_defects(product_image): # 首先定位产品区域 product_mask = locate_product(product_image) # 在ROI内进行精细分割 roi = apply_mask(product_image, product_mask) defect_masks = mask_generator.generate(roi) # 过滤小面积区域 real_defects = [m for m in defect_masks if m['area'] > MIN_DEFECT_AREA] return real_defects