SAM2分割模型微调指南

SAM2(Segment Anything 2)是 Meta 推出的一款新模型,旨在对图像中的任何内容进行分割,而不局限于特定的类别或领域。该模型的独特之处在于其训练数据规模:1100 万张图像和 110 亿个掩码。这种广泛的训练使 SAM2 成为训练新图像分割任务的强大起点。

你可能会问,如果 SAM 可以分割任何东西,为什么我们还需要重新训练它?答案是 SAM 在常见物体方面非常出色,但在罕见或特定领域的任务上表现不佳。

但是,即使在 SAM 给出的结果不足的情况下,仍然可以通过在新数据上对其进行微调来显着提高模型的能力。在许多情况下,这将比从头开始训练模型占用更少的训练数据并获得更好的结果。

本教程演示了如何在仅 60 行代码(不包括标注和导入)中对新数据进行微调 SAM2。

完整的训练脚本可以在这里找到。

NSDT工具推荐: Three.js AI纹理开发包 - YOLO合成数据生成器 - GLTF/GLB在线编辑 - 3D模型格式在线转换 - 可编程3D场景编辑器 - REVIT导出3D模型插件 - 3D模型语义搜索引擎 - Three.js虚拟轴心开发包 - 3D模型在线减面 - STL模型在线切割 

1、SAM的工作原理

SAM 的主要工作方式是获取图像和图像中的一个点,并预测包含该点的片段的掩码。这种方法无需人工干预即可实现完整的图像分割,并且对片段的类别或类型没有任何限制(如上一篇文章所述)。

使用 SAM 进行完整图像分割的过程:

  • 选择图像中的一组点
  • 使用 SAM 预测包含每个点的片段
  • 将生成的片段组合成单个地图

虽然 SAM 还可以利用其他输入,例如掩码或边界框,但这些输入主要与涉及人工输入的交互式分割相关。在本教程中,我们将重点介绍全自动分割,并且仅考虑单点输入。

有关该模型的更多详细信息,请访问SAM项目网站

2、下载 SAM2 并设置环境

SAM2 可从这里下载。如果你不想复制训练代码,也可以下载我分叉的版本,其中已经包含TRAIN.py脚本。

然后按照 github 存储库上的安装说明进行操作。通常,你需要 Python >=3.11 和 PyTorch。

此外,我们将使用 OpenCV,可以使用以下方式安装:

pip install opencv-python

你还需要从这里下载预训练模型,你可以从多种模型中进行选择,所有模型都与本教程兼容。我建议使用训练速度最快的小型模型。

下一步是下载用于微调模型的数据集。在本教程中,我们将使用 LabPics1 数据集来分割材料和液体。你可以从这里下载数据集。

4、准备数据读取器

我们需要编写的第一件事是数据读取器。它将读取并准备网络数据。

数据读取器需要生成:

  • 图像
  • 图像中所有片段的蒙版。
  • 以及每个蒙版内的随机点

让我们从加载依赖项开始:

import numpy as np
import torch
import cv2
import os
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor

接下来我们列出数据集中的所有图像:

data_dir=r"LabPicsV1//" # Path to LabPics1 dataset folder
data=[] # list of files in dataset
for ff, name in 
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值