DenseFusion 6D位姿估计Python复现与排错

简介Python-DenseFusion6D物体姿态估计项目面向计算机视觉学习者及机器人抓取、AR/VR与工业自动化开发者聚焦RGB-D输入下的6D位姿估计任务提供从特征提取、像素级融合到姿态优化的端到端实现适合具备深度学习基础、希望复现并理解DenseFusion的中高级读者。压缩包共55个文件约3.51MB以20个Python脚本为主体配合sh运行脚本、png结果图、txt说明及少量cu、so、h等CUDA与C扩展文件并含预训练模型目录覆盖数据加载、网络构建、损失计算、训练与评估等环节目录划分清晰。已有1872人浏览学习。读者可据此梳理RGB与深度特征融合、Linemod/YCB数据流程、损失函数与姿态细化网络的设计借助训练评估脚本和结果图核对复现思路并积累6D姿态估计项目的环境配置、模块调试与结果分析经验对复现实验与工程落地均有参考价值。1. 从一张 RGB-D 图到 6D 位姿DenseFusion 到底解决了什么在无序抓取和装配线上机械臂要抓的那个零件往往没有纹理还会被旁边的工件挡住半边但工程上仍然需要知道它在相机坐标系下的完整六自由度位姿。用 Python 复现 DenseFusion 这类 6D 物体姿态估计方案价值就在这里它不要求物体整块可见也不指望表面有丰富角点而是把 RGB 外观特征和深度点云的几何特征在逐像素这一层拼起来让 mask 里每个像素独立投一次票再用置信度把票数最高的位姿挑出来。如果你手上有 RGB-D 相机、有物体的 CAD 模型或标好的位姿数据想在遮挡、弱纹理场景下把平移误差压到厘米级、旋转误差压到几度以内这条路线值得投入。下面按数据、网络、损失、排错、评测的顺序把它拆到能直接开工的程度。2. RGB-D 数据先落地标注字段、内参反投影与 mask 采样在 6D 姿态估计里数据准备占用了将近一半的调试时间。网络本身是黑匣子但数据不是——内参错一位、单位差一千倍、mask 膨胀多三个像素网络照样能训练损失照样在降只是你永远拿不到能上机的位姿。所以这一章先把数据链路对齐再谈模型。2.1 一份能直接训练的数据集该有哪些字段不管数据是自己采的还是标注的落到磁盘上时建议统一成下面这套最小字段。字段名可以不同但语义必须一一对应否则后面写 Dataset 时会出现“看起来能跑、实际错位”的情况。提示位姿统一约定为“物体模型坐标系到相机坐标系”平移单位统一为米。整条链路只保留这一种约定不要在不同脚本里混用。字段类型说明scene_idstr 或 int场景编号用于划分训练/测试集frame_idint帧号同一场景内唯一rgb_pathstr彩色图8bit 三通道depth_pathstr深度图16bit 单通道 PNGmask_pathstr实例掩码0 为背景非 0 为该物体K3×3 float相机内参含 fx、fy、cx、cydepth_scalefloat深度原始值到米的除数常见 1000R、t3×3、3物体位姿相机坐标系下单位米最容易出问题的两处一是深度图存的是 uint16 毫米值如果 depth_scale 写成 1点云整体放大一千倍网络仍然“训练得动”但预测的平移永远在几百米量级看着像模型学不会实际是单位错了二是 t 的语义有些标注工具存的是物体中心相对于模型原点的偏移有些存的是相机坐标系下的绝对平移两者差一个旋转混用会让损失在原地打转。划分数据集时也要注意按帧随机切分会造成严重的评估虚高同一场景的相邻帧几乎一样训练集见过测试集自然分高。常见做法是按 scene_id 切分训练场景和测试场景完全不重叠这样测出来的数字才能代表真实泛化能力。2.2 用内参把深度图反投影成点云有了 K 和深度图点云不是“顺便”生成的它是整个几何分支的输入反投影公式必须写对。下面这段是每次开工我都会先单测一遍的函数。import numpy as np def depth_to_points(depth, K, depth_scale1000.0, min_d0.1, max_d3.0): 深度图反投影到相机坐标系返回 (H*W, 3) 点云和有效掩码。 depth: (H, W) uint16 K: (3, 3) 内参 depth_scale: 原始值 - 米 的除数 min_d/max_d: 有效深度区间米过滤 0 值和深度噪声 h, w depth.shape d depth.astype(np.float32) / depth_scale # 统一到米 valid (d min_d) (d max_d) u, v np.meshgrid(np.arange(w), np.arange(h)) # 关键加 0.5 取像素中心否则整幅点云会系统性偏移半个像素 x (u 0.5 - K[0, 2]) * d / K[0, 0] y (v 0.5 - K[1, 2]) * d / K[1, 1] z d pts np.stack([x, y, z], axis-1).reshape(-1, 3) return pts.astype(np.float32), valid.reshape(-1)这段代码的逻辑很直白先把深度值归一化到米再按针孔模型把每个像素沿射线投到三维空间。三个参数值得多说两句。min_d 和 max_d 不是可有可无的装饰深度相机在物体边缘和反光面会输出 0 或明显外飞的值不过滤会直接在点云里制造一片离群点把融合特征带偏。depth_scale 必须和采集设备一致常见的结构化光相机输出毫米、部分设备输出 0.1 毫米这个值写错后面的平移标签和点云尺度就对不上。至于 0.5属于典型的“知道就没事、不知道永远查不出来”的细节——省掉它训练能跑精度会莫名其妙卡在某个水平上不去。2.3 mask 的处理与点数采样把整张图的点云送进网络是最常见的第一版做法也是最先该推翻的做法。桌面、背景、相邻工件会占掉绝大部分点网络学到的“全局几何”其实是场景几何不是物体几何。正确做法是用实例 mask 把物体区域抠出来并且对 mask 做 12 像素的膨胀。膨胀的理由很实在深度图和彩色图之间存在轻微配准误差mask 是照着彩色图标出来的边缘那一圈像素如果严格按 mask 取会把物体表面的点切掉一圈导致估计出的位姿在深度方向上系统性偏一点。膨胀两个像素能把这个偏差吃掉代价是引入少量背景点对逐像素投票的架构来说完全可以接受。采样点数一般固定成一个常数这类方法常用 6890 个点也是很多人形模型的顶点数不够就重复采样多了就随机采样。固定点数的目的是让 batch 里每个样本形状一致避免 collate 时写一堆 pad 逻辑。import torch from torch.utils.data import Dataset import cv2 class PoseDataset(Dataset): def __init__(self, items, num_points6890, mask_dilate2): self.items items # 每条含 rgb/depth/mask/K/R/t 路径 self.num_points num_points self.mask_dilate mask_dilate def __getitem__(self, idx): it self.items[idx] rgb cv2.imread(it[rgb_path])[:, :, ::-1] / 255.0 depth cv2.imread(it[depth_path], cv2.IMREAD_UNCHANGED) mask cv2.imread(it[mask_path], cv2.IMREAD_GRAYSCALE) 0 K it[K] # 1) mask 膨胀抵消彩色/深度配准误差 if self.mask_dilate 0: kernel np.ones((self.mask_dilate * 2 1,) * 2, np.uint8) mask cv2.dilate(mask.astype(np.uint8), kernel) 0 # 2) 只用 mask 内的像素生成点云 pts_all, valid depth_to_points(depth, K, it[depth_scale]) keep (mask.reshape(-1) valid).nonzero()[0] if keep.size 0: # 空 mask 直接跳过别让 NaN 污染 loss return self.__getitem__((idx 1) % len(self)) pts pts_all[keep] # 3) 固定点数采样 choice np.random.choice(len(pts), self.num_points, replacelen(pts) self.num_points) return { rgb: torch.from_numpy(rgb.transpose(2, 0, 1)).float(), pts: torch.from_numpy(pts[choice]).float(), R: torch.from_numpy(it[R]).float(), t: torch.from_numpy(it[t]).float(), K: torch.from_numpy(K).float(), }逻辑说明先把 mask 和有效深度取交集保证进去的每个点既是物体又是有几何意义的再统一采样到固定点数。参数上mask_dilate 取 12物体小、边缘薄的时候取 1物体大、深度噪声重的场景取 2num_points 取 6890 或 4096 都可以前者与人形模板顶点数一致做可视化时方便。2.4 位姿归一化把物体搬到原点附近再学网络直接回归相机坐标系下的绝对平移效果通常很差因为同一个物体在不同场景里的平移跨度可能从 0.3 米到 1.5 米回归头要去拟合一个动态范围极大的量。工程上的常规做法是先减掉点云质心让网络只学“相对于物体中心的偏移”推理时再加回去。旋转部分建议用四元数或六维向量表示别直接回归欧拉角。欧拉角在 90 度附近有万向锁两个数值上差很多的角在物理上是同一个姿态损失函数会被这种不连续折腾得来回震荡。四元数要注意归一化后再转旋转矩阵否则会出现缩放。3. 网络怎么搭两路特征、像素级对齐与每点一个位姿DenseFusion 的结构说起来不复杂一个分支看颜色一个分支看几何然后逐像素融合最后每个点输出一个位姿加一个置信度。但真正把代码写对难点全在“对齐”两个字上——RGB 特征图是二维的点云是一维序列你得让第 i 个点精确地拿到它在图像上对应像素的特征差一行一列融合就变成了随机拼接。3.1 RGB 分支要的是分辨率而不是语义颜色分支用常见的卷积主干就行但不要把它当成分类网络用。分类网络最后一层是全局池化后的 512 维向量那个东西对位姿估计几乎没有信息量因为空间结构全被压没了。这里要取的是中间层的特征图下采样倍率控制在 8 倍以内输出形如 (B, 64, H/8, W/8)。为什么不取 32 倍下采样因为一个小零件在 640×480 的图里可能只占 60×60 像素下采样 32 倍之后只剩 2×2逐像素融合就没有“像素”可言了。分辨率与感受野之间的取舍本质上是“每个点能不能拿到足够有区分度的外观描述”和“这个描述能不能看到上下文”之间的取舍。实际调参时如果发现小物体精度明显差先看的就是特征图分辨率而不是模型容量。训练时这个分支要么冻结预训练权重要么用很小的学习率慢慢放因为姿态回归的梯度相对分类任务要小得多一上来就用同样的学习率主干会在前几个 epoch 里被彻底打乱。3.2 点云分支逐点 MLP 加一层全局上下文几何分支对每个点做几层逐点感知机然后做一次全局最大池化拿到整体形状描述再把这个全局描述拼回每个点的局部特征上。这一步很关键只有局部特征网络看不出“这个点是杯子的把手还是杯身”只有全局特征又退化成了一个物体一个位姿。局部加全局才让每个点既知道自己长什么样也知道自己在整体里的位置。3.3 像素对齐与稠密融合这是整套方案最核心的一段代码。点云里第 i 个点在采样前的原始像素坐标 (u, v) 必须一路传下来采样时同步采样坐标融合时用它去特征图上取特征。import torch import torch.nn as nn import torch.nn.functional as F def sample_pixel_feature(feat_map, uv, size_wh): 按像素坐标从特征图上双线性采样。 feat_map: (B, C, Hf, Wf) uv: (B, N, 2)原始图像像素坐标 (u, v)u 向右、v 向下 size_wh: (W, H)原始图像宽高 w, h size_wh # 归一化到 [-1, 1]align_cornersFalse 下必须用 (u 0.5) / W nx (uv[..., 0] 0.5) / w * 2.0 - 1.0 ny (uv[..., 1] 0.5) / h * 2.0 - 1.0 grid torch.stack([nx, ny], dim-1).unsqueeze(2) out F.grid_sample(feat_map, grid, modebilinear, align_cornersFalse) return out.squeeze(-1) # (B, C, N) class PixelFusion(nn.Module): def __init__(self, c_rgb64, c_pts64, c_fuse64): super().__init__() self.mlp nn.Sequential( nn.Conv1d(c_rgb c_pts, c_fuse, 1), nn.ReLU(), nn.Conv1d(c_fuse, c_fuse, 1), nn.ReLU(), ) def forward(self, f_rgb, f_pts, uv, size_wh): f_rgb_s sample_pixel_feature(f_rgb, uv, size_wh) return self.mlp(torch.cat([f_rgb_s, f_pts], dim1)) class PoseHead(nn.Module): def __init__(self, c_in64): super().__init__() self.body nn.Sequential(nn.Conv1d(c_in, 128, 1), nn.ReLU(), nn.Conv1d(128, 128, 1), nn.ReLU()) self.out nn.Conv1d(128, 8, 1) # 四元数 4 平移 3 置信度 1 def forward(self, x): y self.out(self.body(x)).transpose(1, 2) q F.normalize(y[..., :4], dim-1) # 四元数必须归一化 t y[..., 4:7] c torch.sigmoid(y[..., 7:8]) # 置信度压到 (0, 1) return q, t, c逻辑上分三步先按像素坐标把二维特征图“拉”成与点云一一对应的序列再做 1×1 卷积把两路特征压到同一通道数并融合最后每点输出 8 个值。参数上最需要注意的是采样坐标。grid_sample 的 grid 约定最后一维是 (x, y) 而不是 (row, col)且 align_cornersFalse 时坐标映射要用 0.5 的像素中心约定写错会整体偏移半个特征图表现是位姿精度上不去、可视化时预测框和物体永远差一点点——这个坑非常隐蔽因为代码不会报错。另一个是置信度用 sigmoid 压到 (0, 1)方便后面做加权和聚类。3.4 为什么要给每个像素一个位姿和一个置信度如果只做全局回归一个遮挡严重的物体其可见部分的几何一致性很好不可见部分的推测就全靠先验误差会集中在那里。逐像素预测等于让几千个点各抒己见其中大量落在高置信度表面上的点会给出非常一致的答案落在边缘、反光、深度缺失处的点给出离谱答案但它们的置信度也会被训练得比较低。最后的投票环节就是把这种“内部共识”提取出来。这也解释了为什么这套方案对遮挡鲁棒只要还有一块足够大的表面可见就总有足够多的点投出正确的一票。反过来说如果物体只露出一个特别窄的角可用像素太少共识建立不起来精度必然掉——这是方法本身的下界不是调参能解决的。4. 损失与训练策略置信度加权、对称物体和迭代细化网络能跑通不等于能收敛收敛了不等于精度能用。损失函数的设计直接决定了那些“没投对票”的点会不会把整体带偏而迭代细化决定了最终精度能不能从厘米级压到毫米级。4.1 置信度加权位姿损失每个点都有一个预测位姿和一个置信度损失要同时推动两件事位姿往真值靠置信度往“这个点靠不靠谱”靠。标准写法是置信度加权的位姿误差减去置信度的对数。def densefusion_loss(q_pred, t_pred, c_pred, R_gt, t_gt): q_pred: (B, N, 4) t_pred: (B, N, 3) c_pred: (B, N, 1) R_gt: (B, 3, 3) t_gt: (B, 3) B, N, _ q_pred.shape R_pred quat_to_matrix(q_pred) # (B, N, 3, 3) # 位姿距离用逐点几何误差近似量纲统一到米 # 直接比较 R 的 Frobenius 范数会和米制平移量纲不一致需要加权 rot_err torch.linalg.matrix_norm(R_pred - R_gt[:, None], dim(2, 3)) tra_err torch.linalg.norm(t_pred - t_gt[:, None], dim-1) pose_err rot_err * 0.1 tra_err # 0.1 是量纲平衡系数 c c_pred.squeeze(-1).clamp(1e-4, 1.0) # 防止 log(0) - -inf loss pose_err * c - 0.01 * torch.log(c) return loss.mean()这段损失的含义需要掰开看。pose_err 越小、c 越大加权项越小说明“这个点又准又自信”pose_err 大、c 小乘积也不大说明“这个点不准但它自己知道”不惩罚。真正被惩罚的是“不准但很自信”的点。后面那个 -log(c) 项是防止网络偷懒把所有 c 都压到接近 0 的后悔药0.01 是平衡系数调大会让置信度整体偏高、区分度下降调小则容易让网络倾向于压低置信度。实际调的时候可以先固定成 0.01观察训练中 c 的均值落在 0.30.7 之间比较健康长期贴着 0.95 或者 0.05 都说明这一项没起到作用。rot_err 和 tra_err 量纲不同前者是无量纲的矩阵差后者是米。直接把两者相加会出现“旋转已经学得很好、平移还在飘”的情况因为平移项的梯度被旋转项淹没了。常见做法是给旋转项乘一个平衡系数上面写的 0.1或者把旋转误差换算成角度、把平移误差除以物体直径让两者都在 01 附近。4.2 对称物体别让损失追一个不存在的唯一解杯子、盒子、法兰盘这类物体存在对称性同一个外观对应多个数学上等价但数值完全不同的旋转矩阵。如果按普通物体那样算损失网络会在这些等价解之间来回横跳损失曲线看着在降但降不下去最终精度也上不去。处理办法有两种。一种是把位姿误差换成对对称友好的形式对每个预测点取它到真值模型上所有点的最近距离均值这样任何等价姿态都能得到低损失。另一种是检测物体的对称轴在损失里对旋转的对应角度做取模。我一般倾向第一种因为它对任何物体都成立不需要预先知道对称性代价是每个 batch 都要算一次点对点距离矩阵显存和算力消耗大一些。实现上要把物体模型顶点预先算好batch 内复用不要让这部分成为瓶颈。def symmetric_pose_error(R_pred, t_pred, R_gt, t_gt, model_pts): model_pts: (M, 3) 物体模型顶点单位米 逐点取最近距离对任意对称等价姿态都给出低损失 pred torch.einsum(bnij,mj-bnmi, R_pred, model_pts) t_pred[:, :, None] gt torch.einsum(bij,mj-bmi, R_gt, model_pts) t_gt[:, None] d torch.cdist(pred.reshape(-1, model_pts.shape[0], 3), gt) return d.min(dim-1).values.mean(dim-1).reshape(R_pred.shape[:2])4.3 迭代细化第二轮把残差补回来第一轮网络输出的位姿通常已经“大致对”但精度卡在一个平台上下不去。迭代细化的思路很朴素把预测出来的位姿作用到原始点云上得到一组“已经对齐到模型坐标系”的点再把这组点连同原来的 RGB 特征喂进第二轮让它只学剩下的残差。第二轮的网络结构和第一轮基本一致区别在于输入点云换成了变换后的点输出的是对位姿的修正量。常见做法是让第二轮输出的仍然是一个完整的位姿然后和第一轮的位姿做复合而不是直接回归增量——复合形式在数值上更稳也更容易调试。训练时有个细节第二轮的真值应该用某个中间位姿来构造而不是永远用最终真值否则第二轮见到的输入分布和推理时不一致。工程上常从“真值加随机扰动”和“第一轮的实际预测”里各采一半让模型既能从接近真值的输入里学精细修正也能从真实误差里学纠偏。4.4 关键超参与收敛判断参数常用取值说明采样点数40966890太小共识不稳太大显存吃紧RGB 特征下采样倍率8再大丢空间细节融合通道数64128直接影响参数量batch size816受点数与特征图尺寸限制初始学习率1e-4主干更小1e-5 量级置信度正则系数0.01观察 c 均值调整迭代细化轮数2第三轮收益递减明显收敛判断不要只看总损失。总损失是几千个点误差的平均掩盖了很多信息。我习惯同时看四个量置信度的均值、置信度的方差、位姿误差的中位数以及验证集上的 ADD-S。如果总损失在降但置信度方差在缩小说明网络在“把所有点训成一样自信”这时候融合就退化成了全局平均后面投票环节会挑不出好假设。5. 避坑与排查五个把 loss 骗过去的地方这一节记录的是我在不同项目里反复踩到的坑共同特点是训练不报错、损失看着正常、但结果就是不能用。5.1 位姿“对了”但方向反了现象可视化时预测的坐标系和物体几乎重合但绕某个轴转了 180 度ADD 指标奇高。原因四元数存在 q 与 -q 表示同一旋转的二义性如果损失函数直接比较四元数分量的欧氏距离网络会学到“两种都行”在两者之间反复跳。解决先把预测四元数统一到与真值同半球点积为负就取反再计算损失或者干脆用旋转矩阵或六维表示来计算损失。5.2 训练集精度很高换一个相机就崩现象在自己采的数据上 ADD 能到 1 厘米以内换一台内参不同的相机、或者分辨率改一下误差立刻翻倍。原因网络在训练中隐式记住了训练相机的内参和视野范围点云的绝对尺度和分布成了捷径。解决训练时对点云做随机缩放和小角度旋转增强把内参作为显式输入或者在反投影后把点云按物体直径归一化换相机后必须重新标定 K并至少用少量目标相机数据微调。5.3 深度图有空值损失变成 NaN现象训练几个 epoch 后突然 loss 变成 nan之后再也恢复不了。原因某些帧的 mask 区域内一个有效深度都没有全反光、超量程采样时拿到空的点数组后续归一化除零。解决在 Dataset 里对空样本直接跳过并打印计数训练日志里跟踪被跳过的帧数如果比例超过 5%说明采集环节的参数设置有问题需要调整相机的深度范围或物体摆放。5.4 置信度全部趋近于 1现象训练稳定后所有点的置信度都在 0.98 以上投票环节失效最终精度取决于运气。原因置信度正则系数太小-log(c) 项压不住网络把 c 拉满的倾向本质上网络发现“全自信”也能让加权损失降下去。解决把正则系数上调同时在验证时统计置信度的分布理想状态是有明显的双峰——靠谱的点高、边缘点低。如果一直是单峰说明融合特征没学到区分度要回头检查 RGB 分支的分辨率是不是太低。5.5 评估指标虚高实际抓取失败现象测试集 ADD-S 很漂亮但真机上抓取成功率很低。原因评测用的位姿是相机坐标系下的而机械臂需要的是基坐标系下的手眼标定矩阵的误差被完全掩盖另外评测时只算了成功检测的样本漏检的那些被静默丢弃。解决把评估拆成两步先算检测率有多少帧成功给出位姿再在检测到的样本上算精度两个数字一起报真机上线前用手眼标定做一次闭环验证让机械臂实际走到预测位姿上看末端和物体的偏差。6. 推理与评测从几千个假设里选一个以及 ADD-S 怎么算才不错训练完之后网络对每个样本输出几千个位姿假设必须有一个选优环节。最直接的做法是按置信度取最大值实测很不稳单个点的偶然高置信度就能把结果带偏。我一般用均值偏移聚类以每个假设的平移为特征带宽设成物体直径的 5%10%迭代收敛后取包含最多假设的那一簇簇内按置信度加权平均。def cluster_pose_hypotheses(t_pred, c_pred, bandwidth0.02, iters10): t_pred: (N, 3) c_pred: (N,) 带宽单位米 centers t_pred.clone() for _ in range(iters): d torch.cdist(centers, t_pred) # (N, N) w torch.exp(-(d / bandwidth) ** 2) * c_pred[None] new (w t_pred) / (w.sum(dim1, keepdimTrue) 1e-8) if (new - centers).norm(dim1).max() 1e-5: break centers new d torch.cdist(centers, t_pred) members (d bandwidth).float() * c_pred[None] best members.sum(dim1).argmax() return t_pred[d[best] bandwidth].mean(dim0)带宽是最关键的一个参数设太小同一个位姿会被拆成好几簇选出来的是偶然形成的小簇设太大不同位姿的假设会被并到一起平均之后就飘了。经验值取物体直径的 5%10%物体小取上限物体大取下限。评测环节对称物体必须用 ADD-S非对称物体用 ADD。区别只有一行ADD 是预测点与真值对应点一一配对求距离ADD-S 是每个预测点找真值点集里最近的那个。def add_s(R_pred, t_pred, R_gt, t_gt, model_pts): pred model_pts R_pred.T t_pred gt model_pts R_gt.T t_gt d torch.cdist(pred, gt) # (M, M) return d.min(dim1).values.mean().item() # 对称版本这里有个容易算错的地方model_pts 必须和训练时用的是同一套模型顶点单位也必须是米否则算出来的数字没有可比性另外报告指标时要说清楚用的是物体直径归一化后的值还是绝对值前者跨物体可比后者更贴近实际精度。真机上线前我会做三件事用同一张图分别跑离线和在线推理确认前后处理完全一致拿十帧标注数据算一遍 ADD-S和训练日志里的验证集数字对得上最后让机械臂空跑一次预测位姿看末端执行器和物体中心的实际偏差。这三步能挡掉大部分“指标很好、现场翻车”的情况。踩过的坑里最耗时间的从来不是网络结构而是坐标系和单位——它们不会报错只会让结果安静地错下去。我的习惯是每接一个新数据集先用十行脚本把点云和深度图叠一起可视化一次确认两者的轮廓严丝合缝再开始训练。希望帮到你。本文还有配套的精品资源点击获取