【视觉分割新SOTA】一种最先进的图像分割模型!Segment Anything Model (SAM)如何使用?附代码和解释。
【视觉分割新SOTA】一种最先进的图像分割模型!Segment Anything Model (SAM)如何使用?附代码和解释。
文章目录
欢迎宝子们点赞、关注、收藏!欢迎宝子们批评指正!
祝所有的硕博生都能遇到好的导师!好的审稿人!好的同门!顺利毕业!
大多数高校硕博生毕业要求需要参加学术会议,发表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)
参数解释:
points_prompt
: 遥感图像中用户提供的地表对象提示点,如建筑物、河流等,通过这些点进行分割
6. 解释各参数和模块
sam_segmentation
函数:
image
: 输入图像,要求为 RGB 格式。prompt_type
: 提示的类型,如点提示或框提示。prompt_data
: 提示信息的具体数据,如点的坐标或框的边界。
4)返回值: 分割掩码,用于显示分割后的区域。
display_mask
函数:
- 用于显示原始图像和叠加的分割掩码,使用
matplotlib
可视化。
prompt_type
和prompt_data
:
prompt_type
决定了提示的类型,目前支持点提示和框提示。prompt_data
是具体提示的数值,例如点的位置[x, y]
或框的边界[x1, y1, x2, y2]
。
总结
上述代码框架展示了如何使用 SAM 模型在不同场景下进行图像分割。通过输入不同的提示,SAM 可以快速完成复杂图像中的精确分割。
欢迎宝子们点赞、关注、收藏!欢迎宝子们批评指正!
祝所有的硕博生都能遇到好的导师!好的审稿人!好的同门!顺利毕业!
大多数高校硕博生毕业要求需要参加学术会议,发表EI或者SCI检索的学术论文会议论文:
可访问艾思科蓝官网,浏览2024年即将召开的学术会议列表。会议入口:https://ais.cn/u/mmmiUz