简介这份资源面向计算机视觉方向的研究者与开发者聚焦基于神经网络的RGB-D图像分割任务通过融合彩色图像与深度信息提升复杂几何场景下的分割精度可应用于机器人导航、增强现实与三维重建等场景。项目以DepthAware CNN为核心涵盖预处理融合、基础网络特征提取、深度处理模块、融合解码与分类后处理等关键环节并支持调整网络结构、适配自建数据集、优化超参数、更换损失函数以及评估可视化等二次开发。压缩包共70个文件以41个Python脚本为主体辅以C、CUDA与头文件等底层实现另有少量编译产物与说明文档整体约86KB结构紧凑便于快速上手。目前已有1413人学习下载适合希望深入理解RGB-D分割算法并在此基础上开展定制实验的读者参考。1. 从一张 RGB 图加一张深度图说起这套分割方案到底解决什么问题做机器人抓取、室内导航或者 AR 遮挡的同行大概率都遇到过同一个尴尬单靠 RGB 图像做分割遇到白墙、纯色桌面、光照突变就翻车模型把背景和前景糊成一团。RGB-D 图像分割就是冲着这个痛点来的——在彩色三通道之外再叠一路深度通道让网络同时看到「长什么样」和「离多远」。这次拆的资源是一套基于神经网络的 RGB-D 语义分割实现核心思路是把深度图当作额外特征与 RGB 分支做融合输出逐像素类别。它适合三类人正在做多模态融合课程设计的学生、需要快速搭一个分割 baseline 的算法工程师、以及想搞清楚深度信息到底怎么进网络的自学者。整套代码结构清晰改起来不费劲下面按「是什么 → 怎么跑 → 坑在哪 → 怎么调」的顺序拆开讲。2. 网络结构与数据流RGB 和 Depth 到底在哪一层握手2.1 双分支编码器的选型理由这套实现走的是双分支编码器路线RGB 走一个分支Depth 走另一个分支各自提特征然后在中间层做融合。为什么不在输入层直接拼成四通道我试过四通道输入在浅层还能用但深度图和 RGB 的统计分布差异太大——RGB 是 0 到 255 的纹理强度深度是米为单位的距离值直接 concat 会让 BN 层的均值方差被深度值带偏训练前期 loss 震荡得厉害。双分支的好处是每个模态有自己的 BN 和卷积核互不干扰融合放在特征已经抽象过的层级语义对齐更自然。常见做法是编码器用 ResNet 或 MobileNet 做 backbone把最后的分类头去掉取 stage3 和 stage4 的特征图。深度分支结构可以照搬 RGB 分支也可以砍掉一半通道数省显存。融合方式有几种逐元素相加、通道 concat 后 1x1 卷积降维、或者用注意力模块算权重。这套资源里用的是 concat 1x1 卷积实现简单效果在 NYU Depth v2 这类室内数据集上够用。2.2 融合模块的实现细节融合层的位置很关键。放太浅两个模态特征还没抽象好融合等于白搭放太深空间分辨率已经降了 32 倍小物体分割精度掉得厉害。我一般会在 stage3 输出stride 16和 stage4 输出stride 32各做一次融合然后上采样到原图尺寸做逐像素分类。import torch import torch.nn as nn class FusionBlock(nn.Module): RGB 与 Depth 特征融合concat 后 1x1 卷积降维 def __init__(self, rgb_channels, depth_channels, out_channels): super().__init__() self.rgb_proj nn.Conv2d(rgb_channels, out_channels, 1) self.depth_proj nn.Conv2d(depth_channels, out_channels, 1) self.fuse nn.Sequential( nn.Conv2d(out_channels * 2, out_channels, 1), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) def forward(self, rgb_feat, depth_feat): # 先各自投影到同一通道空间避免 concat 后通道爆炸 r self.rgb_proj(rgb_feat) d self.depth_proj(depth_feat) # 深度特征上采样对齐 RGB 空间尺寸 if d.shape[-2:] ! r.shape[-2:]: d nn.functional.interpolate( d, sizer.shape[-2:], modebilinear, align_cornersFalse) return self.fuse(torch.cat([r, d], dim1))这段代码里三个参数要盯住out_channels控制融合后的特征维度一般设 256 或 512depth_proj的输入通道取决于深度分支 backbone 的输出interpolate的align_cornersFalse是 PyTorch 分割任务里的标准做法设 True 会导致边缘像素偏移。融合完的特征再送进解码器解码器可以用简单的双线性上采样加卷积也可以用 FPN 那种横向连接结构。2.3 数据加载与深度图预处理深度图的预处理是这套流程里最容易埋雷的地方。原始深度图可能是 16 位 PNG单位是毫米也可能是 32 位浮点单位是米。网络输入前必须统一量纲否则不同数据集之间迁移直接崩。我一般会把深度值裁剪到 [0.5m, 5m] 区间然后归一化到 [0, 1]超出范围的置零当作无效深度。import numpy as np import cv2 def preprocess_depth(depth_path, min_depth0.5, max_depth5.0): 读取深度图并归一化到 [0,1]无效值置零 depth cv2.imread(depth_path, cv2.IMREAD_UNCHANGED) if depth is None: raise FileNotFoundError(depth_path) # 16 位 PNG 通常是毫米单位转成米 if depth.dtype np.uint16: depth depth.astype(np.float32) / 1000.0 else: depth depth.astype(np.float32) # 裁剪并归一化 depth np.clip(depth, min_depth, max_depth) depth (depth - min_depth) / (max_depth - min_depth) # 无效深度0 或 NaN置零 depth[~np.isfinite(depth)] 0.0 return depthmin_depth和max_depth要根据你的传感器调。Kinect v2 的有效范围大概 0.5 到 4.5 米RealSense D435 可以到 10 米但远距离噪声大。归一化之后深度图和 RGB 图一样是单通道浮点张量送进 DataLoader 时记得和 RGB 做相同的随机裁剪、翻转增强但深度图不能用颜色抖动。3. 训练流程与损失函数从数据到权重的完整链路3.1 损失函数选型与类别不平衡处理语义分割的标配损失是交叉熵但室内数据集类别极度不平衡——墙面和地板占了 60% 以上像素椅子、书本这些小类经常被淹没。直接上交叉熵模型会学会把所有像素预测成多数类准确率看着高IoU 惨不忍睹。常见做法是交叉熵加 Dice Loss 做加权组合或者用 Focal Loss 压制易分类样本。class CombinedLoss(nn.Module): 交叉熵 Dice Loss缓解类别不平衡 def __init__(self, num_classes, ce_weight0.5, dice_weight0.5): super().__init__() self.ce nn.CrossEntropyLoss(ignore_index255) self.dice_weight dice_weight self.ce_weight ce_weight self.num_classes num_classes def dice_loss(self, logits, targets): probs torch.softmax(logits, dim1) targets_onehot nn.functional.one_hot( targets, self.num_classes).permute(0, 3, 1, 2).float() dims (0, 2, 3) intersection torch.sum(probs * targets_onehot, dims) cardinality torch.sum(probs targets_onehot, dims) dice (2. * intersection 1e-6) / (cardinality 1e-6) return 1 - dice.mean() def forward(self, logits, targets): ce_loss self.ce(logits, targets) d_loss self.dice_loss(logits, targets) return self.ce_weight * ce_loss self.dice_weight * d_lossignore_index255是把未标注像素排除在损失计算外NYU Depth v2 里有些区域没有标签不设这个参数会引入噪声梯度。ce_weight和dice_weight我一般从 0.5 比 0.5 起步如果小类 IoU 还是低把 dice 权重提到 0.7。3.2 优化器配置与学习率调度优化器用 AdamW 还是 SGD这套实现里我建议 AdamW初始学习率 1e-4weight decay 1e-4。SGD 虽然泛化可能好一点但需要更精细的调参对新手不友好。学习率调度用 cosine annealing训练 100 个 epoch前 5 个 epoch 做 warmup 线性升温。from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR optimizer AdamW(model.parameters(), lr1e-4, weight_decay1e-4) warmup LinearLR(optimizer, start_factor0.01, total_iters5) cosine CosineAnnealingLR(optimizer, T_max95, eta_min1e-6) scheduler SequentialLR(optimizer, schedulers[warmup, cosine], milestones[5])start_factor0.01表示 warmup 从 1e-6 开始避免训练初期梯度爆炸。T_max95是 cosine 周期的总步数和总 epoch 减去 warmup 对齐。每个 epoch 结束后调scheduler.step()注意 SequentialLR 的 milestone 是切换点不是总长度。3.3 训练循环与验证指标训练循环里除了常规的 forward-backward-step还要定期在验证集上算 mIoU。mIoU 比像素准确率靠谱得多尤其在你关心小类分割效果的时候。计算时用混淆矩阵累加最后再统一算不要在 batch 级别平均否则每个 batch 类别数不一样会引入偏差。def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss 0 for imgs, depths, labels in loader: imgs, depths, labels imgs.to(device), depths.to(device), labels.to(device) optimizer.zero_grad() logits model(imgs, depths) loss criterion(logits, labels) loss.backward() # 梯度裁剪防止深度分支梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() total_loss loss.item() return total_loss / len(loader)clip_grad_norm_的max_norm1.0是经验值深度图归一化后梯度一般不会太大但融合层 concat 后通道翻倍偶尔会冲高加个裁剪稳一点。验证时记得model.eval()和torch.no_grad()不然显存直接翻倍。4. 避坑与排查那些让我重跑过训练的血泪经验4.1 深度图全黑或全白现象训练 loss 不下降可视化深度输入发现整张图接近 0 或 1。原因深度图读取时没做单位转换16 位 PNG 直接当 0-255 处理或者归一化区间设错把有效深度全裁掉了。解决先打印深度图的 min/max/mean确认原始量纲用cv2.IMREAD_UNCHANGED读 16 位图归一化区间根据传感器手册设别照搬别人的 0.5-5m。4.2 融合后特征图尺寸对不上现象torch.cat报错提示两个张量除通道外其他维度不一致。原因RGB 和 Depth 分支的输入分辨率不同或者下采样倍数不一致导致同一 stage 输出的空间尺寸有偏差。解决在融合前统一做interpolate对齐到 RGB 特征的尺寸别假设两个分支天然对齐。如果深度分支用了不同的 backbonestride 可能差一倍要手动补下采样或上采样。4.3 验证集 mIoU 远低于训练集现象训练集 mIoU 冲到 0.7验证集只有 0.3。原因最常见的是深度图增强和 RGB 不同步——RGB 做了随机裁剪深度图没做或者做了不同的随机种子。另一个原因是验证集的深度图分布和训练集不同比如训练用 Kinect验证用 RealSense。解决把 RGB 和 Depth 绑在同一个 transform 里用相同的随机参数跨传感器时先做深度直方图对齐或者干脆在目标域上微调几个 epoch。4.4 显存溢出但 batch size 已经降到 1现象batch size 设为 1 还是 OOM。原因融合层 concat 后通道数翻倍中间特征图显存占用比单分支高不少另外如果解码器用了全分辨率特征显存峰值在最后几层。解决把融合后的通道数从 512 降到 256解码器用逐步上采样而不是一步到原图开启混合精度训练torch.cuda.amp能省 30% 到 40% 显存。4.5 深度分支梯度消失现象训练几个 epoch 后深度分支的权重几乎不变RGB 分支正常更新。原因深度图归一化后数值范围小经过多层卷积后梯度衰减快或者深度分支学习率没单独设跟着全局学习率走太小。解决给深度分支单独设一个更大的学习率比如 RGB 用 1e-4Depth 用 5e-4或者在深度分支加残差连接让梯度有捷径可走。5. 进阶调优把 mIoU 再往上推几个点的具体手法5.1 深度图补全与噪声抑制消费级深度相机在物体边缘和反光表面会丢深度这些空洞区域如果直接置零送进网络等于告诉模型「这里没有信息」分割边缘就会毛糙。我一般会做一个简单的形态学闭运算填小洞大洞用邻域中值填充。更讲究的做法是跑一个深度补全网络但那又是另一个模型了课程设计级别没必要。def fill_depth_holes(depth, kernel_size5): 形态学闭运算填充小空洞 kernel cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) # 先把无效区域标记出来 invalid (depth 0).astype(np.uint8) # 闭运算填小洞 filled cv2.morphologyEx(depth, cv2.MORPH_CLOSE, kernel) # 只替换原来无效的像素有效区域保持原值 result depth.copy() result[invalid 1] filled[invalid 1] return resultkernel_size别设太大5 或 7 就够了再大边缘会被糊掉。这个操作在预处理阶段做一次就行不用每个 epoch 重复。5.2 多尺度融合与注意力加权如果基础版 mIoU 卡在 0.5 上不去可以试试在融合模块里加通道注意力。SE 模块的实现很简单全局平均池化后过两个全连接层算通道权重再乘回特征图。我实测在 NYU Depth v2 上能涨 1 到 2 个点代价是参数量增加不多。class SEBlock(nn.Module): 通道注意力全局池化 两层 FC 算权重 def __init__(self, channels, reduction16): super().__init__() self.pool nn.AdaptiveAvgPool2d(1) self.fc nn.Sequential( nn.Linear(channels, channels // reduction), nn.ReLU(inplaceTrue), nn.Linear(channels // reduction, channels), nn.Sigmoid() ) def forward(self, x): b, c, _, _ x.shape w self.pool(x).view(b, c) w self.fc(w).view(b, c, 1, 1) return x * wreduction16是原论文的默认值通道数少的时候可以改成 8。把 SEBlock 插在融合层之后、解码器之前效果最明显。5.3 验证 mIoU 的正确计算方式最后说一个容易被忽略的点mIoU 的计算。很多人用sklearn.metrics.jaccard_score逐 batch 算再平均这是错的因为每个 batch 出现的类别不一样平均权重不对。正确做法是维护一个num_classes x num_classes的混淆矩阵整个验证集跑完后一次性算。def compute_miou(confusion_matrix): 从混淆矩阵算 mIoU intersection np.diag(confusion_matrix) union confusion_matrix.sum(axis1) confusion_matrix.sum(axis0) - intersection # 避免除零union 为 0 的类别跳过 iou np.where(union 0, intersection / union, np.nan) return np.nanmean(iou)np.nanmean会自动跳过没有出现的类别比手动过滤干净。混淆矩阵在验证循环里用confusion_matrix np.bincount(...)累加注意把 ignore_index 的像素排除掉。从那以后我每次跑分割训练都强制在第一个 epoch 结束后可视化一遍深度输入和融合特征图确认没有全黑、没有尺寸错位、没有梯度消失再往下跑。这三步检查花不了五分钟但能省下重跑一整晚的电费。希望帮到你。本文还有配套的精品资源点击获取
