SAM+RGB-D双模态图像分割实战：给模型添加深度通道

简介图像分割是计算机视觉的基础任务但仅靠RGB颜色信息在弱光、逆光或目标与背景同色时难以鲁棒。深度图作为物理证据不随光照和颜色变化能有效填补单模态的盲区。将深度图与RGB融合的双模态分割在无人机航拍、自动驾驶、遥感等领域展现出显著优势。本文以SAMSegment Anything Model为基础介绍如何通过为编码器增加深度通道实现RGB-D双模态语义分割的完整工程方案。从数据预处理、模型改造、训练调参到部署优化全方位解析这一方法的落地细节帮助开发者在复杂场景中获得更稳定、更精细的分割结果。 前阵子有朋友拿了个无人机航拍项目来找我说想把画面里的广告牌和围挡区域抠出来。他们先试了官方开源的 Segment AnythingSAM模型结果逆光、阴影、广告牌颜色接近背板时分割效果惨不忍睹。聊到最后我发现一个关键信息他们手里其实还有深度传感器和立体匹配生成的深度图。我当时的判断很直接——这个问题不该用单模态硬撞把深度当第二路输入喂给 SAM 才是正经解法。于是就有了这个基于 SAM 架构的双模态图像分割项目完整代码、预处理好的数据、已经训练好的权重都打包好了亲测可跑。这个方案的核心就一句话给 SAM 加一个输入通道让“颜色证据”和“物理证据”一起参与分割。对 RGB 单模态表现不佳的场景弱光、逆光、目标与背景同色系、近色遮挡深度图可以稳稳托底。项目适合做遥感、无人机巡检、室内语义分割、自动驾驶、工业质检和医学影像相关工作的朋友尤其是手里已经有 RGB-D 或热红外数据的人。1. 双模态 SAM 想解决什么问题1.1 单模态 SAM 的边界在哪里SAM 的训练数据以自然图像为主它的零样本能力来自海量图像语料和 prompt 机制但本质上还是在做“看起来像什么”的推理——颜面、纹理、边缘。这个思路在普通场景很强大可一旦落进专业领域就容易翻车。举个实际的例子无人机俯拍建筑外立面广告牌是红色背景配白色大字但墙面也是砖红色。在 RGB 图像里广告牌和墙面几乎是一团连续的颜色区域SAM 就算给它一个精确的框也倾向于把整面墙一起分割出来因为“颜色上它们确实是一体的”。这种场景在广告牌分割、违章搭建识别、遥感建筑提取里非常常见。更麻烦的是光照。逆光拍摄时阴影会把目标区域压暗或者目标本身处于大面积阴影里RGB 通道里根本看不到清晰轮廓。SAM 只能依赖依稀可辨的纹理和边缘结果往往是一半目标被吞进背景另一半被切得支离破碎。1.2 第二模态的本质是“物理证据”深度图不具备颜色信息但它记录了场景中每个点到相机的距离。光照变化影响不了它颜色近似影响不了它阴影更是跟它没关系。在深度剖面里广告牌是建筑立面上一个凸起的平面选出来就是一个小台阶在 RGB 里它可能跟墙面融为一体。我习惯把这种关系理解成“猜颜色”和“量距离”的差别。RGB 是颜色证据告诉模型“这里看起来像”深度/热红外/多光谱是物理证据告诉模型“这里实际上是什么”。两者交叉验证分割的鲁棒性会明显上一个台阶。双模态的好处还不止于此。热红外图像对温度敏感适合夜间巡检和人员检测多光谱图像能反映材质差异适合植被分类和水体提取。在不同场景下第二模态可以灵活替代深度图。项目代码里通道设计是通用的核心思路不变只是换一个数据源。1.3 哪些场景适合用这套方案结合我个人经验以下情况优先考虑双模态 SAM手头已经有 RGB-D 相机、激光雷达投影、双目立体匹配或热红外设备数据。目标任务中目标与背景颜色接近比如广告牌、灰白建筑物、伪装目标。场景光照不稳定阴影、逆光、夜视环境频繁出现。希望利用 SAM 的预训练能力但不想从零训练一套完整模型。反过来如果任务里目标颜色非常突出、背景干净稳定单用 RGB 模型已经到 95% 以上准确率那确实没必要上双模态。硬上反而增加数据采集、标定和对齐的复杂度成本大于收益。2. 项目整体设计与环境搭建2.1 项目组成与代码结构这套项目我按“数据、模型、训练、推理”四块整理成了清晰目录。你拿到手后不需要东翻西找目录结构如下sam_dm/ ├── checkpoints/ │ ├── sam_dm_best.pth # 已训练好的双模态SAM权重 │ └── sam_vit_b_original.pth # 官方SAM ViT-B原始权重 ├── configs/ │ └── sam_dm.yaml # 训练与推理配置 ├── data/ │ ├── dataset.py # RGB-D数据集加载 │ ├── transform.py # 同步数据增强 │ └── preprocess_depth.py # 深度补全与归一化 ├── models/ │ ├── sam.py # SAM核心模块ViT图像编码器mask解码器 │ ├── fusion.py # 双模态融合模块 │ └── heads.py # 语义分割头 ├── train.py ├── infer.py ├── export_onnx.py └── requirements.txtcheckpoints 目录里放的就是“已训练好”的权重。infer.py 里默认加载sam_dm_best.pth你不需要重新训练直接按后面的推理命令就能看到双模态分割效果。2.2 环境依赖与一条命令跑通运行环境是我验证过的直接对齐不会有兼容性问题Python 3.9PyTorch 2.0CUDA 11.8 或更高torchvision、opencv-python、numpy、tqdm、pyyaml、pillowrequirements.txt 内容大致如下torch2.0.0 torchvision0.15.0 opencv-python4.6.0.66 numpy1.23.0 tqdm4.64.0 pyyaml6.0 pillow9.0.0如果机器显存足够8GB 以上直接按下面命令跑推理# 推理模式加载训练好的权重对单张RGB-D图像做分割 python infer.py --config configs/sam_dm.yaml \ --checkpoint checkpoints/sam_dm_best.pth \ --rgb demo/rgb.png --depth demo/depth.png输出的 mask 会保存到outputs/目录同时会把原图、深度可视化、预测 mask 放到一张图上方便你直观检查效果。2.3 模型与数据集的基本情况模型底座用的是 SAM ViT-B图像编码器参数量约 91M。相比 ViT-L/H 它更轻双模态训练时显存压力小推理速度也更容易接受。我改造后的输入从 3 通道变成 4 通道输出端换成了语义分割头不保留 SAM 原始的 prompt 交互逻辑。数据方面训练集使用了一套 RGB-D 语义分割数据覆盖建筑物、车辆、行人、广告牌、围挡、植被等类别。做航拍场景时我参考了公开航空影像分割数据集的类别映射方式但去掉了过细的类只保留业务需要的类别。广告牌单独作为一类处理而且这个类在画面中占比很小训练时如果不做针对性处理很容易被模型忽略后面训练章节会讲怎么处理。3. 数据准备双模态项目真正的硬骨头3.1 数据集来源与标注策略双模态分割项目里代码反而是最省心的部分数据才是真正决定成败的地方。用公开数据集的话室内可以直接用 NYU Depth V2、SUN RGB-D航拍场景需要自己整理或者从航空影像数据集里筛。我建议先拿公开数据集把整个训练链路跑通再切换到自己业务数据。标注策略上有一个非常实用的提效方法用 SAM 自动预标注再人工修正。先用原始 SAM 框/点 prompt 把明显目标粗略切出来再把粗糙 mask 导入标注工具微调。SAM 对外观清晰的目标零样本效果很好人工只需要修正边缘和漏检几百张图的标注量可以压缩到几个小时。3.2 深度图补全与归一化很多人第一步就栽在深度图上。消费级 RGB-D 相机在透明玻璃、高反光表面、黑暗区域会产生空洞深度值直接是无效值。如果拿这种数据直接训练模型会把“无数据”区域当成一种特殊语义预测边界会出现大量奇形怪状的伪影。处理方案分两步。第一步补洞我用的方法是邻近有效像素填充加中值滤波。空洞较小时OpenCV 的inpaint配合cv2.INPAINT_TELEA就能补得比较自然空洞较大时建议换采集角度重新拍硬补出来的几何信息没有意义。第二步是归一化这一步如果出错训练 loss 会抖到怀疑人生。深度值的量级和传感器有关有的是毫米有的是米有的范围是 0 到 4095。我的做法是先把深度裁剪到有效范围比如 0.5 米到 10 米超出部分截断再除以范围最大值统一映射到 [0,1] 区间import numpy as np def normalize_depth(depth, min_val0.5, max_val10.0): depth np.clip(depth, min_val, max_val) depth (depth - min_val) / (max_val - min_val) return depth.astype(np.float32)如果没有归一化深度数值可能是 1000 量级而 RGB 通道是 0 到 1第一层卷积学习时梯度会完全被深度通道主导整个训练直接失衡。3.3 RGB 与深度图的空间对齐双模态项目里最隐蔽的坑是“两张图没对齐”。RGB 和深度来自不同传感器时分辨率、视差、时间戳都可能不一致。如果直接喂给模型它会额外学习一个隐式的位移修正不仅浪费容量而且模型对拍摄角度极其敏感换一个机位效果立刻崩。怎么验证对齐是否合格有一个土办法把深度图做伪彩色可视化和 RGB 图各自半透明叠加显示肉眼看广告牌、建筑边缘、行人轮廓是否完全重合。如果两层有明显重影说明没有对齐训练前必须先做标定或者直接用已经出厂对齐的 RGB-D 相机数据。另外要特别注意时间同步。动态场景里RGB 和深度如果相差几十毫秒运动物体会出现“重影边缘”。拍摄时尽量固定相机或者选择三脚架、慢速移动的无人机场景先保证数据干净再谈模型精度。4. 模型改造给 SAM 加上第二模态4.1 常见改造方案对比给 SAM 增加第二模态核心问题是“深度信息从哪里进入网络”。我整理过三种常见做法融合方案改动量优势风险输入级拼接4通道输入小复用预训练权重改动最少最容易复现通道间信息混合需要较多数据双编码器独立融合中深度特征更独立融合设计灵活显存翻倍需要额外设计融合层深度作为 Prompt 注入中保留 SAM 原始语义能力对 prompt 机制耦合太深调试复杂我最终选择的是输入级拼接。原因很实际改动量最小代码可维护性高在中小规模数据上效果也足够好。你如果想在项目里快速验证双模态思路这是性价比最高的方案。4.2 4 通道输入改造的具体操作SAM 图像编码器最早的一层是 PatchEmbed 卷积把输入从 3 通道映射到 embedding 维度。ViT-B 里输出通道数是 768卷积核 16x16stride 16。改造思路是把第一层卷积从3 - 768改成4 - 768并让第 4 个通道的权重不是随机初始化而是复用原有权重。核心代码大概是这样的import torch import torch.nn as nn def convert_patch_embed_to_4ch(patch_embed): old_weight patch_embed.proj.weight.data # [768, 3, 16, 16] new_conv nn.Conv2d(4, old_weight.shape[0], kernel_sizeold_weight.shape[2], strideold_weight.shape[3]) with torch.no_grad(): # 前三通道保留原RGB权重 new_conv.weight[:, :3] old_weight # 第4通道复用三通道的均值比随机初始化稳定 new_conv.weight[:, 3] old_weight.mean(dim1) patch_embed.proj new_conv return patch_embed为什么推荐复用均值权重而不是随机初始化因为 SAM 预训练权重的特征空间已经被充分学习随机初始化第 4 通道等于在起点就引入一个偏离预训练分布的大扰动。前几个 epoch 的 loss 会跳得很厉害甚至不收敛。复用均值权重可以保证网络最初的行为接近原始 SAM然后再逐步学习深度特征收敛速度会快很多。4.3 解码器与 Loss 的配套改动原始 SAM 的 mask decoder 是 prompt-based它会根据用户输入的点框生成对应 mask。做双模态语义分割时没有“用户 prompt”这个概念所以我把解码器改造为一个固定输出的语义分割头。具体做法是保留 SAM 解码器里 patch embedding 和 attention 结构作为特征精修模块去掉了 prompt 相关分支最后接一个 1x1 卷积输出类别数对应的分割 mask。class SemanticHead(nn.Module): def __init__(self, in_chans256, num_classes8): super().__init__() self.head nn.Conv2d(in_chans, num_classes, kernel_size1) def forward(self, feats): return self.head(feats)Loss 方面我采用了 Dice Loss BCE Loss 的组合def bce_dice_loss(pred, target, alpha0.5): bce nn.functional.binary_cross_entropy_with_logits(pred, target) pred_prob torch.sigmoid(pred) smooth 1.0 intersection (pred_prob * target).sum(dim(2, 3)) dice 1 - (2 * intersection smooth) / (pred_prob.sum(dim(2, 3)) target.sum(dim(2, 3)) smooth) return alpha * bce dice.mean()为什么一定要配 Dice Loss因为分割任务里前景目标占比小尤其广告牌这类目标在整张图中可能只占百分之几。只用 BCE 时模型只需要把所有像素预测为背景loss 就能压得相当低但这样没有任何实际意义。Dice Loss 对前景覆盖区域更敏感能有效对抗类别不平衡问题。5. 训练与调参从 loss 不收敛到 mIoU 稳步提升5.1 训练配置与策略完整跑通一个双模态分割训练我把关键配置固定成下面这组参数配置项数值输入分辨率512x512Batch Size8训练/ 16评测优化器AdamW初始学习率1e-4Weight Decay1e-4损失函数0.5 * Dice 0.5 * BCE混合精度fp16学习率调度CosineAnnealing 5 epoch warmup总轮次50冻结策略前6层冻结后6层微调PatchEmbed和解码器全量微调冻结前 6 层的理由是SAM ViT 的浅层主要负责低级特征边缘、纹理这和深度特征有一定通用性深层语义特征更需要针对双模态输入做适配。冻结浅层可以省显存也降低过拟合风险。数据量足够大时比如 2000 对以上可以尝试解冻更多层效果上限会更高。训练命令python train.py --config configs/sam_dm.yaml --data /path/to/rgbd_data5.2 训练过程中要盯哪些信号训练日志示例[Epoch 01/50] train_loss0.5321, train_dice0.612, val_mIoU0.521 [Epoch 10/50] train_loss0.2314, train_dice0.801, val_mIoU0.671 [Epoch 20/50] train_loss0.1834, train_dice0.857, val_mIoU0.712 [Epoch 35/50] train_loss0.1221, train_dice0.912, val_mIoU0.738 [Epoch 48/50] train_loss0.0912, train_dice0.941, val_mIoU0.735这个日志是我实验里一个非常典型的情况。前 20 个 epoch 双模态对 mIoU 的提升很明显之后训练 loss 还在下降但验证 mIoU 平台期了说明开始过拟合。这时候不用纠结跑满 50 轮保存验证集上最好的权重就行。如果 loss 一直不降优先排查数据而不是换模型。检查深度归一化是不是有问题、RGB-D 是否对齐、数据增强有没有破坏配对关系这三个点占了 80% 的不收敛原因。5.3 两个调参实战结论第一个结论Backbone 不能整段全冻。我试过把整个图像编码器冻结只训练后面的分割头效果比混合微调差不少。原因很直观SAM 的编码器是在 RGB 单模态上训练的它的特征不会自动适配深度通道的引入。只有微调深层让深层特征感知到“第 4 通道的存在”双模态信息才真正进入网络。第二个结论数据量少时别一上来就全量微调。正确顺序是先冻结浅层、微调深层和解码器跑到验证集提升不明显之后再解冻一部分浅层继续微调。这样做既能保证收敛稳定也比一开始就全量微调省显存。6. 实测效果与性能瓶颈6.1 定量效果对比在我的测试集上把广告牌分割作为一个重点类别做了对比方法广告牌 IoU整体 mIoU原始 SAMRGB onlybox prompt0.550.61UNetRGB-D 4通道输入0.580.63本项目双模态 SAM0.670.70注意原始 SAM 需要 prompt所以我在测试时给每个目标提供了精确的外接框条件已经很理想。即便如此它在颜色混杂场景下还是不如双模态 SAM。双模态的提升主要来自广告牌、围挡这类颜色信息可靠性差的类别整体 mIoU 也被带动起来。6.2 典型失败 case 与原因分析第一类失败是深度空洞大面积出现。玻璃幕墙、镜面材质会产生大量无效深度值即使补洞也不能完全恢复真实的凸起结构。这种场景下双模态会退化成“半个 RGB 模式”有时甚至比纯 RGB 还差因为模型试图从噪声深度中提取不存在的几何信息。解决办法是优先换数据源或者对深度补洞做更强约束。第二类失败是小目标漏检。输入分辨率 512x512 时画面远处的小广告牌只有几十个像素无论 RGB 还是深度都无法提供足够的判别信息。切 patch 会有效果也可以提高输入分辨率到 640 或 768但训练和推理速度都会明显下降。第三类失败是边缘过紧或过松。深度图边缘有时比 RGB 更锐利模型如果过度依赖深度预测出的 mask 边界会非常贴合深度跳变在深度传感器噪声大的地方反而出现锯齿边缘。把 Dice 权重略微调低或者训练时对深度图添加少量随机噪声做对抗能缓解这个问题。6.3 推理速度与部署优化训练好的权重在 512x512 输入下单张图片在 RTX 3090 上推理约 45ms 左右换成 FP16 可以压到 30ms。如果还要更快有三条路导出 ONNX 后用 TensorRT 做 INT8 量化速度能再提一倍换用 MobileSAM 这类轻量编码器但需要重新训练去掉 SAM 解码器里冗余的 transformer 层直接从 ViT 特征上采样分割精度略有下降速度提升明显。另外提一个很实用的降级策略实际部署时如果深度传感器突然不可用可以把第 4 通道输入置为全 0 或全 1。模型在这种情况下会退化成接近单模态 SAM 的效果但不会崩溃系统不至于完全瘫痪。更优的做法是接一个单目深度估计模型比如 MiDaS在线生成伪深度再做双模态推理效果比置零好很多。7. 踩坑记录双模态 SAM 最容易翻车的五个点坑 1深度归一化不一致导致 loss 不收敛。现象是第一轮 loss 就冲到 7 以上然后一直下不来。排查半天发现不同来源的数据深度范围不一样有些是毫米有些是米直接混合训练。解决方法是统一做 clamped normalization并且训练前把数据分布可视化检查一遍。坑 2新通道权重随机初始化导致 PatchEmbed 崩掉。我最早直接把第 4 通道初始化为随机数前几个 epoch 的 loss 波动非常大相当于在预训练模型上强行插入了一个完全陌生的输入源。改成复用三通道均值权重后训练马上稳定下来。这类问题不细看根本想不到但严重影响收敛。坑 3加载预训练权重时 shape 不匹配报错。把 PatchEmbed 从 3 通道改成 4 通道后直接用load_state_dict(strictTrue)加载官方权重会报 shape mismatch。正确做法是先加载原权重再修改层或者把 strict 设 False 然后手动分配第 4 通道权重。代码里我封装了一个转换函数避免每次加载都手动处理。坑 4RGB 和 Depth 用了不同的数据增强。第一版代码对 RGB 做随机翻转却忘了对 Depth 做同样操作。训练出来的模型在预测时出现一种很怪的现象mask 边缘在左右翻转前后不一致验证 IoU 也上不去。修法很简单所有随机操作必须在一个随机种子下同时作用于配对数据保证空间变换严格同步。def sync_transform(rgb, depth, seed): torch.manual_seed(seed) rgb random_flip(rgb) torch.manual_seed(seed) depth random_flip(depth) return rgb, depth坑 5冻结整个 Backbone 导致双模态没有真正参与语义建模。我试过冻结所有 SAM 编码器只训练分割头最终效果只比单模态 SAM 好一点点。原因前面说过编码器完全不更新深度通道的信息只经过了一个不感知它的网络融合效果非常有限。正确做法是至少让后半段编码器参与微调让网络在深层次特征上建立 RGB 与深度通道的关联。这套方案跑下来我最深的一点体会是双模态分割的成功不是靠把深度图拼进去就完事而是数据和训练细节每一层都要配合得当。深度归一化、通道权重初始化、同步增强、分层微调这些看起来不起眼的地方任何一个出问题都可能让整个项目止步不前。但只要把数据管线打通、模型改造做对双模态 SAM 在 RGB 单模态搞不定的场景里确实能带来肉眼可见的提升。如果你手头恰好有 RGB-D 数据强烈建议直接拿我的配置先跑一遍再根据自己的任务类别做调整。本文还有配套的精品资源点击获取