【视觉分割新SOTA】一种最先进的图像分割模型！Segment Anything Model (SAM)如何使用？附代码和解释。

news2024/10/11 21:20:52

【视觉分割新SOTA】一种最先进的图像分割模型！Segment Anything Model (SAM)如何使用？附代码和解释。

文章目录

【视觉分割新SOTA】一种最先进的图像分割模型！Segment Anything Model (SAM)如何使用？附代码和解释。
前言
1. SAM 模型分割的通用框架
2. 应用场景 1: AI 辅助打标
3. 应用场景 2: 医疗图像分割
4. 应用场景 3: 自动驾驶
5. 应用场景 4: 卫星遥感
6. 解释各参数和模块
总结

欢迎宝子们点赞、关注、收藏！欢迎宝子们批评指正！
祝所有的硕博生都能遇到好的导师！好的审稿人！好的同门！顺利毕业！

大多数高校硕博生毕业要求需要参加学术会议，发表EI或者SCI检索的学术论文会议论文：
可访问艾思科蓝官网，浏览2024年即将召开的学术会议列表。会议入口：https://ais.cn/u/mmmiUz

参考链接：https://ai.meta.com/research/publications/segment-anything/

前言

为了在 Python 中实现 Segment Anything Model (SAM) 的应用，我们可以使用一个深度学习框架（如 PyTorch）结合 SAM 的提示分割功能。下面，我将提供一个示例代码框架，展示如何使用 SAM 模型进行四个应用场景的分割任务：AI 辅助打标、医疗图像分割、自动驾驶以及卫星遥感。

由于 SAM 模型在现实应用中涉及大量计算和资源消耗，并且需要大规模的预训练模型，我们假设 SAM 已经被加载并使用了适当的库，如 Meta 的 FAIR 实验室提供的预训练权重。

首先，我们需要以下库：

torch（PyTorch，深度学习框架）
segment-anything（假设 SAM 提供了相应的库）
opencv（用于图像处理）
matplotlib（用于可视化分割结果）
numpy（用于处理数据）

pip install torch opencv-python matplotlib numpy

1. SAM 模型分割的通用框架

我们将定义一个通用的 sam_segmentation 函数，它接受图像和提示信息，返回分割后的掩码。然后，针对四个具体的场景应用该函数。

import torch
import cv2
import numpy as np
import matplotlib.pyplot as plt
from segment_anything import SamModel

# 加载预训练的 SAM 模型（假设已经下载好）
model = SamModel()

def display_mask(image, mask):
    """辅助函数：显示分割掩码"""
    plt.figure(figsize=(10, 10))
    plt.imshow(image)
    plt.imshow(mask, alpha=0.5, cmap="jet")
    plt.axis("off")
    plt.show()

def sam_segmentation(image, prompt_type, prompt_data):
    """
    SAM 通用分割函数
    :param image: 输入图像
    :param prompt_type: 提示类型，例如 'point', 'box'
    :param prompt_data: 提示数据，例如点的位置或框的坐标
    :return: 分割掩码
    """
    # 将图像转化为 PyTorch tensor 并归一化
    input_image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
    input_image = input_image.unsqueeze(0)  # 增加 batch 维度

    # 根据提示类型生成相应的提示数据
    if prompt_type == 'point':
        prompt = {'points': torch.tensor(prompt_data)}
    elif prompt_type == 'box':
        prompt = {'boxes': torch.tensor(prompt_data)}
    else:
        raise ValueError("未知提示类型")

    # 使用 SAM 模型进行分割
    with torch.no_grad():
        mask = model.forward(input_image, prompt)
    
    return mask

# 读取图像
image = cv2.imread("sample_image.jpg")
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

2. 应用场景 1: AI 辅助打标

# 场景 1: AI 辅助打标
# 用户提供点提示，模型完成分割
points_prompt = [[150, 200]]  # 例如: 在图像中的某个点提示
mask = sam_segmentation(image_rgb, 'point', points_prompt)

# 显示分割结果
display_mask(image_rgb, mask)

参数解释：

points_prompt: 用户在图像中选择的提示点位置，模型通过这些点来完成分割。

3. 应用场景 2: 医疗图像分割

# 场景 2: 医疗图像分割
# 例如医生在肿瘤的边界上选中一个框来提示模型进行分割
box_prompt = [[100, 100, 200, 200]]  # 框住肿瘤区域的坐标 (x1, y1, x2, y2)
mask = sam_segmentation(image_rgb, 'box', box_prompt)

# 显示分割结果
display_mask(image_rgb, mask)

参数解释：

box_prompt: 医生提供的框，框住肿瘤区域的坐标，通过该框模型可以检测并分割出病变区域。

4. 应用场景 3: 自动驾驶

# 场景 3: 自动驾驶
# 在自动驾驶场景中，用户选择框住车辆的区域提示
box_prompt = [[300, 400, 500, 600]]  # 框住车辆的区域
mask = sam_segmentation(image_rgb, 'box', box_prompt)

# 显示分割结果
display_mask(image_rgb, mask)

参数解释：

box_prompt: 自动驾驶场景中车辆所在区域的坐标，用于提示模型进行分割。

5. 应用场景 4: 卫星遥感

# 场景 4: 卫星遥感
# 用户通过点击提示分割建筑物、河流等对象
points_prompt = [[350, 450], [200, 300]]  # 选中建筑物或河流的多个点位置
mask = sam_segmentation(image_rgb, 'point', points_prompt)

# 显示分割结果
display_mask(image_rgb, mask)