原创作者:RS迷途小书童
我上篇博文分享了Segment Anything(SAM)模型的基本操作,这篇给大家分享下官方的整张图片的语义分割代码(全局),同时我还修改了一部分支持掩膜和叠加影像的保存。
1 Segment Anything介绍
1.1 概况
Meta AI 公司的 Segment Anything 模型是一项革命性的技术,该模型能够根据文本指令或图像识别,实现对任意物体的识别和分割。这一模型的推出,将极大地推动计算机视觉领域的发展,并使得图像分割技术进一步普及化。
论文地址:https://arxiv.org/abs/2304.02643
项目地址:Segment Anything
1.2 使用方法
具体使用方法上,Segment Anything 提供了简单易用的接口,用户只需要通过提示,即可进行物体识别和分割操作。例如在图片处理中,用户可以通过 Hover & Click 或 Box 等方式来选取物体。值得一提的是,SAM 还支持通过上传自己的图片进行物体分割操作,提取物体用时仅需数秒。
总的来说,Meta AI 的 Segment Anything 模型为我们提供了一种全新的物体识别和分割方式,其强大的泛化能力和广泛的应用前景将极大地推动计算机视觉领域的发展。未来,我们期待看到更多基于 Segment Anything 的创新应用,以及在科学图像分析、照片编辑等领域的广泛应用。
2 模型代码+注释
2.1 模型预加载
我这里将掩膜生成的函数单独拿出来了,因为里面集成了掩膜保存的代码。所以先给大家看预处理部分。
try:
image = cv2.imread(image_path) # 读取的图像以NumPy数组的形式存储在变量image中
print("[%s]正在转换图片格式......" % datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 将图像从BGR颜色空间转换为RGB颜色空间,还原图片色彩(图像处理库所认同的格式)
print("[%s]正在初始化模型参数......" % datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
except:
print("图片打开失败!请检查路径!")
pass
sys.exit()
sys.path.append("..") # 将当前路径上一级目录添加到sys.path列表,这里模型使用绝对路径所以这行没啥用
sam_checkpoint = model_path # 定义模型路径
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device) # 定义模型参数
mask_generator = SamAutomaticMaskGenerator(model=sam, # 用于掩膜预测的SAM模型
points_per_side=32, # 图像一侧的采样点数,总采样点数是一侧采样点数的平方,点数给的越多,分割越细
# points_per_batch=64, # 设置模型同时运行的点的数量。更高的数字可能会更快,但会使用更多的GPU内存
pred_iou_thresh=0.86, # 滤波阈值,在[0,1]中,使用模型的预测掩膜质量0.86
stability_score_thresh=0.92,
# 滤波阈值,在[0,1]中,使用掩码在用于二进制化模型的掩码预测的截止点变化下的稳定性0.92
# stability_score_offset=1.0, # 计算稳定性分数时,对截止点的偏移量
# box_nms_thresh=0.7, # 非最大抑制用于过滤重复掩码的箱体IoU截止点
crop_n_layers=1, # 如果>0,蒙版预测将在图像的裁剪上再次运行。设置运行的层数,其中每层有2**i_layer的图像裁剪数1
# crop_nms_thresh=0.7, # 非最大抑制用于过滤不同作物之间的重复掩码的箱体IoU截止值
# crop_overlap_ratio=512 / 1500, # 设置作物