SAM自动生成物体mask
由于SAM可以高效处理提示,可以通过在图像上抽样大量的提示来生成整个图像的mask。这种方法被用来生成数据集SA-1B。
类SamAutomaticMaskGenerator实现了这个功能。它通过在图像上的网格中对单点输入提示进行抽样,从每个提示中SAM可以预测多个mask。然后,使用非极大值抑制对mask进行质量过滤和去重。其他选项允许进一步提高mask的质量和数量,例如在图像的多个裁剪上运行预测,或者对mask进行后处理以去除小的不连通区域和孔洞。
设置
import numpy as np import torch import matplotlib.pyplot as plt import cv2 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)
示例图像
image = cv2.imread('images/dog.jpg') image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plt.figure(figsize=(20,20)) plt.imshow(image) plt.axis('off') plt.show()
自动mask生成
要运行自动mask生成,将一个SAM模型提供给SamAutomaticMaskGenerator类。在下面设置SAM检查点的路径。建议在CUDA上运行并使用默认模型。
import sys sys.path.append("..") from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor sam_checkpoint = "sam_vit_h_4b8939.pth" model_type = "vit_h" device = "cuda" sam = sam_model_registry[model_type](checkpoint=sam_checkpoint) sam.to(device=device) mask_generator = SamAutomaticMaskGenerator(sam)
要生成mask,只需在图像上运行generate。
masks = mask_generator.generate(image)
mask生成返回一个mask列表,其中每个mask是一个包含有关mask的各种数据的字典。这些键包括:
- segmentation:mask
- area:mask的面积(以像素为单位)
- bbox:mask的边界框(XYWH格式)
- predicted_iou:模型对mask质量的预测
- point_coords:生成此mask的抽样输入点
- stability_score:mask质量的附加衡量指标
- crop_box:用于生成此mask的图像裁剪(XYWH格式)
print(len(masks)) print(masks[0].keys())
dict_keys(['segmentation', 'area', 'bbox', 'predicted_iou', 'point_coords', 'stability_score', 'crop_box'])
显示所有mask叠加在图像上。
plt.figure(figsize=(20,20)) plt.imshow(image) show_anns(masks) plt.axis('off') plt.show()
自动mask生成选项
自动mask生成中有几个可调参数,用于控制点是如何抽样的,以及移除低质量或重复mask的阈值是什么。此外,可以自动在图像的裁剪上运行生成,以在较小的对象上获得更好的性能,并且后处理可以去除杂散像素和孔洞。以下是一个示例配置,用于抽样更多的mask:
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)
len(masks2)
90
plt.figure(figsize=(20,20)) plt.imshow(image) show_anns(masks2) plt.axis('off') plt.show()