简介这是一份Python实现的DenseFusion六自由度物体姿态估计项目包面向计算机视觉研究者和机器人抓取、增强现实等应用开发人员。核心方法通过融合彩色图像与深度图像对物体位置和三轴旋转进行像素级稠密特征估算从而在堆叠或遮挡场景下得到稳健的位姿输出。压缩包采用ZIP格式共收录55个文件其中包含二十个Python源文件用于模型定义、训练与推理另有Shell部署脚本、CUDA与C扩展、MATLAB精度评估脚本、配置文件以及说明文档整体大小为三点五一兆字节。目前已聚集一千八百七十一位学习者代码结构按主目录划分清楚便于逐模块阅读。资源内提供完整的工程化流程包括网络结构、损失函数、数据加载、预训练模型、训练评估脚本以及关键帧精度对比图像覆盖从环境配置、数据预处理到六自由度结果优化的各个环节适合希望复现或改进DenseFusion算法的工程师。1. 6D 位姿估计为何难一个用逐像素融合把 RGB-D 变成空间位姿的 Python 项目机械臂要抓一个杯子只知道它在图像里的 2D 位置是不够的还必须知道它朝哪个方向、离相机多远——这正是 6D 位姿估计要解决的问题。传统做法先找关键点再解 PnP一到弱纹理、强遮挡的物体就翻车关键点检测的失败率高得吓人。DenseFusion 换了一种思路把物体区域的逐像素颜色特征和几何特征做密集融合每个像素独立投票一个候选位姿再按置信度选出最可信的结果。这个 Python 仓库完整实现了论文里的结构和训练流程覆盖 YCB-Video 数据集的下载、预处理、训练、评估与推理适合做机器人抓取、AR 叠加、工业分拣的算法工程师和研究生拿来当基线也适合想在网络里试融合方式和注意力机制的开发者用于二次开发。2. 技术选型与整体架构逐像素融合为什么能打赢全局特征2.1 从 PoseCNN 到 DenseFusion全局特征的问题在哪早期做 6D 位姿估计的深度网络主流方案是把整张图喂给 CNN拿一个全局特征向量再接全连接层直接回归旋转和平移。代表性工作 PoseCNN 就是这么干的先用全局特征做语义分割再分别预测物体的三维平移和旋转。这种范式的核心问题在于全局特征是整张图像的平均它天然丢失了空间位置信息。物体边缘的轮廓、物体表面的细节纹理、被遮挡区域旁边的干扰像素全部被压进同一个特征里网络很难决定“到底该信图像的哪一块”。DenseFusion 的改法非常直接既然全局平均不好那就做 per-pixel 级别的密集融合。它对物体区域内的每个像素都构造一条颜色特征和一条几何特征用两层 MLP 融合成该像素自己的特征紧接着让每个像素独立预测一个候选位姿和一个置信度最后把所有像素的预测结果投票出一个最终位姿。这样做等于让网络先学会判断每一个点对姿态估计的贡献边缘点提供轮廓约束平面点提供几何约束被遮挡的脏点置信度自然会被压低。后续很多工作把这种隐式置信度升级成显式注意力模块但基础框架仍然是 DenseFusion 这一套。这个设计还有一个容易被忽略的优点几何分支不是直接吃深度图而是吃由深度图反投影出来的三维点云。点云的每个点都自带空间坐标对遮挡和弱纹理天然比深度图鲁棒。深度图本质上是一个强度图需要 CNN 自己隐式学习反投影公式而点云把这一步显式做掉了网络省下的容量可以用来学更复杂的东西。这也是我为什么会把这个项目当成位姿估计的默认基线——它在设计上就把几何和颜色摆在平等位置而不是让网络从零去猜它们的关系。2.2 目录结构源码包从数据集到训练脚本的布局整个项目按功能分块看一眼目录就大概能猜到作者的设计意图DenseFusion/ ├── datasets/ │ ├── ycb/ # YCB-Video 数据集的加载与读取 │ ├── linemod/ # LINEMOD 数据集的加载与读取 │ └── dataset.py # 通用数据读入、mask 处理、点云采样 ├── lib/ │ ├── network.py # 初始位姿估计网络双分支 融合 投票 │ ├── refine.py # 迭代细化残差网络 │ ├── loss.py # ADD / ADD-S 损失函数 │ ├── transform.py # 位姿与点云变换工具 │ └── pspnet.py # 颜色分支PSPNet 特征提取 ├── tools/ │ ├── download_ycb.sh # 下载与解压数据集的脚本 │ ├── evaluate.py # YCB-Video 协议下的评估 │ └── visualize.py # 可视化预测位姿与 GT ├── experiments/ # 实验配置与日志 ├── train.py # 训练入口 ├── test.py # 测试入口 ├── demo.py # 单张图片推理演示 └── README.md我拿到一个视觉项目习惯先读三个文件lib/network.py看网络结构lib/loss.py看训练信号怎么设计train.py看数据怎么灌进去。这三个文件读懂了整个项目的骨架就清楚了。剩下的数据加载和评估流程大多是体力活哪里报错改哪里就好。值得留意的是tools/下面放了下载脚本和可视化脚本说明作者是考虑了完整复现链路的。很多开源项目只给训练代码数据准备全凭猜这个项目把下载、可视化、评估都补齐了二次开发时省去不少从零造轮子的时间。2.3 依赖环境PyTorch 版本、GPU 与 visdom项目基于 Python 3 和 PyTorch 实现PyTorch 0.4 到 1.x 的版本我基本上都试过代码层面的兼容性没什么大问题。除了 PyTorch 之外还需要 opencv-python、numpy、scipy以及训练可视化用的 visdom。GPU 显存是硬性要求8GB 起步想在 4GB 卡上跑完整训练基本不现实但用它做推理、跑通 demo 是足够的。这里有一个新手最容易踩的流程坑train.py启动后会默认连接 visdom 服务如果你没有先启动 visdom程序会一直停在“等待连接”的状态看起来就像卡死了。正确的顺序是先在另一个终端执行python -m visdom.server -port 8097确认日志输出Its Alive!之后再启动训练脚本。如果服务器上有多个人在用最好把端口换掉比如-port 8098否则可能连到别人的可视化面板上看到一堆奇怪的 loss 曲线。3. 核心模块拆解一条数据从进入网络到输出位姿的完整链路3.1 预处理把深度图变成逐像素点云YCB-Video 的每一帧数据包含 RGB 图、深度图和 mask 标注。第一步是根据相机内参把深度图反投影成逐像素点云再按 mask 取出物体区域在物体区域上固定采样 1024 个点作为几何分支的输入。这一步最常见的坑是深度单位YCB-Video 原始深度图以米为单位如果代码里多做了一次缩放点云会整体膨胀或收缩位姿估计直接废掉。def depth_to_cloud(depth, fx, fy, cx, cy, mask, num_points1024): 把深度图转成逐像素点云并按 mask 采样固定数量的点。 h, w depth.shape rows, cols np.meshgrid(np.arange(h), np.arange(w), indexingij) # 反投影公式x (u - cx) * z / fx x (cols - cx) * depth / fx y (rows - cy) * depth / fy z depth.copy() cloud np.stack((x, y, z), axis2) # (H, W, 3) cloud cloud[mask] # 只保留物体区域 # 点数不足时重复采样超过时随机采样 if cloud.shape[0] num_points: idx np.random.choice(cloud.shape[0], num_points, replaceTrue) else: idx np.random.choice(cloud.shape[0], num_points, replaceFalse) return cloud[idx].T # (3, num_points)参数说明fx、fy、cx、cy 是相机内参YCB-Video 的合成数据使用固定内参真实场景必须换成自己标定的值num_points 是采样点数默认 1024降到 512 可以显著减少显存占用在 YCB-Video 上精度损失通常不到一个百分点。注意这里只输入 XYZ 坐标没有把 RGB 颜色拼进点云——颜色特征已经由 RGB 分支单独提取几何分支的任务就是纯粹的几何形状建模。3.2 双分支特征提取PSPNet 与 PointNet颜色分支用的是 PSPNet 的特征层引入金字塔池化让特征带上前期多尺度上下文几何分支则是一个用 1×1 卷积实现的 PointNet 式 MLP对每个像素独立提取特征。两个分支的输出在通道维上拼接进入融合模块。class PoseNet(nn.Module): def __init__(self, num_points, num_obj, num_channel3): super().__init__() self.num_points num_points # 几何分支逐像素 PointNet输入 xyz 坐标输出 1024 维特征 self.conv1 nn.Conv1d(num_channel, 64, 1) self.conv2 nn.Conv1d(64, 128, 1) self.conv3 nn.Conv1d(128, 1024, 1) # 颜色分支加载 PSPNet输出 (B, 512, H, W) 特征图 self.color_fea build_pspnet() # 融合 MLP512 维颜色特征 1024 维几何特征 - 256 维 self.fusion nn.Sequential( nn.Conv1d(512 1024, 256, 1), nn.BatchNorm1d(256), nn.ReLU(), nn.Conv1d(256, 256, 1), nn.BatchNorm1d(256), nn.ReLU(), ) # 逐像素位姿预测旋转四元数 三维平移 self.pred nn.Sequential( nn.Conv1d(256, 128, 1), nn.BatchNorm1d(128), nn.ReLU(), nn.Conv1d(128, 7, 1), ) # 逐像素置信度 self.conf nn.Sequential( nn.Conv1d(256, 128, 1), nn.BatchNorm1d(128), nn.ReLU(), nn.Conv1d(128, 1, 1), )这里有几个实现细节值得展开。位姿网络的输出是 7 维——前 4 维是单位四元数表示旋转后 3 维是平移不是直接回归 6 维欧拉角。四元数避免了万向锁问题但训练时必须约束其模长为 1否则旋转预测会漂移。forward 里会先用 mask 把 PSPNet 输出的 H×W 特征图转成 N 个采样点对应的特征向量再和几何特征拼接因此 batch 内每个物体的点数是一致的方便做矩阵运算。3.3 密集融合与置信度投票融合之后每个像素都拿到一个 256 维特征经过 pred 分支独立输出一个位姿经过 conf 分支输出一个置信度。合并这些候选位姿时原论文的做法是直接取置信度最高的点对应的预测。实际工程里我一般会改成 top-k 加权平均把置信度前 8 的像素对应位姿按置信度加权融合这样对遮挡边界更稳。def merge_pose(poses, conf, k8): 按置信度取 top-k 像素做加权平均返回最终四元数 平移。 conf conf.squeeze(1) # (B, N, 1) - (B, N) topk_conf, topk_idx torch.topk(conf, k, dim-1) topk_poses torch.gather(poses, dim2, indextopk_idx.unsqueeze(1).expand(-1, 7, -1)) topk_conf topk_conf.unsqueeze(1) # (B, 1, k) # 四元数部分先单位化再做加权平均 q topk_poses[:, :4, :] / torch.norm(topk_poses[:, :4, :], dim1, keepdimTrue) t topk_poses[:, 4:, :] q (q * topk_conf).sum(-1) q q / torch.norm(q, dim-1, keepdimTrue) t (t * topk_conf).sum(-1) return torch.cat([q, t], dim-1) # (B, 7)参数说明k 值在 YCB-Video 上取 8 和取 16 差别不大但对对称物体建议配合 ADD-S 损失一起调四元数平均不能对欧拉角平均也不能直接对四元数矩阵做普通平均必须先把每个候选四元数归一化到单位长度。这段代码里每个候选四元数在加权前都做了归一化虽然会影响后续加权权重但实际影响很小。3.4 迭代细化用残差网络修正上一轮的姿态DenseFusion 在初始位姿估计之后还带一个 refine 网络。它的输入是上一轮预测的位姿、模型点云和融合特征输出一个位姿增量然后把增量叠加到当前位姿上迭代两次。每次迭代都会让物体模型点云逐渐向观测点云对齐有点像 ICP 迭代最近点但用的是学习出来的残差而不是几何距离匹配。class RefineNet(nn.Module): def __init__(self, num_points): super().__init__() self.mlp nn.Sequential( nn.Conv1d(256 128, 128, 1), nn.BatchNorm1d(128), nn.ReLU(), nn.Conv1d(128, 128, 1), nn.BatchNorm1d(128), nn.ReLU(), nn.Conv1d(128, 7, 1), # 输出残差位姿 ) def forward(self, fused_feat, model_cloud, cur_pose): # 用当前位姿把模型点云变换到相机坐标系再提取几何特征 transformed transform_points(model_cloud, cur_pose) # (B, 3, N) geo pointnet(transformed) # (B, 128, N) delta self.mlp(torch.cat([fused_feat, geo], dim1)) return compose_delta(cur_pose, delta) # 增量作用到位姿上常见理解误区refine 网络预测的不是绝对位姿而是残差。它把当前位姿下的模型点云再送进 PointNet 提取一次几何特征和原有融合特征拼接后输出一个微小的修正量。训练时也要注意阶段——先用第一阶段 loss 训好 PoseNet再固定 PoseNet 单独训 RefineNet最后联合微调。如果一开始就联合训练梯度信号会在两个网络之间互相干扰两个都训不好。3.5 损失函数对称物体的 ADD 与 ADD-S位姿估计的损失函数决定收敛方向。论文用的是模型点经预测位姿和 GT 位姿变换后的对应点距离即 ADD。对对称物体还要引入 ADD-S遍历所有点取最近距离允许网络在绕对称轴旋转的等价位姿之间自由选择。def add_loss(model_pts, pred_pose, gt_pose, symmetricFalse): 模型点经 pred/gt 位姿变换后的平均距离。 pred_pts transform(model_pts, pred_pose) # (B, N, 3) gt_pts transform(model_pts, gt_pose) if symmetric: dist torch.cdist(pred_pts, gt_pts) # 每对点距离 return dist.min(dim-1)[0].mean() # 每个预测点找最近 GT else: return (pred_pts - gt_pts).norm(dim-1).mean()细节说明ADD-S 的 min 发生在第二维是“每个预测点找它对应的最近 GT 点”不是所有点全局取最小。不少人自己实现时会把维度搞错导致对称物体的损失曲线异常低但位姿看起来完全不对。对称标志可以从数据集的模型点云里预先计算也可以直接查标注文件碗、圆柱这类物体一定要开 symmetricTrue否则网络会在多个等价位姿之间来回摇摆训练后期 loss 震荡得很厉害。4. 训练与评估让网络在 YCB-Video 上从零跑到可用精度4.1 数据集准备下载、目录组织与 mask 依赖YCB-Video 是一个大规模 6D 位姿数据集包含多段真实视频序列和对应的合成渲染数据。项目在tools/下提供了下载脚本但下载只是第一步数据下载后需要按特定目录组织datasets/ycb/data/下每个帧目录里放*-color.pngRGB、*-depth.png深度、*-label.pngmask 标注还要准备物体模型点云、相机内参文件和训练帧索引。不少人在这一步卡住——官方数据集的命名规则在不同版本里略有差异需要对着 dataset.py 里的读取逻辑写一个小脚本把文件名对齐。mask 是整个预处理的关键依赖。DenseFusion 训练时直接用 GT mask 从全图特征里抠出物体区域所以训练阶段的精度上限和 mask 质量强相关。如果是做自有数据集mask 可以用分割模型自动生成再人工修正但一定要保证 mask 能完整包住物体否则采样点会混入背景点云位姿估计直接偏移。4.2 训练命令与参数调优环境准备好、数据就位后训练指令很简洁python train.py --dataset ycb --num_points 1024 --batch_size 8 --lr 0.0001 --epochs 20 --gpu 0 --visdom参数说明num_points 是物体区域采样点数影响显存和几何分支精度batch_size 在单卡上建议 8 以下显存不足先降到 4不要贸然缩小 num_points精度损失更明显lr 在 YCB-Video 上 1e-4 比较稳超过 3e-4 经常出现 loss 震荡。训练脚本会自动按三阶段流程走先只训 PoseNet再固定 PoseNet 训 RefineNet最后联合训练但阶段切换的 epoch 需要在配置文件里手动指定。一个小建议是先用 2 个 epoch 做冒烟测试确认 loss 在合理下降、模型能正常保存和恢复再放开 20 个 epoch 的全量训练。很多人上来就跑 full training结果第 3 天发现 visdom 没连上、数据路径配错白白浪费了算力。冒烟测试这个习惯在视觉项目里特别有用训练管线越复杂越要先跑通再跑长。4.3 评估指标ADD 与 ADD-S 的 AUC 曲线评估时模型对每帧预测一个位姿计算预测点与 GT 点的平均距离距离小于阈值就算正确。YCB-Video 的官方协议里通常报告 ADD-S 的 AUC 值——对不同阈值下的准确率做积分得到一个 0 到 1 之间的综合分数比单一阈值更稳定。测试命令python test.py --dataset ycb --checkpoint ./checkpoints/ycb_best.pth --eval这里的 eval 模式会输出每个物体的 ADD-S AUC。要看结果是否可信需要区分两种评估方式以 GT 为中心裁剪物体区域再评估和全图检测后评估。前者只评价位姿精度后者包含了目标检测的误差分数会低不少。如果你的场景是已知物体位置、只需要算姿态那应该重点关注前者。5. 常见问题与避坑显存、收敛、数据单位与评估指标5.1 一跑训练就爆显存现象train.py 刚跑几十个 step 就直接 OOM进程被杀掉日志最后一行停在某个 DataLoader worker 上。原因输入图是 640×480PSPNet 本身显存占用不小再加上 1024 点的几何分支和融合模块峰值显存很容易超过 12GB。有些情况下 DataLoader 开启了 pin_memory还会额外占一份显存拷贝。解决先把 batch_size 从 8 降到 2再把 num_points 从 1024 降到 512最后清一次缓存torch.cuda.empty_cache()再启动。我一般不会一上来就缩小输入分辨率因为分辨率变了相机内参也要同步改位姿结果容易整体偏移。5.2 loss 稳在半山腰不下降预测位姿完全随机现象loss 曲线很快进入平台期数值不再明显变化可视化输出的模型点云和输入点云完全对不上。原因最常见的是深度图单位没转成米点云坐标整体缩放了几百倍网络根本无法收敛。其次可能是 PSPNet 没有加载 ImageNet 预训练权重随机初始化下颜色分支学不动。还有一个容易漏的细节四元数分支没有做单位化约束梯度方向是歪的。解决先打印点云 z 值的分布如果量级是几百到几千说明深度图是毫米为单位而代码按米处理需要统一缩放。再检查颜色分支是否加载预训练权重没有就补上下载。最后在损失函数里确认对四元数做了归一化。这三个检查完loss 平台期的问题基本能排除。5.3 对称物体评估分数异常低现象训练没崩但碗、圆柱这类对称物体的 AUC 比非对称物体差一大截甚至趋近于随机。原因对称物体的 GT 位姿不唯一绕对称轴旋转 90 度同样是正确的。如果用 ADD 损失去约束网络会在多个等价位姿之间摇摆学习信号冲突收敛自然差。解决把对称物体的 symmetric 标志置为 True损失切换到 ADD-S评估指标也同步换成 ADD-S。要注意 ADD 和 ADD-S 不能混用——训练用 ADD-S、评估用 ADD或者反过来都会得到误导性的结果。5.4 多线程数据加载导致训练中断现象DataLoader 的 num_workers 设成 4 或 8 之后训练中途报死锁、段错误或数据形状不一致。原因老版本 PyTorch 在 Windows 上对 DataLoader 多进程支持一般文件读取在子进程下容易出现句柄问题部分预处理用了 numpy 随机数多进程并发时也会偶发越界。解决先把 num_workers 设为 0跑通完整流程再加并行。如果一定要并行把打开文件的逻辑放在getitem里而不是init里每次读取时独立打开文件句柄避免父子进程共享文件指针。5.5 visdom 端口被占用训练卡在等待连接现象train.py 启动后一直打印等待 visdom 连接没有任何训练输出看起来像死循环。原因visdom 服务没有启动或者端口被另一个实验实例占用导致当前进程连不上可视化后端。解决启动训练前先执行python -m visdom.server -port 8097确认服务就绪再跑训练。多人共用服务器时建议给每个实验分配独立端口避免相互干扰。如果不依赖可视化也可以直接把 train.py 里的 visdom 相关调用注释掉训练照常进行。6. 二次开发思路把 DenseFusion 改成你自己的位姿估计基线如果要做自己的物体先别急着改网络结构把数据链路打通更重要。首先把自己的物体模型导出成点云格式点数最好用最远点采样统一到 1024 左右然后制作每个物体的 mask 标注可以通过分割模型自动生成再人工修正最后把相机内参写进配置。这三步做完训练脚本基本不用动直接换数据集路径就能跑通。想提升模型表现可以从特征融合下手。原版用的是 Concat MLP一个简单的升级是把两个分支的特征各自先过一个 SE 模块再把加权后的特征拼接这样每个像素的融合权重是自适应学出来的在弱光、遮挡多的场景里能明显减少位姿抖动。也可以调整 refine 的迭代次数从 2 次改成 1 次推理速度大约提升 1.8 倍精度下降有限适合部署场景。改动目标改哪个文件改法换自己的物体datasets/dataset.py修改模型加载与 mask 读取逻辑更新内参和对称标志加快推理lib/refine.py将 refine 迭代次数从 2 改成 1速度提升明显加注意力融合lib/network.py在 fusion 层前后各插一个 SE 或 CBAM 模块多尺度输入lib/pspnet.py把 PSPNet 输出改成两尺度特征拼接改完网络之后最怕的就是直接上全量训练。我自己在这上面吃过亏改完 loss 后直接跑 20 个 epoch结果权重爆炸白白浪费一周的卡时。后来给自己定了一条纪律——任何改动先跑一个单物体的小实验batch_size 调小、epoch 砍到 5确认 loss 能下降、可视化位姿能对齐再放开全量训练。从那以后我每次改模型都强制走一遍“单物体小实验到全量训练”的流程翻车的概率肉眼可见地下降了。这个项目模块边界清楚、结构简洁特别适合往里塞新想法希望帮到你。本文还有配套的精品资源点击获取