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