简介这是一份基于神经网络实现RGB-D图像分割的完整工程代码包面向计算机视觉方向的研究者、算法工程师以及需要借助深度信息提升分割准确率的一线开发与科研人员。方案以深度感知卷积神经网络Depth-Aware CNN为核心将彩色图像与深度图联合输入弥补传统二维卷积对几何结构感知不足的缺陷从而在机器人导航、增强现实、3D重建等复杂场景中获得更精细的像素级分割结果。压缩包共70个文件主要由41个Python源码、6个C/C头文件与源码、2个CUDA加速文件、4个Shell运行脚本以及说明文档等组成整体大小仅86KB轻量且便于快速部署。目前已有1413人学习关注。工程覆盖从数据加载、模型搭建、损失定义到训练评估、可视化输出的完整流程内置NyUv2、SUNRGBD等主流数据集适配并提供Deeplab、VGG_Deeplab等多种网络结构每个模块都留有扩展接口可灵活调整基础网络、深度融合策略与超参数适合二次开发、算法对比与科研教学。1. RGB-D分割为什么值得单独做多一个深度通道多出一套麻烦手里有一台能出深度图的相机之后很多人的第一直觉是把深度图当第四个通道跟RGB一起塞进卷积神经网络跑一个现成的图像分割算法。结果通常不太好看训练收敛得慢物体边缘出现一圈一圈的锯齿深度传感器在暗处吸光时分割结果整片丢目标。RGB-D图像分割和纯RGB分割的关键差别不在于多了一个输入而在于深度图是另一种统计分布的几何信号——它稀疏、带孔洞、边界突变对光照不敏感却对材质敏感。这篇内容会把深度图预处理、双流网络结构、训练部署的坑讲清楚。适合手里有RGB-D传感器、正准备自己训练分割模型的工程师参考也适合做图像分割算法选型时想评估「深度分支到底值不值得加」的团队。2. 深度图不是第四个通道先搞清几何信号和纹理信号的差异再决定输入策略2.1 深度图与RGB在统计特性上的三个本质差异很多做纯RGB分割的工程师拿到RGB-D相机后的第一反应是把深度图归一化到0-255然后直接拼在RGB后面当成四通道输入。这是最省事的做法也最容易翻车。原因在于深度图和RGB表面上看都是二维图像但它们的统计特性差别非常大。第一RGB是三个通道的反射亮度值域相对稳定局部平滑深度图是相机到表面的距离单通道值域由传感器量程决定常见0.5到8米噪声模型跟距离强相关。Kinect一代的深度标准差随距离近似平方增长距离越远测距值的抖动越明显。第二深度图天然稀疏。黑色吸光表面、镜面、透明材质、超出量程的近处物体会产生无效像素通常标记为0或NaN。如果直接把0当有效最小距离卷积核会学到「离相机最近的是黑窟窿」这种错误先验。第三深度图的边缘语义和RGB边缘不一样。遮挡边界处深度会突变但物体表面纹理变化的地方深度是完全平滑的反过来反光表面在RGB里有高亮边缘深度图对应位置却直接缺失。一个在纯RGB数据集上表现良好的卷积神经网络默认并不知道这两种边的区别。所以RGB-D分割的核心工作首先不是改网络结构而是想清楚深度信号以什么形态进网络。我一般会准备两个输入通道给深度分支一个归一化深度图一个valid mask表示像素深度是否有效。这个预处理改变往往比换一个更大的骨干网络提升更明显。2.2 三种融合策略的取舍early fusion、late fusion 和多级融合RGB和深度信息在哪个位置融合是决定整个模型行为的分水岭。常见做法可以分成三类第一种是early fusion把RGB三通道和深度单通道拼接成四通道输入后面接一个常规分割网络。优点是实现简单、显存开销小、代码改动最少缺点是网络第一层卷积的权重同时被RGB和深度「拉扯」而两种输入的分布差异太大预训练权重基本只能迁移到RGB部分深度通道等于从零学。个人经验是除非训练数据量特别大否则early fusion很容易出现深度分支学不动的情况。第二种是late fusion两个分支各自编码到最后再拼接。这样两个分支互不干扰各自的预训练权重都能发挥作用缺点是几何细节在浅层编码阶段就走丢了等到最后融合时深度图最有价值的遮挡边界信息已经被池化抹得差不多。实测结果通常是边缘碎、小目标召回率低。第三种是多级融合两个分支各自独立编码同时在多个分辨率层级做特征拼接或相加。这是RGB-D分割里最常用、效果最稳的做法。浅层融合保留深度突变边界深层融合补充语义解码器再逐步恢复分辨率。三种方式的对比如下融合方式实现难度预训练迁移边界质量显存开销适用场景early fusion低RGB可迁移深度难一般小数据量大时的快速验证late fusion中双流各自独立边缘碎中特征差异极大、需要分支解耦多级融合中高双流独立且逐层互补好大精度优先的正式项目2.3 编码器骨架怎么选双流ResNet、轻量化共享编码器与图神经网络的边界网络骨架的选择常见做法是RGB流和深度流各用一个ResNet34或ResNet50作为编码器预训练权重在ImageNet上加载。需要注意深度流加载预训练权重的收益远小于RGB流因为深度图的局部统计特征跟自然图像完全不同。我一般会让深度流的第一层卷积重新初始化或者只迁移浅层几个卷积的权重其他层随机初始化并配一个稍高的学习率。这样做比强行让网络适应一个错误初始值要快得多。如果显存紧张可以改成共享前两层浅层编码器、中后期再分叉的结构。因为浅层学到的是边缘、角点这类通用几何特征RGB和深度都适用深层语义差异大分开编码更合理。另外图神经网络近年来在物体关系建模、点云语义理解里很热但在RGB-D逐像素分割这种稠密预测任务上工程落地还是以卷积神经网络为主。要把图神经网络作为主分割通路需要构建像素级或超像素级的图结构推理成本和工程复杂度都翻倍实际收益往往不如多级融合来得直接。它更适合作为辅助分支处理长距离依赖而不是替代CNN主干。2.4 融合点的位置与融合方式浅层补细节、深层补语义融合点选在哪一层决定了深度几何信息能不能活到解码阶段。我的经验是浅层特征分辨率高RGB纹理细节丰富深度图的遮挡边界在浅层最清楚深层特征语义强、感受野大但深度细节已经被下采样抹平。所以融合不应该只做一次而是每个编码器层级各做一次和U-Net的skip connection配合使用。融合操作本身也有讲究。拼接concat是最常用的信息保留完整但通道数会翻倍后面通常跟一个1x1卷积降维相加add实现简单、显存友好但要求两个分支的特征图通道数一致且响应尺度要相近。RGB分支经过BN和ReLU后的激活值范围通常比深度分支大不少直接相加容易被RGB侧主导。一个稳妥做法是在相加前给深度分支加一个可学习的缩放权重初始值设为0.5让网络自己学两个分支的可信度。另一个容易被忽略的细节是融合前的深度特征降噪。深度图在遮挡边界附近有±1像素的随机偏移直接拼接会放大边缘噪声。我一般会在融合前对深度分支的浅层特征加一个3x3的中值滤波或者用一层可分离卷积做边缘保持平滑能明显改善分割边界的稳定度。3. 最小可复现方案用PyTorch把NYU Depth V2跑出第一版分割模型3.1 数据准备NYU Depth V2的读取、对齐与预处理NYU Depth V2是RGB-D分割最常用的公开数据集1449张对齐好的RGB-D图像40类语义标签。先从预处理说起代码里已经把最容易踩的坑标出来了import h5py import cv2 import numpy as np def load_nyu_sample(path, index, size(320, 240)): with h5py.File(path, r) as f: rgb f[images][index].transpose(2, 0, 1) # (3, 480, 640), RGB顺序 depth f[depths][index].astype(np.float32) # (480, 640), 单位是米 label f[labels][index].astype(np.int64) # (480, 640), 0为无效, 1~40为类别 h, w depth.shape rgb cv2.resize(rgb, size, interpolationcv2.INTER_LINEAR) depth cv2.resize(depth, size, interpolationcv2.INTER_NEAREST) label cv2.resize(label, size, interpolationcv2.INTER_NEAREST) max_depth 8.0 depth np.clip(depth / max_depth, 0.0, 1.0) depth[depth 0.0] 1.0 # 无效深度填到最远距离 valid (depth 1.0).astype(np.float32) # valid mask: 1为有效 return rgb, depth, valid, label这套预处理有四个关键点。第一RGB和深度在NYU里已经用相机内参对齐过了直接resize即可但如果是自己用Kinect或RealSense采集的数据必须先把深度图做重投影对齐否则后续融合全是错位。第二深度图和标签用INTER_NEAREST最近邻插值不要用线性插值线性插值会把深度突变边界抹平还会把标签插出原本不存在的类别值。第三深度归一化的分母max_depth固定为8.0这是传感器量程映射不要每帧动态除以当前帧最大值否则网络学到的是相对距离而不是绝对距离。第四valid mask作为额外通道输入深度分支告诉网络哪些位置深度可信。3.2 双流编码器与逐层融合一个可以直接跑的结构下面给一个精简但完整的双流分割网络。结构上就是两个独立编码器每层做完卷积后把RGB和深度特征拼起来融合import torch import torch.nn as nn import torch.nn.functional as F class ConvBN(nn.Module): def __init__(self, in_c, out_c): super().__init__() self.block nn.Sequential( nn.Conv2d(in_c, out_c, 3, padding1, biasFalse), nn.BatchNorm2d(out_c), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.block(x) class Fuse(nn.Module): def __init__(self, in_c, out_c): super().__init__() self.conv ConvBN(in_c * 2, out_c) # 拼接后降维 def forward(self, a, b): return self.conv(torch.cat([a, b], dim1)) class RGBDSegNet(nn.Module): def __init__(self, n_classes40): super().__init__() # 两条独立编码器深度分支输入2通道(深度valid mask) self.rgb_e1 ConvBN(3, 32) self.rgb_e2 ConvBN(32, 64) self.rgb_e3 ConvBN(64, 128) self.depth_e1 ConvBN(2, 32) self.depth_e2 ConvBN(32, 64) self.depth_e3 ConvBN(64, 128) # 三个层级融合 self.fuse1 Fuse(32, 32) self.fuse2 Fuse(64, 64) self.fuse3 Fuse(128, 128) # 轻量解码器 self.up3 nn.ConvTranspose2d(128, 64, 2, stride2) self.up2 nn.ConvTranspose2d(64, 32, 2, stride2) self.head nn.Conv2d(32, n_classes, 1) def forward(self, rgb, depth): x F.max_pool2d(self.rgb_e1(rgb), 2) # 32通道, H/2 y F.max_pool2d(self.depth_e1(depth), 2) f1 self.fuse1(x, y) # 第一层融合 x F.max_pool2d(self.rgb_e2(x), 2) # 64通道, H/4 y F.max_pool2d(self.depth_e2(y), 2) f2 self.fuse2(x, y) x F.max_pool2d(self.rgb_e3(x), 2) # 128通道, H/8 y F.max_pool2d(self.depth_e3(y), 2) f3 self.fuse3(x, y) out F.relu(self.up3(f3)) # 解码到H/4 out F.relu(self.up2(out f2)) # 加上中层融合特征解码到H/2 out self.head(out) # (40, H/2) out F.interpolate(out, scale_factor2, modebilinear, align_cornersFalse) return out这个结构里最值得关注的是融合位置和通道配比。三个层级各做一次拼接融合浅层的f1负责遮挡边界几何细节中层的f2负责局部结构深层的f3负责语义。解码时把f2直接加回来是为了让高分辨率路径保留中层特征。通道数32/64/128是在精度和显存之间取的常见默认值如果你输入分辨率只有320x240这个体量在12GB显卡上可以轻松跑起来。forward里depth参数已经包含了valid mask也就是调用时传入的depth形状是(B, 2, H, W)。实际训练时还可以把RGB流换成预训练的ResNet34前几层作为编码器深度流保持随机初始化这样精度会进一步上来但网络结构复杂度也随之增加。3.3 训练循环损失函数、优化器与关键参数class_weights torch.tensor([0.0] [1.0] * 40, devicecuda) # 下标0忽略 criterion nn.CrossEntropyLoss(weightclass_weights, ignore_index0) optimizer torch.optim.SGD( model.parameters(), lr0.01, momentum0.9, weight_decay1e-4 ) scheduler torch.optim.lr_scheduler.PolynomialLR( optimizer, total_itersepochs, power0.9 ) for epoch in range(epochs): model.train() for rgb, depth, valid, label in loader: rgb, depth, label rgb.cuda(), depth.cuda(), label.cuda() pred model(rgb, depth) # depth已经拼接valid mask loss criterion(pred, label) optimizer.zero_grad() loss.backward() optimizer.step() # 每个epoch结束跑一次验证集的mIoU两个参数值得单独说明。CrossEntropyLoss的ignore_index0用来跳过未标注区域这比手动mask损失要干净class_weights的位置0设0.0保证忽略类不会参与梯度计算。优化器选SGD加momentum而不是Adam在分割任务上SGD配合poly学习率衰减的收敛曲线更平滑最终精度通常也更好。PolynomialLR的power0.9是DeepLab系列常用的配置等价于学习率按(1 - iter/total)^0.9衰减到0。3.4 推理与可视化从logits到彩色分割图palette np.array([ [128, 128, 128], [128, 0, 0], [128, 0, 128], # 固定40类颜色映射 # ... 实际工程里写满40行 ]) model.eval() with torch.no_grad(): logits model(rgb, depth) # (1, 40, H, W) pred logits.argmax(dim1).squeeze(0).cpu().numpy() vis palette[pred] # (H, W, 3)argmax直接取最大logits对应的类别索引效果和softmax后取argmax完全等价省掉一次指数运算。palette必须是固定数组不能每次推理随机生成否则同一类别在不同帧里颜色不一致后续人工审查分割质量时非常痛苦。可视化文件可以用cv2.imwrite直接存PNG注意CV2默认是BGR顺序调色板切换到RGB顺序后需要转一下通道。4. 训练策略与调参顺序损失函数、增强方式和BN设置该怎么配4.1 损失函数怎么配加权交叉熵、Dice Loss与边界惩罚NYU Depth V2的40个类别里床、沙发这类大物体占了大量像素杯子、书本、台灯这些小物体像素占比极低。直接用默认的CrossEntropyLoss网络会偏向大类别小目标mIoU经常是0。常见做法是给每个类别设置权重具体用median frequency balancing统计每个类别在训练集中的像素频率然后用中位数频率除以各类别频率再截断到[0.01, 10]之间。直接inverse frequency会被极小类别的权重顶到几千训练直接震荡。加权交叉熵对类别不均衡很有效但对边界质量没有直接帮助。分割任务的另一个常见做法是叠加Dice Loss它对像素数量不敏感天然适合小目标class DiceLoss(nn.Module): def __init__(self, n_classes, ignore_index0): super().__init__() self.n_classes n_classes self.ignore_index ignore_index def forward(self, logits, target): pred F.softmax(logits, dim1) mask target ! self.ignore_index # 忽略0类 target_onehot F.one_hot(target * mask.long(), num_classesself.n_classes).permute(0, 3, 1, 2).float() intersection (pred * target_onehot).sum(dim(2, 3)) union pred.sum(dim(2, 3)) target_onehot.sum(dim(2, 3)) dice (2.0 * intersection 1.0) / (union 1.0) return 1.0 - dice.mean()这里平滑因子1.0是经验值目的是防止某些类别在训练初期完全没预测对时分母为0。建议把最终损失配成加权交叉熵和Dice的组合loss 0.8 * ce_loss 0.2 * dice_loss前几个epoch可以只跑CE等类别分布稳定后再引入Dice项。如果对边缘锐度有要求还可以在GT边界做一次Canny或拉普拉斯提取对边界像素的CE loss加权2.0但这个操作会显著拖慢收敛属于锦上添花。4.2 数据增强RGB和深度必须用同一套几何变换RGB图像增强的标准套路到RGB-D这里有一半要推翻。随机翻转、随机缩放、随机裁剪、随机旋转这些几何变换RGB和深度必须用完全同一组参数如果分别随机做RGB和深度对齐关系就破坏了。实现时把变换矩阵先算好再同时应用到两个输入上。深度图增强有几个专属细节。第一几何变换的插值方式要注意RGB用线性插值没有副作用深度图用最近邻否则深度边界被平滑后分割边缘的锐度会下降。第二归一化不要动态做训练时固定max_depth8.0inference时用同一个值模型才学到一致的绝对距离语义。第三可以给深度图加随机高斯噪声对深度值做随机腐蚀膨胀模拟不同传感器硬件差异这个技巧对跨设备泛化非常管用。另一个容易犯的错是把RGB的color jitter亮度、对比度、饱和度扰动直接套到深度图上。深度图反映的是几何距离不受光照影响对它做颜色扰动只会引入噪声。HSV抖动只作用于RGB分支深度分支保持几何变换同步即可。4.3 训练参数细节双流结构下的学习率、BN与batch size双流网络参数量大不同部位的收敛速度差异很大。常见做法是把优化器参数分组预训练加载的RGB分支主干用lr1e-4随机初始化的深度分支和融合层用lr1e-3解码器也取1e-3。如果只设一个全局学习率要么RGB分支被扰动得太厉害要么深度分支学得太慢。BatchNorm在双流结构里有自己的坑。BN统计量依赖batch内的样本分布双流等于每个分支各有一组统计量batch size太小统计噪声会很大。在8或更小的batch下我一般会把标准BN替换成GroupNorm或者用PyTorch的SyncBatchNorm在更大规模上同步统计量。显存不够时梯度累积也是一个可行方案累积4个step的梯度再更新一次参数等效于把batch size放大4倍注意此时BN统计量仍然按原始batch计算所以它解决的是优化稳定性问题解决不了BN的统计噪声问题。学习率调度建议直接用poly策略而不是StepLR。分割任务训练后期需要精细微调StepLR在每段step边界处容易把loss弹出来poly是渐进衰减到0曲线平滑得多。训练初期观察一下深度分支和RGB分支各自的loss下降速度如果深度分支长时间不下降先怀疑输入预处理而不是网络结构。4.4 评估指标pixel accuracy、mIoU和边界F1pixel accuracy在背景占比高的分割数据集上非常骗人。假设一张图70%是墙面模型把所有像素都预测成墙面pixel accuracy也有70%看起来不错但桌子、椅子全部丢失。所以评估必须看mIoU按类别逐类算IoU再取平均对每个类别的贡献是平等的。计算时ignore_index0的区域要排除否则未标注像素会稀释所有类别的IoU。除了mIoU如果应用场景对边缘敏感建议额外看Boundary F1。计算方法是分别提取预测和GT的边缘像素看边缘像素集合的重叠程度它能捕捉到mIoU里难以反映的「物体整体正确但边缘锯齿明显」问题。不同的网络结构在mIoU上可能只差零点几个点但Boundary F1可能差出三五个点这个差距在工业检测场景里往往更致命。5. 避坑记录RGB-D分割训练中最常踩的五个坑与排查方法5.1 深度分支loss不下降分割结果等于纯RGB分割现象加了深度分支后验证集mIoU和单用RGB几乎一样甚至稍微更差训练日志里深度分支对应的loss长期不变化。原因最常见的是深度输入里大量0值被当成了有效近距。自己用Kinect采集的数据黑色物体和远处背景会有很多无效深度0网络经过若干轮训练发现深度通道不可信于是学到一个恒等跳过深度分支的局部最优。另一个常见原因是valid mask没有作为输入传给网络深度分支被迫从一堆无效值里猜测信号。解决无效深度统一填充到max_depth对应的最大值1.0同时把valid mask作为额外通道输入让网络显式知道哪些位置深度可信。训练时打印深度分支第一层卷积的激活均值如果始终接近0说明输入信号没有被利用此时优先复查数据预处理管线而不是换网络结构。5.2 分割边缘锯齿明显同一物体边缘内外波动现象可视化结果里物体边缘呈锯齿状尤其出现在RGB边缘和深度边界重叠的位置人眼一眼就能看出轮廓不对。原因深度图在遮挡边界处存在±1像素的随机偏移RGB和深度的对齐在硬件上不可能完全精确融合时网络把深度突变当成了强特征边缘就被拉成了噪声形状。解决在融合前对深度分支的浅层特征用3x3中值滤波或双边滤波做边缘保持平滑如果不想加额外计算可以在融合层后接一个可学习的attention gate让网络自己学习RGB和深度特征在每个位置的可信度。实测里这个操作比调整融合权重阈值有效得多。5.3 batch size只有4或者更小BN统计量崩了现象训练loss震荡明显验证集mIoU在epoch之间大幅跳动甚至验证时用batch的BN统计量跑出比合理值低5个点以上的结果。原因双流结构让BN统计量翻倍batch4时每个统计量只由4张图估计噪声太大。小batch下BN的running mean和running val不稳定是根因。解决把小batch下的BN全部替换成GroupNorm分组数取8或16模型就能在batch2下稳定训练如果必须用BN改用SyncBatchNorm并确保多卡同步或者先在大batch下预训练一段再切到小batch微调。显存允许时输入分辨率降到224x224也能缓解batch过小的问题。5.4 小目标类别mIoU为0或者接近0现象杯子、书本、台灯这类小物体在训练集里样本不少但验证时mIoU几乎为0分割结果里完全找不到这些类别。原因类别不均衡是主因默认CrossEntropy对大类别像素多的类更友好另外MaxPooling逐层下采样会让小目标的特征在下采样两次后直接被邻域大目标的特征淹没。解决损失换成加权CE加Dice的组合权重用median frequency balancing数据增强里增加随机小尺寸crop让小目标在一个crop里占更大比例如果显存够可以把编码器的前两层stride从2改成1减少小目标空间信息的损失。这三个手段叠加后小目标mIoU通常能拉回到正常水平。5.5 验证集mIoU很好换一个传感器或场景直接掉点现象在NYU Depth V2划分的验证集上mIoU达到不错水平但用另一款深度相机或者换一个房间采集的数据测试指标掉了两三个点。原因不同深度传感器的噪声分布差异很大包括深度精度、孔洞比例、边界偏移量如果训练集来自单一设备场景模型已经对这套硬件的噪声形态过拟合了。解决训练时对深度图做传感器域随机化——随机加高斯噪声、随机对深度图做腐蚀膨胀、随机把一部分有效像素置为无效再填到max_depth。这套数据增强不用采集任何新数据就能让模型对深度噪声的形态不再敏感。这是RGB-D分割工程落地里投入产出比最高的一个技巧。6. 从分割掩码到三维点云CRF后处理、反投影与部署形状分割网络输出mask后工作远没有结束。我习惯把mask反投影到三维点云里做进一步处理这样分割结果才真正变成可用的几何信息。核心代码是把深度图按相机内参反算回三维坐标def mask_to_pointcloud(depth, mask, fx, fy, cx, cy): h, w depth.shape ys, xs np.meshgrid(np.arange(h), np.arange(w), indexingij) z depth x (xs - cx) * z / fx y (ys - cy) * z / fy pts np.stack([x, y, z], axis-1).reshape(-1, 3) return pts[mask.reshape(-1) 0]fx、fy、cx、cy是相机内参从传感器标定结果里读取。反投影前我通常会先对mask做一次形态学闭运算3x3或5x5矩形核把分割边缘的孤立小洞补掉再做一个开运算去掉孤立的噪点区域然后再投影。这两步比在网络上加后处理头划算得多CV的形态学操作几十行就能完成。关于CRF后处理DenseCRF确实是经典后悔药在离线场景下能把mIoU提升0.5到2个点尤其是边缘锐度改善明显。但实时推理时CRF的迭代优化开销很大如果模型要跑在边缘设备上我一般放弃CRF改用测试时增强——把输入水平翻转再推理一次两次softmax取平均再argmax。虽然推理耗时翻倍但稳定性提升明显边缘也不会因为单次随机噪声出现异常抖动。部署时还有两个现实问题。一是双流模型体积通常比单流大50%以上嵌入式端跑不动常见做法是训练时保持双流部署时用知识蒸馏把深度分支压成一个浅层网络让RGB分支和浅层深度分支的组合在推理阶段更快。二是深度归一化、valid mask生成这些预处理算子必须跟着模型一起部署进推理框架的预处理管线不能在应用层用另一套语言重写否则参数不一致会导致推理结果和训练时对不上。我早期做RGB-D分割时最先改的是网络结构注意力机制、图神经网络分支都试过涨点都有限。后来发现把数据预处理和融合层位置调对再配一个简单的后处理涨点比任何结构改动都大。从那以后接到新需求我第一件事是检查深度输入格式、归一化方式和融合点位置再考虑要不要改网络。这个顺序能帮你避开大半无用功希望这些经验能帮到你。本文还有配套的精品资源点击获取