1.拉取sam2的源码()

git clone https://github.com/facebookresearch/sam2.git

拉取后,该代码需要PyTorch、TorchVisionpython>=3.10torch>=2.5.1TorchVision 。请按照此处的torchvision>=0.20.1,要安装对应库。

安装后运行

python -c "import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.device_count())"

python -c "import torch, torchvision; print('torch version:', torch.__version__); print('torchvision version:', torchvision.__version__)"

正确会打印版本信息,和是否支持cuda。错误重新安装。conda环境下建议conda直接安装

2按步骤安装依赖:

cd  sam2

pip install -e .

3.然后下载训练权重

cd checkpoints && \
./download_ckpts.sh && \
cd ..

4.测试运行sam2 

import torch
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor

checkpoint = "./checkpoints/sam2.1_hiera_large.pt"
model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"

# 创建 predictor
predictor = SAM2ImagePredictor(build_sam2(model_cfg, checkpoint))

# 加载图片
image_path = "~/sam2/notebooks/images/cars.jpg" #自己的项目路径
image = np.array(Image.open(image_path).convert("RGB"))
predictor.set_image(image)

# 整图预测,不提供提示点
input_prompts = None

# 执行预测
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
    masks, scores, logits = predictor.predict(input_prompts)

# 显示第一个 mask
mask = masks[0]  # numpy array
plt.imshow(image)
plt.imshow(mask, alpha=0.5)  # 半透明叠加
plt.show()

5.上面的操作后,下面才是重点。第一下载项目自己的训练集

https://github.com/henghuiding/MOSE-api?tab=readme-ov-file#downloadhttps://github.com/henghuiding/MOSE-api?tab=readme-ov-file#download下载一个小版本的,我下载的是valid.tar.gz。然后解压到文件夹中。

我直接放 ~/sam2trainDir/valid/valid.解压后有JPEGImages和Annotations两个文件夹。里面是原始图片和掩码图片

6.复制或者直接使用项目中的已有配置文件~/sam2/sam2/configs/sam2.1_training/sam2.1_hiera_b+_MOSE_finetune.yaml

这个文件是写好的视频分帧数据训练。我们只要修改其中的配置参数就可以运行。

7 修改参数

 修改路径为你解压的文件。

img_folder: /home/bing/sam2trainDir/valid/valid/JPEGImages
  gt_folder: /home/bing/sam2trainDir/valid/valid/Annotations
  file_list_txt: /home/bing/sam2trainDir/validtest.txt

其中validtest.txt写入你要训练的视频名:比如JPEGImages/002845.可以一次写多个,按列写好就行。然后记得修改scratch:
  resolution: 640   #这个是分辨率,太高gpu吃不消
  train_batch_size: 1
  num_train_workers: 10
  num_frames: 1 #这个是每个视频读取的帧数,主要文件中你处理了几帧数据写几,valid只有一帧,记得改成1
  max_num_objects: 3  #每帧图像中最多训练或预测的对象数量。
  base_lr: 5.0e-6
  vision_lr: 3.0e-06
  phases_per_epoch: 1 #每个 epoch 的阶段数。
  num_epochs: 40  #总训练轮数(epoch)。 

详细如下

resolution: 640

输入图像的分辨率(短边或长边?通常是短边 resize 至 640 像素)。

影响训练时图像缩放、内存占用以及模型对小目标的感知能力。

较小分辨率:训练快,占用内存少,但细节损失;

较大分辨率:保留更多细节,但显存占用增加。

train_batch_size: 1

每次训练迭代送入模型的样本数。

你这里是 1,说明 GPU 显存可能有限或者你使用单帧训练。

batch size 太小会导致梯度不稳定,但适合显存受限的情况。

num_train_workers: 10

数据加载器(DataLoader)的工作线程数。

用于并行读取数据,提高训练数据加载速度。

太多线程可能导致 CPU 占用过高或内存竞争。

num_frames: 1

每个样本包含的帧数。

如果是视频分割任务,num_frames 可以大于 1 表示一次输入多帧进行训练。

你这里是 1,说明每次训练只用单帧图像。

max_num_objects: 3

每帧图像中最多训练或预测的对象数量。

超过这个数量的对象可能被忽略或者限制在 3 个实例。

对小目标场景或资源受限训练有用。

base_lr: 5.0e-6

主学习率(base learning rate),控制整个模型参数更新的步长。

很小的学习率说明训练很谨慎,通常用于从头训练或者微调大型模型。

vision_lr: 3.0e-06

图像编码器(Vision Encoder)的学习率。

SAM2 中视觉编码器可能需要比解码器更小的学习率,以防破坏预训练权重。

当训练从零开始(scratch)时,可以设置与 base_lr 相近;微调时通常更小。

phases_per_epoch: 1

每个 epoch 的阶段数。

SAM2 训练可能分成多个阶段(phase),例如 warm-up、fine-tune 等。

这里设为 1,表示每个 epoch 只有一个训练阶段。

num_epochs: 40

总训练轮数(epoch)。

每个 epoch 遍历一次完整训练集。

40 个 epoch 对于小数据集来说可能足够,但大数据集通常需要更多轮。

然后运行训练语句 在sam2文件夹下,注意是第一层sam2文件夹,这个文件夹下还有一个sam2文件夹

python -m training.train -c sam2.1_my_hiera_b+_MOSE_finetune.yaml     --use-cluster 0     --num-gpus 2

这样就开始训练了。

如果是自己的数据集,记得下载labelme ,进行数据标注。得到的json文件使用python,调用的也是labelme的接口转成同名的图片放在Annotations中做掩码图片。

#!/usr/bin/env python3
import os
import json
import glob
import numpy as np
from PIL import Image
import labelme

# 数据路径
input_dir = "/mnt/hgfs/fileShare/sam2Data/result"  # JSON 文件所在目录
output_dir = os.path.join(input_dir, "masks")      # 保存黑白遮罩的目录
os.makedirs(output_dir, exist_ok=True)

# 遍历所有 JSON 文件
json_files = glob.glob(os.path.join(input_dir, "*.json"))

for json_file in json_files:
    with open(json_file, "r") as f:
        data = json.load(f)
    
    # 原图大小
    img_height = data["imageHeight"]
    img_width = data["imageWidth"]
    img_shape = (img_height, img_width)
    
    # 构造 label_name_to_value,所有标注统一为 1
    label_name_to_value = {shape["label"]: 1 for shape in data["shapes"]}
    
    # 生成 mask
    mask, _ = labelme.utils.shapes_to_label(img_shape, data["shapes"], label_name_to_value)
    
    # 转成黑白 mask
    mask = (mask > 0).astype(np.uint8) * 255
    
    # 保存
    mask_img = Image.fromarray(mask)
    base_name = os.path.splitext(os.path.basename(json_file))[0]
    mask_path = os.path.join(output_dir, f"{base_name}.png")
    mask_img.save(mask_path)
    
    print(f"生成黑白遮罩: {mask_path}")

print("全部完成!")

保证文件格式是这样的

~/sam2trainDir/newDateSet/
├── JPEGImages/
│   ├── 112233/
│   │   ├── 00000.png
│   │   ├── 00001.png
│   │   └── ...
└── Annotations/
    ├── 112233/
    │   ├── 00000.png
    │   ├── 00001.png
    │   └── ...

这样训练接口调用不报错。

训练结束后会生成新的权重

注意:

1.sam2.1_my_hiera_b+_MOSE_finetune.yaml文件要复制到sam2/sam2文件夹下。不然找不到

2.训练集图片要使用或改成jpg格式,不然训练集找不到视频图片

Logo

DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐