Vision - 视觉分割开源算法 SAM2(Segment Anything 2) 配置与推理 教程 (1)
欢迎关注我的CSDN：https://spike.blog.csdn.net/
本文地址：https://spike.blog.csdn.net/article/details/143220597
免责声明：本文来源于个人知识与公开资料，仅用于学术交流，欢迎讨论，不支持转载。
SAM2(Segment Anything Model 2) 视觉分割算法是 计算机视觉(CV) 的关键技术，精确的将图像中的不同对象区分开来，使用深度学习模型来分析图像中的像素分布，生成 Mask，标识出每个对象的边界。通过在多层裁剪和不同尺度下计算 Mask 的稳定性评分，确保结果的高精度和稳定性。
Paper: Segment Anything in Images and Videos
1. 环境配置
运行代码：
git clone https://github.com/facebookresearch/sam2.git
注意：项目文件比较大，可以直接使用 zip 包，或者使用 GitHub 代理。
构建环境：
conda create -n sam2 python=3.10
conda activate sam2
安装 PyTorch 包：
pip3 install torch torchvision torchaudio
python
import torch
print(torch.__version__) # 2.5.0+cu124
print(torch.cuda.is_available()) # True
exit()
环境依赖：
Python ≥ 3.10、PyTorch ≥ 2.3.1
配置 CUDA 环境变量：
export CUDA_HOME=/usr/local/cuda # change to your CUDA toolkit path
echo $CUDA_HOME
安装 SAM2 项目：
pip install --no-build-isolation -e .
pip install --no-build-isolation -e ".[notebooks]" # 适配 Jupyter
--no-build-isolation是禁用构建隔离，避免 CUDA 无法访问。
将 conda 导入 Jupyter 环境：
pip install ipykernel
python -m ipykernel install --user --name sam2 --display-name "sam2"
环境变量
export PYTORCH_ENABLE_MPS_FALLBACK=1，PyTorch 将会在遇到 MPS 不支持的操作时，自动切换到 CPU 处理。
2. 测试推理
导入 Python 包：
import os
import numpy as np
import torch
import matplotlib.pyplot as plt
from PIL import Image
# os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
os.environ["CUDA_VISIBLE_DEVICES"] = "4" # 自定义使用的卡
配置 Torch 运行设备，与运行精度，即：
# select the device for computation
if torch.cuda.is_available():
device = torch.device("cuda")
elif torch.backends.mps.is_available():
device = torch.device("mps")
else:
device = torch.device("cpu")
print(f"using device: {device}")
if device.type == "cuda":
# use bfloat16 for the entire notebook
torch.autocast("cuda", dtype=torch.bfloat16).__enter__()
# turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)
cuda_major = torch.cuda.get_device_properties(0).major
print(f"[Info] cuda_major: {cuda_major}")
if cuda_major >= 8:
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
elif device.type == "mps":
print(
"\nSupport for MPS devices is preliminary. SAM 2 is trained with CUDA and might "
"give numerically different outputs and sometimes degraded performance on MPS. "
"See e.g. https://github.com/pytorch/pytorch/issues/84936 for a discussion."
)
注意：
- 启用 bfloat16 数据类型，自动混合精度计算，有助于提高模型训练的速度和效率，同时保持较高的精度。
torch.autocast("cuda", dtype=torch.bfloat16).__enter__()
- 获取第一个 CUDA 设备（通常是 GPU）的主要版本号，判断设备是否支持 Ampere 架构的 GPU (>=8 支持)
cuda_major = torch.cuda.get_device_properties(0).major
- 如果检测 CUDA 设备的主要版本号 >= 8，即 Ampere 架构的 GPU，则启用 tfloat32 (TensorFloat-32)，在 Ampere 设备上提高矩阵运算性能的优化技术。
torch.backends.cuda.matmul.allow_tf32 = Truetorch.backends.cudnn.allow_tf32 = True
显示标注的 mask 信息 show_anns() 即：
np.random.seed(3)
def show_anns(anns, borders=True):
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.5]])
img[m] = color_mask
if borders:
import cv2
contours, _ = cv2.findContours(m.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
# Try to smooth contours
contours = [cv2.approxPolyDP(contour, epsilon=0.01, closed=True) for contour in contours]
cv2.drawContours(img, contours, -1, (0, 0, 1, 0.4), thickness=1)
ax.imshow(img)
读取图像，显示图像：
# image = Image.open('notebooks/images/cars.jpg')
image = Image.open('[your path]/llm/vision_test_data/image2.png')
image = np.array(image.convert("RGB"))
# image.shape (569, 1138, 3)
plt.figure(figsize=(20, 20))
plt.imshow(image)
plt.axis('on') # 现在坐标
plt.show()
构建自动 Mask 生成器，使用默认参数，注意选择模型 sam2.1_hiera_large.pt ，以及配置参数 sam2.1_hiera_l.yaml，即：
from sam2.build_sam import build_sam2
from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
sam2_checkpoint = "sam2.1_hiera_large.pt"
model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
sam2 = build_sam2(model_cfg, sam2_checkpoint, device=device, apply_postprocessing=False)
mask_generator = SAM2AutomaticMaskGenerator(sam2)
生成图像的 Mask：
masks = mask_generator.generate(image)
print(len(masks))
# 43
print(masks[0].keys())
# dict_keys(['segmentation', 'area', 'bbox', 'predicted_iou', 'point_coords', 'stability_score', 'crop_box'])
图像预测效果：
plt.figure(figsize=(20, 20))
plt.imshow(image)
show_anns(masks)
plt.axis('on')
plt.show()
默认效果：
自定义生成器：
mask_generator_2 = SAM2AutomaticMaskGenerator(
model=sam2,
points_per_side=64,
points_per_batch=128,
pred_iou_thresh=0.7,
stability_score_thresh=0.92,
stability_score_offset=0.7,
crop_n_layers=1,
box_nms_thresh=0.7,
crop_n_points_downscale_factor=2,
min_mask_region_area=25.0,
use_m2m=True,
)
参数说明：
model (Sam)：用于生成 mask 预测的 SAM2 模型。points_per_side (int or None)：沿图像一侧采样的点数。总点数为points_per_side**2。如果为 None，则需要point_grids提供显式点采样。points_per_batch (int)：模型同时处理的点数。点数越多，速度越快，占用更多 GPU 内存。pred_iou_thresh (float)：使用模型预测的 mask 质量的过滤阈值，范围在[0,1]。stability_score_thresh (float)：使用 mask 在二值化过程中，变化的稳定性作为过滤阈值，范围在[0,1]。stability_score_offset (float)：计算稳定性评分时用于调整 mask 的偏移量。box_nms_thresh (float)：非极大值抑制中使用的Box IoU阈值，用于过滤重复 mask。crop_n_layers (int)：如果 >0，会在图像裁剪后，再次运行 mask 预测。设置要运行的层数，每层有2**i_layer个图像裁剪。crop_n_points_downscale_factor (int)：第 n 层采样的每边点数按crop_n_points_downscale_factor**n缩放。min_mask_region_area (int)：如果 >0，后处理将移除面积小于min_mask_region_area的分离区域和 mask 中的孔洞。需要 OpenCV。use_m2m (bool)：是否使用以前的 mask 预测进行一步优化，即在 mask 中，继续进行分割 mask。
运行图像分割：
masks2 = mask_generator_2.generate(image)
plt.figure(figsize=(20, 20))
plt.imshow(image)
show_anns(masks2)
plt.axis('off')
plt.show()
示例图像分割，更加细腻，即：
