简介面向计算机视觉学习者与开发者的RGB-D图像分割项目代码包围绕彩色图像与深度信息融合及Depth-Aware CNN展开适合希望理解三维场景语义分割、开展网络改造与二次开发的中高级实践者。压缩包共70个文件、约86KB以41个Python源码为主另含少量编译缓存、C/CUDA源码、头文件、Shell脚本及许可和说明文档整体覆盖数据读取、模型定义、损失函数、训练测试入口与可视化工具等模块目录分层较清晰。已有1414人学习下载。读者可从中获得深度图与RGB特征融合的实现思路、基础骨干网络与深度处理模块的组合方式以及跳跃连接、上采样解码和IoU等指标评估的参考也可基于现有框架替换数据集、调整网络结构与损失函数并参考训练测试流程完成可复现实验用于机器人导航、增强现实或3D重建等场景的算法改进。1. 深度图不是第四通道这套双流 CNN 工程解决的是对齐与融合问题“把深度图拼到 RGB 后面变成四通道丢进一个分割网络mIoU 经常只比纯 RGB 高一个点甚至掉点。”这是我拆这类 RGB-D 图像分割项目时最常见的翻车现场。深度图的价值不在多一个通道而在它和彩色图的空间对应关系。这套基于神经网络的 RGB-D 分割工程核心是用双流 CNN 分别吃 RGB 和 Depth在网络中间做融合最后输出逐像素类别。它适合做室内场景理解、机器人抓取前的物体区域提取、AR 遮挡判断也适合缺陷检测里需要深度辅助的像素级分类。如果你手里有对齐好的 RGB-D 数据和一张像样的显卡照着工程结构跑一遍比从零搭编码器省事。2. 工程骨架与数据管线从目录分层到 RGB-D 对齐采样RGB-D 分割项目最容易乱的地方不是网络而是数据和模型耦合在一起。今天换个深度截断阈值明天加一种增强后天发现标签映射错了代码一改就是十几处。我一般会先把工程按职责切开让配置、数据、模型、损失各管一件事训练入口只负责串流程。下面这个目录分层是这类项目里比较常见、也比较好扩展的做法。2.1 目录分层configs、datasets、models、losses 各管什么目录/文件职责二次开发常改点configs/YAML 配置数据路径、backbone、融合位置、训练超参新增数据集配置、调 batch_size 和学习率datasets/Dataset 与 transform负责 RGB-D 读取、对齐和增强类别映射、深度尺度、无效值处理models/backbone、融合模块、分割头换 backbone、改融合层通道数losses/CE、Dice、组合损失类别权重、ignore_indexutils/指标、日志、学习率、可视化mIoU 计算、颜色表train.py训练入口分布式、混合精度、断点续训eval.py评估入口测试集推理、导出混淆矩阵infer.py单张或文件夹推理输出彩色 mask、叠加原图这样分的好处是换数据集时基本只动datasets/和configs/换模型时只动models/和配置里的backbone字段。训练脚本不关心深度图是 16 位 PNG 还是npy也不关心融合模块是相加还是拼接。配置文件里通常会把num_classes、ignore_index、depth_max、crop_size这几个参数显式写出来避免后面调参时靠记忆。configs/下我一般按数据集分文件比如configs/rgbd_indoor.yaml、configs/rgbd_selfcollect.yaml。每个配置里至少要有数据根目录、训练/验证划分文件、RGB 均值和方差、深度截断上限、backbone 名字、预训练权重路径、融合方式、损失权重、优化器和学习率策略。配置文件不要写得太深两层 key 够用太深了改一个参数要翻半天。models/下通常再分backbones/、fusion/、heads/。backbone 只负责提特征fusion 只负责把两个模态的特征合在一起head 只负责把特征图上采样到原图分辨率。这个切法在换融合模块时特别省事只要输入输出通道对齐直接替换类名就行。2.2 RGB-D 四通道读入与对齐别让深度图和彩色图错位深度图常见的存储格式是 16 位 PNG单位毫米也有数据集直接给 32 位浮点单位米。读的时候如果不统一单位后面归一化就会出玄学问题。另一个坑是无效值深度相机在玻璃、黑色物体、超出量程的区域会返回 0 或者一个很大的哨兵值。这些像素直接参与归一化会把整张图的数值分布拉偏。下面这段Dataset是我会先写出来的骨架重点看深度图的单位转换、无效值 mask 和 RGB/Depth 分开归一化。import cv2 import numpy as np import torch from torch.utils.data import Dataset from PIL import Image class RGBDDataset(Dataset): RGB-D 语义分割数据集返回 image、depth、label 和无效值 mask def __init__(self, samples, rgb_mean, rgb_std, depth_max10.0, ignore_index255): self.samples samples self.rgb_mean np.array(rgb_mean, dtypenp.float32) self.rgb_std np.array(rgb_std, dtypenp.float32) self.depth_max depth_max # 深度截断上限单位米 self.ignore_index ignore_index def __len__(self): return len(self.samples) def __getitem__(self, idx): item self.samples[idx] # 1. 读 RGB统一成三通道 rgb np.array(Image.open(item[rgb_path]).convert(RGB), dtypenp.float32) # 2. 读深度16 位 PNG 常见单位是毫米先转成米 depth cv2.imread(item[depth_path], cv2.IMREAD_UNCHANGED) if depth is None: raise FileNotFoundError(item[depth_path]) depth depth.astype(np.float32) / 1000.0 # 3. 无效深度置 0并保留 mask 供融合模块或统计使用 invalid (depth 0) | (depth self.depth_max) depth[invalid] 0.0 # 4. 读标签边界和未标注像素用 ignore_index 标记 label np.array(Image.open(item[label_path]), dtypenp.int64) # 5. RGB 按数据集均值方差归一化深度单独缩放到 [0,1] rgb (rgb / 255.0 - self.rgb_mean) / self.rgb_std depth np.clip(depth / self.depth_max, 0.0, 1.0) # 6. 转 tensor深度补一个通道维 rgb torch.from_numpy(rgb).permute(2, 0, 1).float() depth torch.from_numpy(depth).unsqueeze(0).float() label torch.from_numpy(label).long() return { image: rgb, depth: depth, label: label, invalid: torch.from_numpy(invalid), }这段代码里最关键的是第 2 步和第 3 步。depth / 1000.0把毫米转成米如果你的数据本身就是米这一行要去掉。invalid这个 mask 不只是为了清洗数据后面融合模块可以用它把无效深度的特征压掉避免网络从空洞区域学到噪声。depth_max一般设成 8 到 10 米室内场景够用太远的深度值本身也不准。RGB 的均值和方差要在训练集上统计不要直接套 ImageNet 的。深度归一化不要用 ImageNet 参数深度是单通道物理量按depth_max缩放到[0,1]更稳。标签里的ignore_index通常设 255边界像素、未标注类别都归到这里损失和 mIoU 都要排除它。2.3 Dataset 与 DataLoader把增强做成可配置的流水线RGB-D 的几何增强必须同步。随机裁剪、翻转、缩放RGB、Depth、Label 三张图要用同一组参数否则融合时特征和标签对不上训练 loss 看着在降mIoU 就是上不去。颜色抖动、亮度对比度这些只作用在 RGB 上深度图不要跟着变色。我一般不用两套 transform 分别处理 RGB 和 Depth而是写一个joint_transform输入是三张图输出也是三张图。随机裁剪可以这样写import random def joint_crop(rgb, depth, label, size(480, 640)): RGB、Depth、Label 使用同一组随机参数裁剪避免错位 h, w label.shape th, tw size if h th or w tw: raise ValueError(f图像尺寸 {h}x{w} 小于裁剪尺寸 {th}x{tw}) i random.randint(0, h - th) j random.randint(0, w - tw) rgb rgb[:, i:i th, j:j tw] depth depth[:, i:i th, j:j tw] label label[i:i th, j:j tw] return rgb, depth, label训练阶段用随机裁剪加随机水平翻转验证阶段只做 resize 或中心裁剪不要加随机增强。DataLoader的num_workers设成 4 到 8pin_memoryTrue如果显存够persistent_workersTrue能减少每个 epoch 重启 worker 的开销。RGB-D 读取比单张 RGB 慢尤其是深度图要用cv2.IMREAD_UNCHANGED读 16 位worker 数量不够时 GPU 会一直等数据训练日志里每个 iteration 的时间会忽高忽低。还有一个小细节collate_fn默认会把invalid这个 bool mask 也堆叠成 batch如果融合模块不用它可以在__getitem__里去掉减少一点传输开销。如果用就在模型 forward 里把它 resize 到对应特征图尺寸别直接拿原图尺寸的 mask 去乘深层特征。3. 双流编码器与融合模块RGB 和 Depth 在哪一层握手双流结构听起来简单真正要调的是两个问题两个分支要不要共享权重以及融合放在哪一层。这两个选择直接决定参数量、显存占用和最终 mIoU。我见过把深度图当第四通道直接喂给 ResNet 的写法也见过两个分支完全独立、最后把 logits 相加的写法效果差别不小但都不是“越复杂越好”。3.1 双分支 backbone共享权重还是独立权重常见做法是 RGB 分支用 ResNet-50 或 MobileNetV2 的预训练权重深度分支第一层卷积改成单通道输入后面的层要么随机初始化要么把 RGB 第一层权重按通道求平均复制过来。完全共享权重的情况比较少因为 RGB 是纹理和颜色统计深度是几何距离统计分布差异太大共享浅层还能忍共享到 layer3、layer4 通常会拖累 RGB 分支。如果显存紧张可以只让两个分支共享 stem 和 layer1layer2 之后分开。这样参数量比完全独立少一截又能保留深层特征的模态差异。深度分支的预训练权重不是必须的室内 RGB-D 数据量通常不大深度分支随机初始化加 BN 也能训起来但学习率要比 RGB 分支低一点避免一开始就把预训练特征冲掉。backbone 的输出一般是四个 stage 的特征stride 分别是 4、8、16、32。融合模块可以只接 layer4 的输出也可以在每个 stage 都融合一次。前者显存省后者对小物体更友好。我一般先在 layer4 做一次融合跑基线如果小类 IoU 明显偏低再把 layer3 的输出也接上融合。3.2 融合位置选择早期、中期、后期融合的工程取舍融合方式做法优点缺点早期融合RGB 三通道加深度一通道直接输入同一个 backbone实现最简单参数量最小深度和颜色分布差异大浅层卷积容易学乱中期融合两个分支各自提特征在 stage3 或 stage4 做融合兼顾模态差异和计算量最常用需要对齐特征图尺寸和通道数后期融合两个分支各自出 logits再相加或取平均改动最小适合快速验证融合发生得太晚跨模态互补信息利用不充分室内场景分割里中期融合是比较稳的选择。深度图在边界、平面区域和遮挡关系上有优势RGB 在纹理和类别外观上有优势这两类信息在 stride 16 的特征图上已经比较抽象融合起来比原始像素层更有效。后期融合可以作为消融实验的基线但不建议作为最终方案。融合的具体操作也分几种通道拼接后接 1x1 卷积、逐元素相加、门控加权、注意力加权。拼接最直接但通道数翻倍显存涨得快相加最省但要求两个分支通道数一致且默认两个模态同等重要门控和注意力能学出逐像素的权重代价是多一点参数。下面这个门控融合块是我常用的写法RGB 特征做主路深度特征做门控训练稳定也方便插入到不同 stage。3.3 一个可插拔的注意力融合块实现import torch import torch.nn as nn class ConvBNReLU(nn.Module): def __init__(self, in_ch, out_ch, k3, s1): super().__init__() self.block nn.Sequential( nn.Conv2d(in_ch, out_ch, k, s, k // 2, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.block(x) class GatedFusion(nn.Module): 中期融合RGB 特征做主路Depth 特征经过门控后残差相加 def __init__(self, channels): super().__init__() self.rgb_proj ConvBNReLU(channels, channels, 1) self.depth_proj ConvBNReLU(channels, channels, 1) self.gate nn.Sequential( nn.Conv2d(channels * 2, channels, 1, biasFalse), nn.BatchNorm2d(channels), nn.Sigmoid(), ) self.out ConvBNReLU(channels, channels, 3) def forward(self, rgb_feat, depth_feat): # 先对齐通道保证两个特征图可以直接相加 rgb_feat self.rgb_proj(rgb_feat) depth_feat self.depth_proj(depth_feat) # 门控信号看的是两个模态的联合信息输出 0~1 的逐通道权重 gate self.gate(torch.cat([rgb_feat, depth_feat], dim1)) # 残差式融合深度特征只补充 RGB 缺失的部分 fused rgb_feat gate * depth_feat return self.out(fused)GatedFusion的输入是两个[B, C, H, W]的特征图C必须一致H和W也要一致。如果两个分支的 stride 不同先对深度特征做一次F.interpolate或调整 backbone 的膨胀率不要直接相加。gate的输出是逐通道的不是逐像素的标量这样每个通道可以独立决定深度信息保留多少比全局一个权重灵活。如果手头有invalidmask可以在depth_proj之后乘上去把无效深度区域的特征清零再进融合。这个操作在深度图空洞多的数据集上很有效能减少网络把空洞当平面学的概率。融合块放在 layer4 之后通道数一般是 2048ResNet-50显存占用不小如果放在 layer3 之后通道数是 1024可以省一些显存但感受野会小一点。4. 训练循环与损失函数类别不平衡下把 mIoU 稳住分割任务里背景像素通常占大多数小类样本少交叉熵很容易把模型带成“全预测背景”。RGB-D 项目还多一层麻烦深度图无效区域如果被当成正常像素参与损失模型会学到一片低置信度的平坦区域mIoU 看起来不低小类 IoU 却接近零。损失函数和指标计算要一起改不能只盯着总 loss。4.1 损失组合CE、Dice 与类别权重我一般用交叉熵加 Dice 的组合。CE 负责逐像素分类Dice 负责拉高整体重叠度两个权重各 1.0 起步。类别权重不要一上来就按频率倒数设得很猛先不加权跑几个 epoch看混淆矩阵里哪些类被背景吞了再针对那几类加权重。权重太大容易让小类过拟合验证集上小类 IoU 波动会很大。import torch import torch.nn as nn import torch.nn.functional as F class SoftDiceLoss(nn.Module): def __init__(self, ignore_index255, eps1e-6): super().__init__() self.ignore_index ignore_index self.eps eps def forward(self, logits, target): num_classes logits.shape[1] prob F.softmax(logits, dim1) # 排除 ignore_index避免边界像素参与 Dice valid (target ! self.ignore_index).unsqueeze(1).float() target_onehot F.one_hot(target.clamp(min0), num_classes) target_onehot target_onehot.permute(0, 3, 1, 2).float() prob prob * valid target_onehot target_onehot * valid dims (0, 2, 3) inter (prob * target_onehot).sum(dims) union prob.sum(dims) target_onehot.sum(dims) dice (2 * inter self.eps) / (union self.eps) return 1 - dice.mean() class CombinedLoss(nn.Module): def __init__(self, class_weightsNone, ignore_index255, ce_weight1.0, dice_weight1.0): super().__init__() self.ce nn.CrossEntropyLoss(weightclass_weights, ignore_indexignore_index) self.dice SoftDiceLoss(ignore_indexignore_index) self.ce_weight ce_weight self.dice_weight dice_weight def forward(self, logits, target): ce_loss self.ce(logits, target) dice_loss self.dice(logits, target) return self.ce_weight * ce_loss self.dice_weight * dice_lossclass_weights传的是一个长度为num_classes的 tensor背景类权重可以设 1.0小类设 2.0 到 5.0不要超过 10.0。ignore_index在 CE 和 Dice 里都要排除CE 直接支持这个参数Dice 需要手动乘validmask。target.clamp(min0)是为了防止one_hot遇到 255 报错真正的排除靠后面的valid。4.2 训练脚本参数epoch、batch size、学习率策略参数常见取值说明optimizerAdamW分割任务常用权重衰减和 Adam 解耦lr1e-4 到 6e-5backbone 可以设更低融合模块用主学习率weight_decay1e-4太大容易欠拟合batch_size4 到 8取决于 crop 尺寸和显存crop_size480x640 或 512x512室内场景常用 480x640epochs80 到 200数据量小取大一点配合早停lr_schedulerpolypower0.9分割里比 cosine 更常见amp开启显存省 30% 左右注意 loss scalingpoly 策略的实现很简单每个 iteration 按lr * (1 - iter / max_iter) ** 0.9更新。混合精度训练时Dice loss 里有sum和除法建议把 loss 计算放在autocast外面或者用GradScaler包住反向传播避免梯度下溢。batch size 太小的时候 BN 统计不稳定可以把 backbone 的 BN 换成 GroupNorm或者冻结 BN 的 running stats。4.3 验证与指标mIoU 的计算别写错mIoU 不是把所有 batch 的 IoU 求平均而是先累积混淆矩阵再按类算 IoU最后对类别求平均。ignore_index必须在累积之前排除否则边界像素会把某些类的分母撑大。验证集不要做随机裁剪和颜色抖动统一 resize 到固定尺寸或者直接跑原图然后 center crop。import torch class ConfusionMatrix: def __init__(self, num_classes, ignore_index255): self.num_classes num_classes self.ignore_index ignore_index self.mat torch.zeros((num_classes, num_classes), dtypetorch.int64) torch.no_grad() def update(self, pred, target): # pred 是 argmax 后的类别索引不是 logits valid target ! self.ignore_index pred pred[valid] target target[valid] idx target * self.num_classes pred self.mat torch.bincount( idx, minlengthself.num_classes ** 2 ).reshape(self.num_classes, self.num_classes) def compute(self): hist self.mat.float() iou hist.diag() / (hist.sum(0) hist.sum(1) - hist.diag() 1e-9) return iou.mean().item(), iouupdate里先做 mask 再bincount顺序不能反。bincount的输入要转成int64在 CPU 上算比在 GPU 上快累积完一个 epoch 再搬到 GPU 算 IoU 也行。如果某些类在验证集里根本没出现hist.sum(0) hist.sum(1) - hist.diag()会是 0分母加1e-9防止 NaN但这类别的 IoU 不应该参与平均最好在compute里加一个valid_classes过滤。5. 避坑与排查RGB-D 分割项目里最容易翻车的五件事RGB-D 分割的坑很少集中在网络结构上大部分出在数据读入、增强同步和指标计算。下面这五条是我踩过或者帮人排查过的按“现象 → 原因 → 解决”写方便对着日志定位。5.1 数据读取与预处理阶段的坑现象训练 loss 前几个 epoch 正常后面突然变成 NaN或者 mIoU 卡在背景类不动。原因多半是深度图里的无效值没有处理。深度相机在玻璃、黑色物体和超量程区域返回 0直接除以depth_max后还是 0本身不会 NaN但如果后面做了标准化或者取对数0 就会变成负无穷。解决方法是读入时先算invalid (depth 0) | (depth depth_max)把无效值置 0再做截断归一化如果用了深度图标准化均值和方差只在有效像素上统计。现象训练时 loss 正常下降验证集上预测的 mask 和原图边缘错位尤其是物体边界附近。原因通常是 RGB 和 Depth 的几何增强没有同步。随机裁剪用了两组随机数或者 resize 时 RGB 用双线性、Depth 用最近邻缩放后的像素中心对不上。解决方法是把 RGB、Depth、Label 放进同一个joint_transform裁剪参数只生成一次resize 时三张图用同一组尺寸和插值方式。Depth 建议用最近邻或双线性但必须和 RGB 保持一致。5.2 模型加载与训练阶段的坑现象加载预训练权重后训练初期 loss 比不加预训练还高或者直接报错说 size mismatch。原因通常是深度分支第一层被改成了单通道但加载时用了strictTrue或者把 RGB 三通道权重硬塞进单通道卷积。解决方法是加载前打印missing_keys和unexpected_keys对第一层卷积单独处理要么随机初始化要么把 RGB 三个通道的权重求平均后赋给深度分支。不要为了省事直接strictFalse然后不检查分类头的权重没加载上也会静默失败。现象总 mIoU 看着还行但小类 IoU 一直是 0混淆矩阵里小类全被预测成背景。原因可能是类别权重设得不够也可能是 Dice 的权重太低被 CE 的梯度淹没了。先不加类别权重跑一轮看混淆矩阵里小类的召回率如果召回率极低再把小类权重从 2.0 开始往上加同时把 Dice 权重提到 1.5 左右。不要一上来就把小类权重设成 20.0模型会为了几个像素过拟合验证集指标反而更差。现象验证集 mIoU 每个 epoch 波动很大同样的模型跑两次结果差好几个点。原因通常是验证集用了随机增强或者DataLoader的shuffleTrue没关。验证集只做 resize、中心裁剪和归一化不要加随机翻转和颜色抖动shuffle设成Falsedrop_last也设成False否则最后一个 batch 的样本会被丢掉指标统计不完整。5.3 评估与复现阶段的坑现象自己算的 mIoU 和论文或 baseline 对不上差两三个点。原因可能是ignore_index没有排除或者混淆矩阵在 batch 之间被重置了。检查update里是不是先 mask 再bincountmat是不是在整个验证集上累积而不是每个 batch 清零。还有一种情况是预测时用了多尺度 TTA但验证脚本里没加对比时要把 TTA 关掉再比。6. 二次开发换数据集、换 backbone 与导出推理的改造路线换数据集时先改三处datasets/里的类别映射表、配置里的num_classes和ignore_index、深度图单位。类别映射不要直接按标签 ID 顺序假设先打印训练集里出现的唯一值确认背景是 0 还是 255未标注像素是 255 还是别的哨兵值。深度图如果从毫米变成米depth_max和除法那行要一起改否则归一化后的数值范围会差一千倍。换 backbone 时重点看输出 stride 和通道数。把新 backbone 的四个 stage 输出打印出来确认和融合模块的输入对齐如果 stride 太大小物体分割会明显变差可以加膨胀卷积或者把融合层提前到 stage3。导出 ONNX 时模型要有两个输入image和depth动态轴至少把 batch 和 H/W 设成动态否则换分辨率推理会报错。导出后拿一张图对比 PyTorch 和 ONNX Runtime 的输出逐像素算最大绝对误差超过 1e-3 就要查是不是有算子不支持或者归一化写在了模型外面。血泪经验是每次改完数据管线先跑一个 overfit 小实验拿 8 到 16 张图训练 200 个 iteration看训练 mIoU 能不能到 0.95 以上。如果连过拟合都做不到不要急着调融合模块先回去查标签对齐和深度无效值。从那以后我每次接新的 RGB-D 数据都先把深度图的无效值比例、RGB 和 Depth 的尺寸一致性、标签唯一值这三项打成日志再谈模型。希望帮到你。本文还有配套的精品资源点击获取