多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

Python实战DenseFusion:从RGB-D图像到6D物体姿态估计

Python实战DenseFusion:从RGB-D图像到6D物体姿态估计 简介Python-DenseFusion6D物体姿态估计是一套面向计算机视觉开发者和机器人研究者的完整项目实现方案用于解决RGB-D图像中物体6自由度位姿的精确估计问题可支撑机器人抓取、AR/VR交互和工业自动化等实际应用场景。资源包共55个文件以Python源码20个py为核心辅以Shell脚本、配置文件、模型文件、图片与文档等辅助内容压缩包仅3.51MB目录划分清晰涵盖网络定义、损失函数、特征提取与融合、姿态优化、数据集加载及评估等关键模块。已有1871人学习下载适合具备一定深度学习与计算机视觉基础、希望直接复现DenseFusion实验或基于此进行二次开发的开发者。通过该资源可系统了解YCB与LineMOD数据集上的训练和评估流程利用预训练模型与可视化对比图快速验证位姿估计效果减少从零搭建环境与调试代码的时间投入。1. 从“认得物体”到“抓得到”DenseFusion 6D姿态估计到底解决什么问题Python-DenseFusion6D物体姿态估计名字很长解决的问题却很具体给一张RGB-D图像和一堆要抓的物体让程序输出目标物体在相机坐标系下的六自由度位姿——三轴旋转加三轴平移合起来就是常说的6D位姿或6d位姿估计。DenseFusion是这个方向上绕不开的经典方法核心思想是把RGB的纹理特征和Depth的几何特征在像素级融合而不是像早期方案那样把两路特征一路拼到尾。它适合机械臂抓取、AR虚实叠加、无人仓分拣等场景这类任务光“认出物体”不够必须知道物体在哪儿、朝哪个方向机械臂才知道怎么下爪。用Python复现这套流程是姿态估计算法工程师最好的入门到进阶路径也是评估自己后期方案时必须对比的基线。2. 为什么逐像素融合比双流拼接更稳先看懂DenseFusion的三段式设计DenseFusion不是python基础语法能搞定的玩具项目它把分割、特征提取、姿态回归、损失设计串成了一条完整流水线。整体看分为三段先用语义分割在RGB图上把目标物体抠出来然后把抠出来的区域同时送进RGB特征网络和点云特征网络最后在融合模块里把每个像素位置的颜色特征和每个三维点的几何特征拼在一起逐点输出姿态候选和置信度。下面把这三段以及背后的选择理由拆开讲。2.1 先拆解RGB分支、几何分支与融合单元第一段的分割模块通常用现成的语义分割网络常见做法是DeepLab、PSPNet或者任意一个能在验证集上跑出高IoU的小网络。分割输出是逐像素的maskmask决定了后面所有计算的边界RGB分支只处理mask包住的区域点云分支也只取mask覆盖的深度像素所以mask质量直接影响位姿精度这点在后面避坑章还会重点说。第二段的RGB分支是一个小型的卷积特征网络输入是crop出来的物体区域输出是一张高维特征图每个像素对应一个颜色特征向量。Depth分支则用PointNet处理物体区域的点云把mask里的每个深度像素反投影成三维点PointNet对每个点提取几何特征。这里要特别注意几何分支用的不是整幅点云的全局特征而是逐点per-point特征只有这样才能跟RGB分支做像素级的对齐。第三段是DenseFusion的融合单元。RGB特征图上的像素和三维点在mask坐标上天然一一对应一个像素反投影过去就是一个三维点。于是可以用最简单的cat把两个特征向量拼起来送进一个MLP输出这个点对应的旋转四元数、平移向量和置信度。所有点都输出一组候选位姿最后按置信度做加权投票得到整个物体的位姿。为什么不直接把整幅RGB的全局特征和整幅点云的全局特征拼起来因为全局特征对遮挡很敏感物体被挡住一半时全局特征会被背景和其他物体污染逐点投票保留的是局部证据哪个区域可信就多信哪个。这也是DenseFusion和更早的PointFusion这类双流拼接方案最本质的区别。额外提醒一点逐点输出姿态候选的真实计算代价并不低如果一个mask内有1.2万个像素相当于同一张图要跑1.2万次MLP虽然实现上只是矩阵乘法但batch里所有物体的总点数会迅速涨到十几万。常见优化是先把mask区域降采样到长边不大于64的网格或者提前按物体mask随机抽500个像素再做融合后面训练章里还会专门讲点采样数这个参数为什么敏感。2.2 用Loss看懂算法意图ADD(-S)与置信度加权的PyTorch实现姿态估计的Loss设计决定了网络到底在优化什么。直接回归旋转矩阵的三个欧拉角是不推荐的欧拉角有周期性跳变网络得先学会处理角度回绕。常见做法是用四元数表示旋转并做归一化训练时额外处理四元数的符号二义性因为q和-q表示同一个旋转。平移则直接回归三维坐标但要把物体点云先做质心归一化让平移头的输出范围落在原点附近。下面这段PyTorch代码是ADD/ADD-S距离的最小实现也是DenseFusion训练时姿态损失的核心。它把所有模型表面点分别用预测位姿和真值位姿变换到相机坐标系再计算平均距离import torch def add_adds_loss(pred_rot, pred_trans, gt_rot, gt_trans, model_pts, symmetricFalse): # pred_rot/gt_rot: [B, 3, 3]训练时由四元数转矩阵得到 # pred_trans/gt_trans: [B, 3] # model_pts: [M, 3]物体表面均匀采样点M常见取1000~2000 # 把模型点云分别用预测位姿和真值位姿变换到相机系 pts_gt torch.einsum(bij,mj-bmi, gt_rot, model_pts) \ gt_trans[:, None, :] # [B, M, 3] pts_pred torch.einsum(bij,mj-bmi, pred_rot, model_pts) \ pred_trans[:, None, :] # [B, M, 3] if symmetric: # 对称物体圆柱、碗不存在严格对应点找最近点作为替代 dist torch.cdist(pts_pred, pts_gt) # [B, M, M] dist dist.min(dim-1).values # [B, M] else: # 非对称物体用对应点距离 dist torch.norm(pts_pred - pts_gt, dim-1) # [B, M] return dist.mean(dim1) # [B]逻辑说明einsum的写法把旋转矩阵同时作用到每个模型点上比手写for循环快得多对称物体必须走cdist分支因为对称物体的两个“不同姿态”在物理上不可区分如果强制按点对应关系计算距离网络会被惩罚到平均值上最后学出一个“各向都差一点”的位姿这就是很多复现Loss不降、姿态却不对的根因。M取多少合适我一般按物体表面积控制小物体1000点够用大物体取2000到5000点点太少ADD对位移和旋转都不敏感。置信度分支的训练思路是让网络学会“哪些点的预测靠谱”。常见实现是把姿态loss按置信度加权并额外加一项约束防止置信度退化成均匀分布conf pred_conf.softmax(dim1) # [B, N] 逐点置信度先归一化 weighted (conf * dist.detach().unsqueeze(1)).sum(dim1) loss_pose weighted.mean() # loss_conf 常见做法约束置信度分布不要变成均匀分布逻辑说明这里对dist做detach是常见技巧让梯度只更新融合网络和姿态头避免置信度头因为梯度过大提前锁死。加权用softmax会让所有点的权重之和恒为1物体点云点数多时单点权重被稀释所以不少人把softmax改成带温度的版本温度在0.5到1.0之间微调。置信度正则项的权重通常给到姿态损失的0.01到0.1倍具体调法在训练章里再说。2.3 表示与选型旋转用四元数、平移直接回归为什么这样组合最省心关于旋转的表示输出四元数再归一化是DenseFusion系列实现里最稳妥的选择。旋转矩阵九个输出但只有三个自由度要满足正交和行列式为1的约束网络很难直接学出来欧拉角的万向锁和回绕问题训练时经常翻车轴角虽然紧凑但在零角度附近有奇点。四元数四个数加上模长归一化约束最简单反算旋转矩阵也方便。平移部分我不建议输出相对于物体中心的绝对三维坐标后直接算L2而是先在数据预处理时把点云减去质心。这样平移头的目标值就是物体中心相对点云质心的偏移通常是一个零点几的小数配合smoothL1损失收敛更快。实际训练里旋转和平移的尺度天然不一致使用ADD距离作为损失的好处在于它把旋转误差通过物体半径换算成了“距离”旋转差3度和小物体上差3毫米对ADD的影响是同一个量纲不需要再手工给两个头配权重。这也是为什么评估指标和训练损失都用同一套ADD口径最省事。3. 用Python搭好环境与数据集YCB-Video到点云采样的落地步骤进入Python环节之前先说结论DenseFusion这类老方法代码能不能跑起来一半看环境一半看数据。数据准备好了训练只是等收敛数据没准备好后面所有环节都会被污染而且很难定位。建议不要跳步按顺序来。3.1 Python环境配置conda、PyTorch与点云库的组合拳这是整个项目里最卡人的一步。不少开源实现的KNN、模型点云采样等模块会用C扩展需要gcc配合CUDA编译。你如果不想在这上面耗时间我的建议是尽量用Python环境配置里最主流的组合Python 3.8、PyTorch 2.0附近版本、gcc版本和CUDA toolkit版本对齐。conda create -n densefusion python3.8 -y conda activate densefusion # 按当前机器的CUDA版本从PyTorch官网复制安装命令 pip install torch torchvision pip install numpy scipy open3d trimesh h5py tensorboard逻辑说明Python 3.8是兼容老代码和编译扩展的最佳档位3.10以上会在某些C扩展的编译处报错open3d和trimesh都用来处理点云与mesh二选一也够主要承担ICP和读取模型文件如果编译扩展一直失败直接绕开它KNN用open3d或scipy的cKDTree替代旋转向量到矩阵用scipy的Rodrigues公式自己写后面避坑章会给具体方案。这一步绕开后环境配置的很多玄学问题都不会再遇到。提示在跑第一个epoch之前先打印一次点云质心和真值平移确认数据通路上没有符号或单位错误。单位错误是最难定位的因为Loss照样会下降。3.2 YCB-Video与LineMOD目录结构、深度图单位与内参对齐姿态估计实验最常用的两个公开数据集YCB-Video量级大、场景杂包含21个日常物体和92个视频序列适合做法消融和最终指标LineMOD只有十几个工业零件纹理弱、更看重几何特征适合快速验证网络和调参。我一般先在LineMOD上挑一个物体跑通再切到YCB-Video跑全部省下的调试时间比什么都值钱。读数据时有两个细节决定后续能不能正常训练。第一个是深度图单位有的数据集版本深度存的是米有的存的是毫米读进来后先打印depth.max()和depth.min()判断如果是几百上千记得在反投影前统一除以1000换算成米。第二个是相机内参K训练数据里给的是整张图的内参但把物体crop出来之后图像坐标系的原点被平移了K也必须跟着平移否则每帧点云都会在x、y方向上错位几个像素对应的距离。import numpy as np K np.array(meta[intrinsic]).reshape(3, 3).astype(np.float32) ys, xs np.where(mask 0) x1, y1, x2, y2 xs.min(), ys.min(), xs.max(), ys.max() # 外扩8个像素避免卷积把边缘截掉 x1, y1 max(x1 - 8, 0), max(y1 - 8, 0) x2, y2 min(x2 8, W - 1), min(y2 8, H - 1) crop_rgb rgb[y1:y21, x1:x21] # python数组切片直接拿到ROI crop_depth depth[y1:y21, x1:x21] crop_mask mask[y1:y21, x1:x21] # 内参随crop原点平移这里是最多人漏掉的一步 K_crop K.copy() K_crop[0, 2] - x1 K_crop[1, 2] - y1逻辑说明外扩8个像素不是随手拍的RGB分支的骨干网络在最上层有下采样不给边缘留余量的话物体边缘的像素会直接参与padding特征被背景污染。K_crop的平移是把图像坐标系原点从全图左上角移到crop左上角不处理的话点云的x、y会整体偏移。mask的类型也要检查连续分割输出可能是0到255的灰度记得先做类型转换和二值化再执行np.where。3.3 从深度图到500个采样点反投影、均匀采样与质心归一化的代码骨架这一步是把像素转成几何输入。我建议用python定义函数把反投影、采样、归一化这三步封装成一个函数任何数据集的读取脚本都调用同一套后续排查数据问题会快很多。常见做法是把mask内的深度像素做小孔成像逆变换得到相机坐标系下的三维点。要注意反投影公式用的是内参而不是图像坐标的简单缩放def back_project(crop_depth, crop_mask, K_crop, n_sample500, seed0): ys, xs np.where(crop_mask 0) d crop_depth[ys, xs].astype(np.float32) # 先做python类型转换 d d / 1000.0 if d.max() 100 else d # 深度图单位统一到米 valid (d 0.3) (d 1.5) np.isfinite(d) xs, ys, d xs[valid], ys[valid], d[valid] # 小孔成像逆变换像素坐标(u, v) 相机三维坐标(x, y, z) x (xs - K_crop[0, 2]) * d / K_crop[0, 0] y (ys - K_crop[1, 2]) * d / K_crop[1, 1] z d pts np.stack([x, y, z], axis1) # 物体区域的点可能有上万固定采样n_sample个点 rng np.random.RandomState(seed) if len(pts) n_sample: pts pts[rng.choice(len(pts), n_sample, replaceFalse)] # 质心归一化让平移头的目标值变小收敛更快 centroid pts.mean(axis0) pts - centroid return pts.astype(np.float32), centroid逻辑说明d.max() 100是单位判断的经验写法深度单位是米时数值一般在0.5到3之间毫米制时是几百到几千valid条件里的0.3到1.5米不是固定的按相机工作距离调整目的是把明显无效的深度剪掉。n_sample用500是DenseFusion系列实现里最常见的选择几何分支能拿到足够多的点又不至于让逐点融合的计算量失控采样时固定seed保证实验可复现训练时可以不固定。pts减去质心后位姿真值里的平移也要对应减掉质心训练代码里必须同步处理否则Loss对不上。到这里输入输出的数据通路已经完整RGB的crop图、mask、采样点云、带质心的平移真值都能按batch送进训练循环了。4. 训练与评估的完整闭环从训练脚本到ADD(-S)指标怎么算数据通路打通后训练本身并不复杂。真正影响结果的是训练循环里的数据增强、学习率计划和评估口径。4.1 训练工程的最小骨架数据增强、batch与学习率的经验值DenseFusion这类一物一模型的做法训练时最有效的增强不是颜色抖动而是位姿扰动。常见做法是对点云和对应真值同时加随机扰动把物体绕相机系z轴随机转一个角度、在x/y方向随机平移一点模拟机械臂在不同角度看待抓物体的差异。RGB图做对应投影增强在实现上比较绕很多人只对点云做增强RGB保持原样效果也够用。训练骨架如下optimizer torch.optim.Adam(model.parameters(), lr1e-3) scheduler torch.optim.lr_scheduler.StepLR(optimizer, step_size20, gamma0.5) for epoch in range(200): for rgb, pts, rot_gt, trans_gt, centroid in train_loader: # batch层面做几何扰动点云绕z轴旋转真值位姿同步变化 ang np.random.uniform(-30, 30) / 180 * np.pi # 对pts施加旋转矩阵Rz并对gt_rot左乘Rztrans_gt做相应旋转 ... pred_rot, pred_trans, pred_conf model(rgb, pts) loss pose_loss(pred_rot, pred_trans, gt_rot, trans_gt) \ 0.05 * conf_loss(pred_conf) optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() if epoch % 20 0: torch.save(model.state_dict(), fckpt_{epoch}.pth)逻辑说明学习率从1e-3起步是Adam配小模型的常用档位StepLR每20个epoch乘0.5到100个epoch附近学习率已经很低适合做精细收敛。置信度损失的权重0.05不是拍脑袋可调范围一般在0.01到0.1之间设太大会让网络把点置信度全部抹平来降低loss。显存不够时先减batch size并用梯度累积而不是直接砍点数原因在避坑章细说。还有一个常被忽略的点batch里最好保持同一物体的不同视角样本不要混入多个物体否则点云特征分布被拉宽收敛变慢。4.2 评估不是看lossADD、ADD-S与AUC怎么用代码落地训练过程中最容易被假象骗到的点是看训练loss。因为姿态loss已经被置信度加权过网络完全可以把难点的置信度调低来压loss所以必须定期在验证集上算ADD。ADD的平均距离把旋转误差和平移误差统一到同一个厘米量级小于阈值就算这一帧成功。import numpy as np def add_auc(add_list, max_th0.1): add_list np.asarray(add_list) # 单位米 ths np.linspace(0, max_th, 100) # 0~10cm百分位扫描 acc [np.mean(add_list t) for t in ths] return np.trapezoid(acc, ths) / max_th # 面积占比即AUC逻辑说明AUC的扫描上限按数据集习惯走YCB-Video我习惯用0.1米LineMOD有人用物体直径的10%当阈值所以对比别人复现前先统一口径。np.trapezoid是numpy 2.0里trapz的新名字老版本换成np.trapz。除了AUC还要用python数据分析与可视化那一套画误差分布直方图如果大量帧误差集中在2到4厘米问题多半是旋转角度差而不是平移漂移如果分布是长尾多半是遮挡严重时的mask出错。逐物体把ADD平均值写成csv或excel留档每次改完参数都能对比上一版这个习惯能省下大量调参时间。4.3 三个影响收敛的参数置信度权重、点采样数、学习率先把最值得调的参数列出来按影响程度排序。第一是置信度权重它决定网络学不学得会“哪些点可靠”。权重太小时置信度退化成均匀分布遮拦严重的点照样参加投票太大时网络会把所有点都预测成低置信度来逃避损失。第二是采样点数它在显存、几何表达能力和融合计算量之间做权衡。第三是学习率与学习率计划除了StepLR也有人用CosineAnnealing在单物体姿态估计这种任务里两者差别不大选择一个后不要再频繁切换。采样点数的坑值得多说一句网络在训练时见过的点数分布会和Loss绑在一起。训练时固定500点推理时为了省时间改成100点ADD会掉得比想象多因为PointNet的特征分布变了。要提速就训练和推理一起改点数不要让两者脱节。5. 避坑Loss下降但ADD不涨五条血泪经验下面五条是最常遇到、也最容易被误解的问题。每条按现象、原因、解决三段写你可以直接对着自己的训练日志检查。排查顺序我一般是这样先查数据通路mask和内参再查Loss口径接着查增强最后才怀疑模型结构。因为模型结构的问题通常所有指标一起崩而数据或Loss口径的问题往往是“loss漂亮、指标不动”这种隐蔽状态。5.1 Loss下降但ADD不涨对称物体用错了损失口径现象训练曲线很漂亮loss从几十降到个位数但验证集ADD一直稳定在一个偏大的值怎么调学习率都没用。原因物体表面存在对称性杯子、圆柱、碗代码里却用了非对称的对应点距离计算ADD。网络学到的是一个“平均姿态”它不会把某个对称方向学得特别准而是所有方向都偏离一点因为这个平均姿态在所有点上的距离误差总和反而最小。解决先对物体模型做对称性分析常见做法是在训练脚本里给每个物体标一个symmetric标志把Loss从ADD切到ADD-S用第2章代码里的cdist分支。如果是圆对称物体评估也要用ADD-S否则训练和评估口径不一致指标永远上不去。5.2 自定义C扩展编译不过gcc、CUDA、PyTorch三者版本相爱相杀现象pip install都成功一跑训练就在from ext import ...报错不是缺头文件就是符号找不到重装PyTorch后反而又坏了。原因老实现里的C扩展跟PyTorch的C ABI强耦合gcc版本太新或太旧都会翻车CUDA_HOME指向的版本和PyTorch编译时的版本不一致也会在编译期崩。解决最简单的路径是绕开编译。Rodrigues公式用scipy.spatial.transform.Rotation实现KNN用open3d或scipy的cKDTree实现点云采样用numpy自己写全部纯Python/PyTorch化单物体场景的性能差距几乎感知不到。如果一定要编译就固定一套经过验证的组合Python 3.8、PyTorch 2.0附近版本、gcc 9并且不要混着pip乱升级。5.3 平移很准、旋转偏数据增强里少了z轴旋转和深度噪声现象验证集误差大多在2到5厘米画出误差分布发现主要是旋转角度差平移误差很小。原因很多复现默认只对点云做水平平移扰动或者增强时随机转的角度范围太小。机械臂实际抓取时会从各种角度观察物体训练集里旋转多样性不足网络自然学不好旋转。另一个原因是深度噪声深度边缘一圈的像素反投影出来误差很大这些点会直接把旋转估计带偏。解决把增强里的绕z轴旋转范围加到±30度以上必要时做轻微随机倾斜反投影前对深度图做一次中值滤波或按mask做孔洞填充把深度边缘一圈的异常值先抹掉。我一般会在训练循环里对每个batch随机选一个角度而不是固定角度列表。5.4 mask边缘毛刺把姿态拉偏分割误差是位姿误差的第一来源现象同一物体、同一网络换了语义分割模型后ADD明显变化且幅度超过预期。原因mask膨胀或收缩几个像素在点云里就是几毫米的几何变形如果mask把背景物体也包含进来点云里会混入完全不属于目标物体的点姿态回归头被这些离群点拉偏。解决crop时不要直接用原始mask先做一次形态学闭运算或腐蚀把细碎毛刺清掉点云采样后加一步统计离群点剔除例如按到质心的距离过滤超过3倍标准差的三维点推理阶段用网络输出的置信度再过滤低分点让最终投票只保留可靠区域。顺序很重要先清mask再剔点最后置信度过滤。5.5 显存不足砍batch_size后效果反而变差现象GPU显存8Gbatch_size只能设4明明模型没变收敛后ADD比用8的时候差一大截。原因逐点融合的计算量随总点数走砍batch_size等于把每个epoch看到的物体姿态数砍半加上batch内BN统计量不稳几何分支没有见过足够多的点云分布特征学偏了。解决保持batch_size不变先用梯度累积模拟大batch如果还是放不下再统一减少n_sample从500降到384或256但训练和推理必须一起改。很多人为了贪推理速度只减推理点数结果训练推理不一致ADD凭空掉两三个点这是典型的部署期翻车。另外用混合精度训练也能省下约一半显存老代码加autocast就能跑收益比砍点数高。6. 收尾三个从实验室走向产线的姿态后处理技巧6.1 给姿态装一个ICP收尾器网络直接输出的位姿在精抓取场景通常差那么一两个厘米或一两度。常见做法是把预测位姿当初值用模型点云对观测点云跑一次点到面ICP# 以open3d为例模型点云与观测点云的精细对齐 reg o3d.pipelines.registration.registration_icp( model_cloud, obs_cloud, max_correspondence_distance0.02, initpred_pose, estimation_methodo3d.pipelines.registration. TransformationEstimationPointToPlane())逻辑说明ICP只能局部收敛所以初值必须由DenseFusion提供max_correspondence_distance按物体尺寸给一般先给2到3厘米迭代后收窄到5毫米再做一次能显著改善旋转精度。这个技巧对遮挡严重的帧尤其有效因为网络给出的初值已经在正确附近ICP只负责修局部偏差。6.2 用“假抓取”验证位姿容差最后一个技巧是部署前离线验证把验证集真值位姿加高斯噪声旋转±3度、平移±5毫米送入抓取规划里统计成功率。成功率在容差范围内不跌说明位姿模块的误差预算够一旦跌得厉害说明抓取策略对位姿太敏感要么加强姿态模块要么换更鲁棒的抓取点。我的习惯是把这套流程固定成回归脚本每次改动网络、数据增强或部署优化后都重跑一遍。位姿估计做到后面性能瓶颈往往不在网络结构而在数据口径、损失口径和验证口径是否一致。这套DenseFusion方案跑通之后你会发现很多所谓“新方法”的增益其实只是把上述某个环节补对了。希望帮到你。本文还有配套的精品资源点击获取
返回列表