SAM-2.1 自定义分割微调实战:从自定义数据集到自动掩码生成
2026/7/12 0:16:24 网站建设 项目流程

SAM-2.1 自定义分割微调实战:从自定义数据集到自动掩码生成


这篇教程是我根据 SAM-2.1 自定义数据集微调和自动 mask 生成流程整理出来的。重点演示如何下载 SAM2 格式数据集,配置 SAM-2.1 训练脚本,训练后对比微调模型和基础模型的分割效果。

SAM-2.1 的训练流程依赖官方仓库、配置文件和 checkpoint。这个 notebook 适合作为完整项目复现模板,尤其适合希望把 SAM-2.1 迁移到特定零件、工业目标或业务数据集的场景。

本文会重点跑通以下流程:

  • 下载 SAM2 格式数据集
  • 克隆并安装 SAM2 官方仓库
  • 修正数据集文件名以适配训练配置
  • 启动 SAM-2.1 单 GPU 微调
  • 对比微调模型和基础模型的自动 mask 结果

如果你正在系统学习多模态微调、目标检测、OCR 或图像分割,建议收藏本文;配套 notebook、示例图片和运行环境说明后续会继续整理。如果环境配置卡住,可以在评论区说明具体报错。

📚 文章目录

  • SAM-2.1 自定义分割微调实战:从自定义数据集到自动掩码生成
    • 📥 下载 SAM-2.1 与数据集
    • 🧹 修改数据集文件名
    • 🏋️ 开始训练 SAM-2.1
    • 🔍 可视化微调模型结果
    • 📌 小结
    • 📚 同系列教程汇总

📥 下载 SAM-2.1 与数据集

先从数据集后台获取 car parts 数据集,再克隆 SAM2 官方仓库和训练配置。注意:原 notebook 中直接写了 API key,发布前建议替换为自己的密钥或 Colab Secrets。

fromtypesimportSimpleNamespace# 从数据集后台下载并解压数据集后,修改 DATASET_DIR 指向数据集目录。DATASET_DIR="/content/dataset"# 修改为数据集后台导出的数据集目录dataset=SimpleNamespace(location=DATASET_DIR)
!git clone https://github.com/facebookresearch/sam2.git
!wget-O/content/sam2/sam2/configs/train.yaml'https://drive.usercontent.google.com/download?id=11cmbxPPsYqFyWq87tmLgBAQ6OZgEhPG3'
%cd./sam2/
!pip install-e.[dev]-q
!cd./checkpoints&&./download_ckpts.sh

🧹 修改数据集文件名

SAM-2.1 训练对文件名格式有要求,这一步把 自定义数据集导出的文件名改成兼容形式。

# 将 自定义数据集文件名改成 SAM 2.1 兼容格式。# 也可以通过调整 sam2/sam2/configs/train.yaml 来避免这一步。importosimportre FOLDER="/content/data/train"forfilenameinos.listdir(FOLDER):# Replace all except last dot with underscorenew_filename=filename.replace(".","_",filename.count(".")-1)ifnotre.search(r"_\d+\.\w+$",new_filename):# Add an int to the end of base namenew_filename=new_filename.replace(".","_1.")os.rename(os.path.join(FOLDER,filename),os.path.join(FOLDER,new_filename))

🏋️ 开始训练 SAM-2.1

使用官方训练脚本和配置文件启动单 GPU 训练,并可通过 TensorBoard 查看训练曲线。

!python training/train.py-c'configs/train.yaml'--use-cluster0--num-gpus1
%load_ext tensorboard%tensorboard--bind_all--logdir./sam2_logs/

🔍 可视化微调模型结果

加载微调 checkpoint 和基础 checkpoint,在验证集图片上生成 mask 并并排对比。

!pip install supervision-q
importtorchfromsam2.build_samimportbuild_sam2fromsam2.automatic_mask_generatorimportSAM2AutomaticMaskGeneratorimportsupervisionassvimportosimportrandomfromPILimportImageimportnumpyasnp# 在整个 notebook 中使用 bfloat16# from Meta notebooktorch.autocast("cuda",dtype=torch.bfloat16).__enter__()iftorch.cuda.get_device_properties(0).major>=8:torch.backends.cuda.matmul.allow_tf32=Truetorch.backends.cudnn.allow_tf32=Truecheckpoint="/content/sam2/sam2_logs/configs/train.yaml/checkpoints/checkpoint.pt"model_cfg="configs/sam2.1/sam2.1_hiera_b+.yaml"sam2=build_sam2(model_cfg,checkpoint,device="cuda")mask_generator=SAM2AutomaticMaskGenerator(sam2)checkpoint_base="/content/sam2/checkpoints/sam2.1_hiera_base_plus.pt"model_cfg_base="configs/sam2.1/sam2.1_hiera_b+.yaml"sam2_base=build_sam2(model_cfg_base,checkpoint_base,device="cuda")mask_generator_base=SAM2AutomaticMaskGenerator(sam2_base)
validation_set=os.listdir("/content/data/valid")# choose random with .json extensionimage=random.choice([imgforimginvalidation_setifimg.endswith(".jpg")])image=os.path.join("/content/data/valid",image)opened_image=np.array(Image.open(image).convert("RGB"))result=mask_generator.generate(opened_image)detections=sv.Detections.from_sam(sam_result=result)mask_annotator=sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)annotated_image=opened_image.copy()annotated_image=mask_annotator.annotate(annotated_image,detections=detections)base_annotator=sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)base_result=mask_generator_base.generate(opened_image)base_detections=sv.Detections.from_sam(sam_result=base_result)base_annotated_image=opened_image.copy()base_annotated_image=base_annotator.annotate(base_annotated_image,detections=base_detections)sv.plot_images_grid(images=[annotated_image,base_annotated_image],titles=["Fine-Tuned SAM-2.1","Base SAM-2.1"],grid_size=(1,2))


📌 小结

这篇教程完整整理了Fine-Tune SAM-2.1的核心复现流程。实际操作时,建议先确认 GPU、依赖版本、数据集路径和模型权限,再逐段运行 notebook。

  • 下载 SAM2 格式数据集
  • 克隆并安装 SAM2 官方仓库
  • 修正数据集文件名以适配训练配置
  • 启动 SAM-2.1 单 GPU 微调
  • 对比微调模型和基础模型的自动 mask 结果

后续我会继续按源项目顺序整理 项目教程 中的目标检测、实例分割、OCR、多目标跟踪和视觉大模型教程。

📚 同系列教程汇总

  • Google Gemini 3.5 Flash 零样本目标检测教程:从提示词到可视化结果

  • GLM-OCR 文档识别实战教程:从验证码、公式到车牌 OCR

  • RF-DETR + ByteTrack 多目标跟踪实战教程:从命令行到 Python 视频轨迹可视化

  • SAM 3 图像分割实战教程:文本、框和点提示的多种分割方式

  • SAM 3 视频分割实战教程:用文本提示分割并跟踪视频中的目标

  • SAM-2.1 自定义分割微调实战:从自定义数据集到自动掩码生成-本文

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询