之前的文章【CV大模型SAM(Segment-Anything)】真是太强大了,分割一切的SAM大模型使用方法:可通过不同的提示得到想要的分割目标,中详细介绍了大模型SAM(Segment-Anything)根据不同的提示方式得到不同的目标分割结果。
后续有小伙伴问我,怎么样一键分割图片中的所有对象,并且保存各个对象的分割结果? 这篇文章解答一下这个问题。
如何保存不同提示的分割结果?请参考我的另一篇博文:【CV大模型SAM(Segment-Anything)】如何保存分割后的对象mask?并提取mask对应的图片区域?
关于Segment-Anything模型的相关代码、论文PDF、预训练模型、使用方法等,我都已打包好,供需要的小伙伴交流研究,获取方式如下:
关注文末名片GZH:阿旭算法与机器学习,回复:【SAM】即可获取SAM相关代码、论文、预训练模型、使用方法文档等
1. 一键分割图片中所有对象
import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2def show_anns(anns):if len(anns) == 0:returnsorted_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)))image = cv2.imread('notebooks/images/dog.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)plt.figure(figsize=(20,20))
plt.imshow(image)
plt.axis('off')
plt.show()import sys
sys.path.append("..")
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictorsam_checkpoint = "./models/sam_vit_b_01ec64.pth"
model_type = "vit_b"device = "cpu"sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)mask_generator = SamAutomaticMaskGenerator(sam)
masks = mask_generator.generate(image)plt.figure(figsize=(20,20))
plt.imshow(image)
show_anns(masks)
plt.axis('off')
plt.show()
原图:
一键分割后的结果:
2.一键分割的SamAutomaticMaskGenerato参数调整与说明
调整SamAutomaticMaskGenerator函数相关的参数,可以得到不同尺度的分割结果,代码示例如下:
mask_generator_2 = 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, # Requires open-cv to run post-processing
)masks2 = mask_generator_2.generate(image)plt.figure(figsize=(20,20))
plt.imshow(image)
show_anns(masks2)
plt.axis('off')
plt.show()
分割结果:
通过调整SamAutomaticMaskGenerator函数的相应参数后,可以发现:相比于默认参数,这个示例分割的结果更加细致了,例如:盆子里面的不同区域等都被独立分割出来了。
SamAutomaticMaskGenerator()参数详细说明:
model (Sam):用于掩膜预测的SAM模型。
points_per_side (int or None): 沿着图像一侧采样的点的数量。点的总数是point_per_side**2。如果没有,'point_grids’必须提供明确的点采样。
points_per_batch (int):设置模型同时运行的点的数量。更高的数字可能会更快,但会使用更多的GPU内存。
pred_iou_thresh (float): 滤波阈值,在[0,1]中,使用模型的预测掩膜质量。
stability_score_thresh (float):滤波阈值,在[0,1]中,使用掩码在用于二进制化模型的掩码预测的截止点变化下的稳定性。
stability_score_offset (float):计算稳定性分数时,对截止点的偏移量。 - box_nms_thresh (float):非最大抑制用于过滤重复掩码的箱体IoU截止点。
crop_n_layers (int):如果>0,蒙版预测将在图像的裁剪上再次运行。设置运行的层数,其中每层有2**i_layer的图像裁剪数。
crop_nms_thresh (float):非最大抑制用于过滤不同作物之间的重复掩码的箱体IoU截止值。
crop_overlap_ratio(float):设置作物重叠的程度。在第一个作物层中,作物将以图像长度的这个分数重叠。在第一个裁剪层中,裁剪物将以图像长度的这一比例重叠,以后的裁剪层中,更多的裁剪物将缩小这一重叠。
crop_n_points_downscale_factor (int):在图层n中每面采样的点数被crop_n_points_downscale_factor**n缩减。
point_grids (list(np.ndarray) or None):用于取样的明确网格的列表,归一化为[0,1]。列表中的第n个网格被用于第n个作物层。与points_per_side排他。
min_mask_region_area (int):如果>0,后处理将被应用于移除面积小于min_mask_region_area的遮罩中的不连接区域和孔。需要opencv。
output_mode (str):掩模的返回形式。可以是’binary_mask’, ‘uncompressed_rle’, 或者’coco_rle’。coco_rle’需要pycocotools。对于大的分辨率,'binary_mask’可能会消耗大量的内存。
SamAutomaticMaskGenerator()函数的参数默认值:
model: Sam,
points_per_side: Optional[int] = 32,
points_per_batch: int = 64,
pred_iou_thresh: float = 0.88,
stability_score_thresh: float = 0.95,
stability_score_offset: float = 1.0,
box_nms_thresh: float = 0.7,
crop_n_layers: int = 0,
crop_nms_thresh: float = 0.7,
crop_overlap_ratio: float = 512 / 1500,
crop_n_points_downscale_factor: int = 1,
point_grids: Optional[List[np.ndarray]] = None,
min_mask_region_area: int = 0,
output_mode: str = “binary_mask”,
3. 保存所有分割后的目标对象
for i, mask in enumerate(masks):mask = ~mask['segmentation']mask = mask + 255mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2)mask = mask.astype(np.uint8)res = cv2.bitwise_and(image, mask)res[res == 0] = 255plt.imshow(res)plt.savefig('res-{}.png'.format(i + 1))plt.show()
下图只展示了4个分割结果:
关于Segment-Anything模型的相关代码、论文PDF、预训练模型、使用方法等,我都已打包好,供需要的小伙伴交流研究,获取方式如下:
关注文末名片GZH:阿旭算法与机器学习,回复:【SAM】即可获取SAM相关代码、论文、预训练模型、使用方法文档等
如果文章对你有帮助,感谢点赞+关注!