基于SAM2项目的自定义图像数据集微调指南

前言

Facebook Research开源的SAM2项目为图像分割任务提供了强大的基础模型。本文将详细介绍如何利用SAM2框架对自定义图像数据集进行微调，帮助研究人员和开发者快速上手。

数据集准备

数据集格式要求

SAM2支持对SA-1B格式的数据集进行直接训练。自定义数据集需要组织为以下结构：

数据集根目录/
├── train/
│   ├── images/
│   │   ├── 1.jpg
│   │   ├── 2.jpg
│   │   └── ...
│   └── labels/
│       ├── 1.json
│       ├── 2.json
│       └── ...
├── val/
└── test/

每个JSON标注文件应包含以下内容：

{
    "image": {
        "image_id": 1,
        "width": 1400,
        "height": 1400,
        "file_name": "1.jpg"
    },
    "annotations": [
        {
            "area": 45499,
            "segmentation": {
                "size": [1400, 1400],
                "counts": "RLE编码的二进制字符串"
            }
        }
    ]
}

数据格式转换

对于使用LabelMe等工具标注的数据集，需要转换为SA-1B格式。转换过程主要包括：

1. 读取原始标注的多边形坐标
2. 创建二值掩码图像
3. 使用RLE（Run-Length Encoding）对掩码进行编码
4. 生成符合SA-1B格式的JSON文件

配置文件修改

基础配置

在SAM2的配置文件中，需要针对自定义数据集进行以下调整：

data:
  train:
    _target_: training.dataset.sam2_datasets.TorchTrainMixedDataset
    phases_per_epoch: ${phases_per_epoch}
    batch_sizes:
    - ${bs1}
    datasets:
    - _target_: training.dataset.vos_dataset.VOSDataset
      training: true
      video_dataset:
        _target_: training.dataset.vos_raw_dataset.SA1BRawDataset
        img_folder: ${path_to_img_folder}
        gt_folder: ${path_to_gt_folder}
      sampler:
        _target_: training.dataset.vos_sampler.RandomUniformSampler
        num_frames: 1
        max_num_objects: ${max_num_objects_per_image}
      transforms: ${vos.train_transforms}

关键参数说明

1. num_frames: 设置为1表示使用单张图像
2. max_num_objects: 控制每张图像处理的最大目标数，建议值为3
3. transforms: 使用默认的图像增强策略即可

模型选择与训练

SAM2提供了不同规模的模型：

1. Tiny模型：计算量小，适合边缘设备
2. Base模型：平衡性能与计算量
3. Large模型：最高精度，计算需求大

训练命令示例

python training/train.py \
  -c configs/sam2.1_training/sam2.1_hiera_b+_MOSE_finetune.yaml \
  --use-cluster 0 \
  --num-gpus 1

训练资源评估

根据实际经验，不同硬件配置下的训练表现：

1. NVIDIA RTX 4090 (24GB): 可处理4帧/批次
2. NVIDIA A6000 (48GB): 同样配置下内存使用情况相似
3. 训练时间：取决于数据集规模和模型大小

常见问题解决

1. 内存不足：减少批次大小或降低max_num_objects值
2. 格式不匹配：确保标注文件严格遵循SA-1B格式
3. 性能不佳：尝试调整学习率或增加训练轮次

结语

通过本文介绍的方法，开发者可以有效地利用SAM2框架对自定义图像数据集进行微调。实际应用中，建议从小规模数据集开始试验，逐步扩大训练规模，以获得最佳的性能与效率平衡。