使用MLX-Examples中的Segment Anything模型实现自动图像分割-CSDN博客

本文链接：https://blog.csdn.net/gitblog_00336/article/details/148464025

使用MLX-Examples中的Segment Anything模型实现自动图像分割

概述

Segment Anything Model (SAM) 是近年来计算机视觉领域的一项重要突破，它能够对图像中的对象进行高质量分割。在ml-explore/mlx-examples项目中，提供了一个基于SAM的自动掩码生成器示例，展示了如何利用该模型实现无需人工干预的全自动图像分割。

技术原理

SAM的核心优势在于其能够通过提示(prompt)高效处理图像分割任务。自动掩码生成器(SamAutomaticMaskGenerator)通过在图像上网格化采样单点输入提示，让SAM为每个点预测多个掩码。随后系统会：

对生成的掩码进行质量过滤
使用非极大值抑制(NMS)去除重复掩码
可选地进行后处理以移除小的不连续区域和空洞

这种方法曾被用于生成著名的SA-1B数据集，证明了其在大规模应用中的有效性。

环境准备

首先需要导入必要的Python库：

import numpy as np
import matplotlib.pyplot as plt
import cv2

同时定义了一个可视化函数show_anns()，用于将生成的掩码叠加显示在原图上：

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)

基本使用流程

1. 加载图像

image = cv2.imread('images/dog.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

2. 初始化自动掩码生成器

from segment_anything import SamAutomaticMaskGenerator
from segment_anything.sam import load

sam_checkpoint = "../sam-vit-base"
sam = load(sam_checkpoint)
mask_generator = SamAutomaticMaskGenerator(sam)

3. 生成掩码

masks = mask_generator.generate(image)

生成的每个掩码包含以下信息：

segmentation: 实际的分割掩码
area: 掩码区域的像素面积
bbox: 边界框(XYWH格式)
predicted_iou: 模型预测的掩码质量分数
point_coords: 生成该掩码的采样点坐标
stability_score: 掩码稳定性评分
crop_box: 用于生成该掩码的图像裁剪区域(XYWH格式)

4. 可视化结果

plt.figure(figsize=(20,20))
plt.imshow(image)
show_anns(masks)
plt.axis('off')
plt.show()

高级配置选项

自动掩码生成器提供了多个可调参数，可以优化不同场景下的分割效果：

mask_generator_2 = SamAutomaticMaskGenerator(
    model=sam,
    points_per_side=32,  # 每边采样点数
    pred_iou_thresh=0.86,  # 预测IOU阈值
    stability_score_thresh=0.92,  # 稳定性分数阈值
    crop_n_layers=1,  # 图像裁剪层数
    crop_n_points_downscale_factor=2,  # 裁剪层点数降采样因子
    min_mask_region_area=100,  # 最小掩码区域面积(需要OpenCV)
)

各参数说明：