用sam2实现图像语义分割
七崽爱吃小饼干2025/05/14阅读 3
简介
SAM 2(Segment Anything Model 2)是Meta推出的用于图像和视频中对象分割的模型,在行人重识别领域有诸多应用,能辅助提取行人特征、优化训练数据、实现跨模态融合等,提升行人重识别的性能。
- SAM 2模型概述
SAM 2是一种能够在任何视频或图像中实现快速、精确对象选择的分割模型,是首个统一的跨图像和视频分割对象的模型。用户可通过点击、框选或掩码等方式选择图像或视频帧中的对象,还能在视频帧间选择和调整对象,且在未见过的对象、图像和视频上具有强大的零样本性能,能进行实时交互并输出结果,在对象分割任务上超越了同类模型。
- SAM 2在行人重识别领域的应用
- 行人区域分割与特征提取:在行人重识别中,SAM 2可精准分割出图像或视频中的行人区域,排除背景干扰,帮助模型聚焦于行人本身。通过这种方式,能更准确地提取行人的特征,如颜色、纹理、姿态等,为后续的身份匹配提供更可靠的特征依据。在复杂背景的监控视频里,SAM 2可以快速分割出行人,让行人重识别模型更有效地学习行人的独特特征,提高特征的辨识度,进而提升行人重识别的准确率。
- 辅助训练数据增强:利用SAM 2对训练数据集中的行人图像进行分割,生成高质量的分割掩码,这些掩码可作为额外的监督信息,辅助行人重识别模型的训练。将分割掩码与行人重识别模型的预测结果相结合,能使模型学习到更具判别力的特征表示。通过对比掩码中行人的类别标签和重识别模型对行人身份的预测,模型可以更好地理解行人特征与身份之间的关系,优化模型参数,增强在复杂场景下对不同行人的区分能力。
- 跨模态融合助力重识别:在涉及多模态数据的行人重识别场景中,SAM 2分割得到的行人区域可与其他模态数据(如红外图像、深度图像等)进行融合。结合可见光图像的分割结果和红外图像的热特征,能够获取更全面的行人信息,弥补单一模态数据的不足,提高重识别的准确性和鲁棒性。在夜间或光线较暗的环境下,红外图像能提供额外的信息,与SAM 2分割后的可见光图像融合,有助于更准确地识别行人。
- 处理遮挡问题:行人在监控场景中常常会出现被遮挡的情况,这给行人重识别带来很大挑战。SAM 2的记忆机制和处理遮挡的能力可以发挥作用,它能在行人被遮挡时,依靠记忆库预测行人再次出现时的位置和外观,这对于行人重识别模型在处理遮挡情况时保持对行人身份的持续跟踪非常关键,有助于提高在遮挡场景下的行人重识别性能。
用SAM2生成语义分割图像
1.安装相关库
bash
pip install git+https://github.com/facebookresearch/segment-anything.git
pip install torch torchvision # 根据CUDA版本调整
pip install matplotlib # 用于可视化
pip install opencv
pip install numpy
2. 下载 SAM2 预训练模型
从 Meta AI 官方仓库下载适合你需求的模型权重:
bash
# SAM2-H (高精度,大模型)
wget https://dl.fbaipublicfiles.com/segment-anything/sam_vit_h_4b8939.pth
# 或 SAM2-L (平衡精度和速度)
wget https://dl.fbaipublicfiles.com/segment-anything/sam_vit_l_0b3195.pth
# 或 SAM2-B (轻量级,适合移动设备)
wget https://dl.fbaipublicfiles.com/segment-anything/sam_vit_b_01ec64.pth
3. 目录结构
codeType
- sam
- data # 数据目录
- 0001
-0016.jpg
- model # 模型数据
- sam_vit_l_0b3195.pth # 模型参数
- main.py # 代码
4.分割图像
python
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
import matplotlib.pyplot as plt
import numpy as np
import cv2
import os
print(f"当前工作目录: {os.getcwd()}")
# 加载 SAM2 模型
model_type = "vit_l"
checkpoint = "./model/sam_vit_l_0b3195.pth"
sam = sam_model_registry[model_type](checkpoint=checkpoint)
# 若使用 GPU
sam.to(device="cuda")
# 创建自动掩码生成器(SAM2 优化版本)
mask_generator = SamAutomaticMaskGenerator(
model=sam,
points_per_side=32, # 控制自动分割的密度
pred_iou_thresh=0.86, # SAM2 推荐阈值
stability_score_thresh=0.92, # SAM2 新增参数,控制掩码稳定性
crop_n_layers=1,
crop_nms_thresh=0.7 # SAM2 优化的非极大值抑制阈值
)
# 加载图像
image = cv2.imread('data/0006/0001.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# 生成掩码
masks = mask_generator.generate(image)
print(f"掩码数据:{masks}")
# 可视化所有掩码
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)
# 为每个掩码随机分配颜色
for ann in sorted_anns:
m = ann['segmentation']
img = np.ones((m.shape[0], m.shape[1], 3))
color_mask = np.random.random((1, 3)).tolist()[0]
for i in range(3):
img[:, :, i] = color_mask[i]
ax.imshow(np.dstack((img, m * 0.35))) # 0.35 是透明度
# 显示结果
plt.figure(figsize=(12, 10))
plt.subplot(1, 2, 1)
plt.imshow(image)
plt.title('原始图像')
plt.axis('off')
plt.subplot(1, 2, 2)
plt.imshow(image)
show_anns(masks) # 叠加所有掩码
plt.title('自动分割结果')
plt.axis('off')
plt.tight_layout()
plt.savefig('sam2_automatic_segmentation.jpg', dpi=300)
plt.show()
5.结果展示
