使用SAM（Segment Anything Model）模型中的SamAutomaticMaskGenerator一键分割图片中所有对象并保存所有分割后的目标对象

君不见~

已于 2024-09-07 12:12:31 修改

阅读量162

点赞数 5

分类专栏： YOLO + SAM 文章标签：计算机视觉深度学习人工智能经验分享

于 2024-07-18 10:19:16 首次发布

本文链接：https://blog.csdn.net/qq_23865133/article/details/140514902

版权

YOLO + SAM 专栏收录该内容

2 篇文章 0 订阅

订阅专栏

下面将介绍使用SAM（Segment Anything Model）模型中的SamAutomaticMaskGenerator一键分割图片中所有对象并保存所有分割后的目标对象，包括保存图像完整的mask和图像中各个mask。

将所有mask保存在一张TIF图像中：

import os
import cv2
import numpy as np
import torch
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
import tifffile as tif

# 设置设备
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
MODEL_TYPE = "vit_h"
CHECKPOINT_PATH = "sam_h.pth"

# 加载SAM模型
sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)
mask_generator = SamAutomaticMaskGenerator(sam)

# 读取图像
image_path = "yolo_sam_test/dimo_0126.png"
image_name = os.path.splitext(os.path.basename(image_path))[0]
image_bgr = cv2.imread(image_path)
image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)

# 为图像生成掩码
sam_result = mask_generator.generate(image_rgb)

# 创建文件夹
output_dir = "sam_all_mask"
os.makedirs(output_dir, exist_ok=True)

# 合并所有掩码为一个二值图像
all_masks = np.zeros_like(image_rgb[:, :, 0], dtype=np.uint8)
for mask in sam_result:
    all_masks += (mask['segmentation'] * 255).astype(np.uint8)

# 保存为TIFF
output_path = os.path.join(output_dir, f"{image_name}_mask.tif")
tif.imwrite(output_path, all_masks)

print(f"All masks saved to {output_path}")

运行结果：

也可以通过调整SamAutomaticMaskGenerator函数中的相关参数，以得到不同尺度的分割结果，下面是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",

参数说明：

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’可能会消耗大量的内存

自己修改

import os
import cv2
import numpy as np
import torch
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
import tifffile as tif

# 设置设备
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
MODEL_TYPE = "vit_h"
CHECKPOINT_PATH = "sam_h.pth"

# 加载SAM模型
sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)
mask_generator = SamAutomaticMaskGenerator(sam)

# 读取图像
image_path = "yolo_sam_test/dimo_0126.png"
image_name = os.path.splitext(os.path.basename(image_path))[0]
image_bgr = cv2.imread(image_path)
image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)

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
)

sam_result = mask_generator_2.generate(image_rgb)

# 创建文件夹
output_dir = "sam_all_mask"
os.makedirs(output_dir, exist_ok=True)

# 合并所有掩码为一个二值图像
all_masks = np.zeros_like(image_rgb[:, :, 0], dtype=np.uint8)
for mask in sam_result:
    all_masks += (mask['segmentation'] * 255).astype(np.uint8)

# 保存为TIFF
output_path = os.path.join(output_dir, f"{image_name}_mask.tif")
tif.imwrite(output_path, all_masks)

print(f"All masks saved to {output_path}")

将所有mask单独保存在一张TIF图像中：

import os
import cv2
import numpy as np
import torch
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
import tifffile as tif

# 设置设备
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
MODEL_TYPE = "vit_h"
CHECKPOINT_PATH = "sam_h.pth"

# 加载SAM模型
sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)
mask_generator = SamAutomaticMaskGenerator(sam)

# 读取图像
image_path = "yolo_sam_test/dimo_0126.png"
image_name = os.path.splitext(os.path.basename(image_path))[0]

image_bgr = cv2.imread(image_path)
image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)

# 为图像生成掩码
sam_result = mask_generator.generate(image_rgb)

# 创建文件夹
output_dir = "sam_all_mask"
os.makedirs(output_dir, exist_ok=True)

# 保存每个掩码为单独的TIFF文件
for idx, mask in enumerate(sam_result):
    # 将掩码转换为二值图像
    binary_mask = (mask['segmentation'] * 255).astype(np.uint8)

    # 构建输出文件名
    output_filename = f"{image_name}_mask_{idx}.tif"
    output_path = os.path.join(output_dir, output_filename)

    # 保存TIFF
    tif.imwrite(output_path, binary_mask)

print(f"Masks saved to {output_dir}")