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