简介:Python-DenseFusion6D物体姿态估计是一份面向计算机视觉开发者与机器学习学习者的完整项目资源,核心实现基于RGB-D图像的6自由度物体位姿估计。通过DenseFusion模型融合彩色与深度特征,适用于机器人抓取、AR/VR交互、工业自动化等场景,适合已有Python与深度学习基础、希望深入理解端到端姿态估计流程的读者。资源压缩包共55个文件,以Python源码为主(20个py文件),涵盖网络定义、特征提取、损失函数、训练与评估脚本,同时包含Shell部署脚本、MATLAB评估工具、CUDA扩展及配置文件,可支撑从数据预处理到模型训练、位姿优化与精度评估的完整流程。包体仅3.51MB,轻量易部署,目录结构清晰。该资源已有1871人学习,内容包含DenseFusion在YCB与LINEMOD数据集上的训练与评估代码、预训练模型说明、实验日志以及可视化对比图,能够帮助读者快速复现实验并掌握密集特征融合、ICP迭代优化等关键实现细节,是学习6D姿态估计的实用参考资料。
1. 6D 物体姿态估计在解决什么:一个比目标检测更难的问题
你是做机械臂抓取或者 AR 叠放的,首先会碰到一个反直觉的事实:2D 检测框在多数场景里并不能支撑一次稳定的抓取。框只告诉你物体在哪片区域,不告诉你它到底怎么躺。用 Python 做 6D 物体姿态估计的团队,最后大多都会绕到一个叫 DenseFusion 的方案上:它把彩色图和高精度深度图逐像素融合,让每个点独立投票估计三维位姿,再靠迭代精化把姿态磨准。这篇文章写给两种人:一种是想把定位从 2D 检测升级到 6D 位姿估计的机器人项目开发者,另一种是对 6D 位姿估计感兴趣、想从原理到 Python 代码把全链路走通的入门者。
2. DenseFusion 的原理:为什么“密集”比“特征拼接”更能抗遮挡
2.1 6D 姿态估计的任务拆解:旋转 1 个矩阵 + 平移 1 个向量
在真实项目里,6D 姿态的数学表达是两个部分:旋转矩阵 R∈SO(3) 和平移向量 t∈R³。物体上任意一个三维点 p 从物体坐标系变换到相机坐标系,只需要做一次 p' = R·p + t。求 R 和 t,就是从图像和深度数据里反推这个刚体变换。
这里面最隐蔽的问题是旋转的自由度表达。欧拉角直观但有万向锁,四元数紧凑但正负符号有歧义,旋转矩阵在训练时最“听话”,但要保证正交性。DenseFusion 的常见做法是让网络直接输出 3×3 矩阵的 9 个数值,再在后处理用 Gram-Schmidt 正交化把它拉回 SO(3)。这个选择会直接影响后面的损失函数设计,因为矩阵元素之间不是独立的。
平移向量相对简单,网络输出 3 个数值,和物体中心的相机坐标直接对应。对于单目 RGB-D 场景,深度图的缩放因子和相机内参一旦对齐,平移量可以直接落到物理单位(毫米或米),这也是 RGB-D 方案比纯 RGB 更稳的根本原因。我见过不少项目在纯 RGB 上加了很多技巧,最后精度还是顶不上去,瓶颈往往就出在单目深度歧义上。
2.2 三条路线横评:全局回归、对应点匹配、密集融合
在 DenseFusion 出现之前,业界主流的姿态估计路线大致分三类。我习惯用下面这张表来对比,方便在项目选型时快速判断该走哪条路:
| 路线 | 代表做法 | 输入 | 抗遮挡能力 | 对低纹理的容忍度 | 工程复杂度 |
|---|---|---|---|---|---|
| 全局回归 | PoseCNN、SSD-6D | RGB | 弱 | 中 | 低 |
| 对应点匹配 | PVN3D、PVNet | RGB(RGB-D) | 中 | 中 | 中 |
| 密集融合 | DenseFusion | RGB-D | 强 | 强 | 中高 |
全局回归路线的思路是让网络把整张图的特征压成一个全局向量,再直接输出 R 和 t。这个做法在物体完整、背景干净时足够用,但一旦物体被遮挡一部分,全局特征就被“污染”了,姿态预测的方向会跟着漂。对应点匹配是另一种思路:先预测物体每个点在图像中的可见位置,再去求解 PnP。它对遮挡有一定抵抗力,但前提是能可靠地建立 3D 点和 2D 点的对应关系,在低纹理物体上效果明显打折。
DenseFusion 选的路子很巧妙:不回归单一全局向量,也不强求找对应点,而是让物体分割区域内的每一个像素(在三维空间里就是每一个点)独立地预测一个姿态,最后再做一次带置信度的投票。每个点看到的都是物体的一小块局部几何,既不会被远处背景污染,又能利用局部邻域的几何纹理信息。即使单个点被遮挡,其他可见点照常投票,这正是密集融合在堆叠物体场景表现出色的原因。
2.3 颜色和几何到底怎么“揉”:像素级融合的工程直觉
RGB-D 方案的核心问题不是“用两份数据”,而是“用什么粒度融合”。最粗糙的做法是把彩色图输入 CNN 提特征、深度图提特征,再把两张特征图 concat 起来。这种特征拼接的毛病在于:卷积网络输出的特征图分辨率已经缩得很小,空间细节模糊掉了,深度图上局部微小的凸起和凹陷对应的几何线索在高层特征里已经消失殆尽。
DenseFusion 的做法是把融合下沉到每个三维点上。深度图反投影出三维点云,每个点先经过一个 PointNet 式的逐点卷积提取几何特征;同时把物体区域的彩色图通过 CNN 提特征,保持特征图与原图分辨率有对应的空间位置;然后把每个三维点按照相机内参重新投影回彩色特征图上,用双线性采样取出该位置的颜色特征,和几何特征拼在一起,这个拼接后的向量就是“这个点自己的 RGB-D 特征”。
这种像素级密集融合有工程上直接的好处:分辨率损失只发生在颜色分支内部,而颜色分支的特征图依然可以对齐到点云坐标;几何分支则完全保留点云的空间精度。两个分支的信息在逐点粒度完成交割,后端的姿态估计头就有了既有纹理又有形状的输入。DenseFusion 这个名字里的 Dense,指的就是这种逐像素的融合粒度,而不是把整张特征图整体拼起来。
3. 用 Python 搭 DenseFusion 训练管线:模型、损失函数与最小训练循环
3.1 环境与依赖:一张表搞定启动配置
我做过几个姿态估计项目,最后都回到 Python + PyTorch 的组合。PyTorch 的自动求导和动态图对旋转矩阵正交化这类非标准操作非常友好,调试也比静态图方案顺手。常见的环境配置可以看作这样一张表:
| 依赖 | 版本建议 | 用途 |
|---|---|---|
| Python | 3.8 及以上 | 运行环境 |
| PyTorch | 1.10+(2.x 也可) | 模型训练与推理 |
| torchvision | 与 PyTorch 版本匹配 | 颜色编码器接口 |
| OpenCV | 4.5+ | 图像读写与预处理 |
| scipy | 1.8+ | 读取 YCB-Video 的 meta.mat 标签 |
| numpy | 1.21+ | 点云计算 |
我没有在表里写死 CUDA 版本,因为显卡驱动、PyTorch 和 CUDA 是配套的,强行固定反而容易翻车。装完环境后,先用一个最简单的命令验证 GPU 可用:
提示:安装完依赖后先跑
python -c "import torch; print(torch.cuda.is_available())",输出 True 再继续,否则后面每一步训练都会慢得让人怀疑人生。
3.2 DenseFusionNet 模型主体:双编码器与逐点姿态头
模型主体按“颜色分支 → 几何分支 → 逐点融合 → 姿态头”四段拆开。下面这个实现以教学可读性优先,保留了论文结构里最关键的部分,省去一些工程压缩技巧。
# model.py import torch import torch.nn as nn import torch.nn.functional as F class DenseFusionNet(nn.Module): def __init__(self, color_feat_dim=256, point_feat_dim=256): super().__init__() # 颜色分支:把物体区域的 RGB patch 压成 1/16 分辨率特征图 self.color_encoder = nn.Sequential( nn.Conv2d(3, 32, 3, 2, 1), nn.BatchNorm2d(32), nn.ReLU(), nn.Conv2d(32, 64, 3, 2, 1), nn.BatchNorm2d(64), nn.ReLU(), nn.Conv2d(64, 128, 3, 2, 1), nn.BatchNorm2d(128), nn.ReLU(), nn.Conv2d(128, color_feat_dim, 3, 2, 1), nn.BatchNorm2d(color_feat_dim), nn.ReLU(), ) # 几何分支:输入 [B, 3, N] 的点云,逐点提取几何特征 self.point_encoder = nn.Sequential( nn.Conv1d(3, 64, 1), nn.BatchNorm1d(64), nn.ReLU(), nn.Conv1d(64, 128, 1), nn.BatchNorm1d(128), nn.ReLU(), nn.Conv1d(128, point_feat_dim, 1), nn.BatchNorm1d(point_feat_dim), nn.ReLU(), ) # 逐点融合 + 姿态头:输出 12 维 = 9 维旋转矩阵 + 3 维平移 self.pose_head = nn.Sequential( nn.Conv1d(color_feat_dim + point_feat_dim, 512, 1), nn.BatchNorm1d(512), nn.ReLU(), nn.Conv1d(512, 1024, 1), nn.BatchNorm1d(1024), nn.ReLU(), nn.Conv1d(1024, 12, 1), ) def forward(self, color_patch, points, uv_feat): """ color_patch: [B, 3, H, W] points: [B, 3, N] 相机坐标系下的点云 uv_feat: [B, N, 2] 每个点在颜色特征图上的采样坐标 """ color_feat = self.color_encoder(color_patch) # [B, C, H/16, W/16] point_feat = self.point_encoder(points) # [B, C, N] B, C, h, w = color_feat.shape # 把 uv 坐标从特征图坐标系归一化到 [-1, 1],给 grid_sample 用 uv_norm = uv_feat.clone() uv_norm[..., 0] = 2.0 * uv_norm[..., 0] / (w - 1) - 1.0 uv_norm[..., 1] = 2.0 * uv_norm[..., 1] / (h - 1) - 1.0 uv_norm = uv_norm.unsqueeze(1) # [B, 1, N, 2] # 双线性采样:每个点从其投影位置取颜色特征 color_feat_sampled = F.grid_sample( color_feat, uv_norm, mode='bilinear', align_corners=True ) # [B, C, 1, N] color_feat_sampled = color_feat_sampled.squeeze(2) # [B, C, N] # 逐点拼接颜色特征和几何特征 fused = torch.cat([color_feat_sampled, point_feat], dim=1) raw_pose = self.pose_head(fused) # [B, 12, N] rot = raw_pose[:, :9] # [B, 9, N] trans = raw_pose[:, 9:] # [B, 3, N] return rot, trans代码逻辑说明
forward 的关键在于uv_feat,它由数据加载器预先算好:先把掩膜内每个三维点反投影回彩色原图的像素坐标,再除以编码器的下采样倍数 16,就得到特征图上的采样坐标。grid_sample对这个坐标做双线性采样,取出该位置的颜色特征,与 PointNet 输出的几何特征拼在一起。最后姿态头对每个点独立输出一组旋转和平移。
这里有几个参数值得注意。color_encoder的四层卷积把分辨率逐级减半,最终到 1/16,所以uv_feat的坐标也必须按 1/16 缩放。双线性采样的align_corners=True要固定,否则 UV 坐标换算会和特征图边界差半个像素,训练时看不出,测试时姿态会莫名其妙偏 1~2 毫米。如果你显存比较紧张,把中间层的通道数从 512/1024 降到 256/512 也能跑,只是精度会跟着掉一点。
3.3 损失函数:旋转和平移联合起来才不会崩
直接对旋转矩阵元素做 L1 损失不是个好主意,因为矩阵元素和空间角度不是线性关系。更贴近几何意义的做法是把点云用预测姿态和真值姿态分别变换一次,计算两个点云之间的距离差,这样旋转和平移被统一在同一套物理单位里。
# loss.py import torch def pose_loss(rot_pred, trans_pred, rot_gt, trans_gt, points): """ 逐点姿态损失:把物体点云分别用预测和真值姿态变换, 计算两个变换点云之间的平均欧氏距离。 """ B, _, N = points.shape # 逐点预测姿态展开成 [B, N, 3, 3] 和 [B, N, 3, 1] rot = rot_pred.permute(0, 2, 1).view(B, N, 3, 3) trans = trans_pred.permute(0, 2, 1).view(B, N, 3, 1) # 点云广播为 [B, 1, 3, N] pts = points.unsqueeze(1) # 第 i 个点用自己预测的姿态去变换整片点云 pts_pred = torch.matmul(rot, pts) + trans # [B, N, 3, N] # 真值姿态只有一个,广播到每个点 r_gt = rot_gt.unsqueeze(1) # [B, 1, 3, 3] t_gt = trans_gt.unsqueeze(1).unsqueeze(-1) # [B, 1, 3, 1] pts_gt = torch.matmul(r_gt, pts) + t_gt # [B, 1, 3, N] # 每个点的预测姿态对整片点云产生的平均距离,就是该点的损失 dist = (pts_pred - pts_gt).norm(dim=2).mean(dim=2) # [B, N] return dist.mean()代码逻辑说明
这个损失函数把每个点的姿态预测都当作一个独立的“候选观测者”,用它把整片点云变换到预测位姿,再和真值变换后的点云比较。距离用欧氏距离的均值,单位与点云单位一致;深度图以米为单位时,loss 的数值就是米。
你可能会问:为什么不直接对旋转矩阵做差?因为旋转矩阵的 9 个分量的差值没有直观的物理含义,9 个分量的微小误差可能在空间角度上放得很大。而这个损失把旋转和平移统一成“空间点距离”,旋转差 5° 和平移差 5mm 对最终抓取位置的影响可以被公平地加权。如果抓取任务对旋转更敏感,可以在 loss 里给旋转部分额外加权,但一般先不加,跑一把基线再说。
3.4 训练循环:数据增强、学习率和遮蔽采样
训练循环本身不长,但数据增强决定了模型能不能在遮挡场景下站稳。我常用的三招是:绕相机光轴(深度轴)随机旋转点云 ±30°,把同样的旋转左乘到旋转真值和平移真值上;给深度值加 0.005m 的高斯噪声,模拟深度传感器的读数抖动;随机保留 60%~100% 的点,模拟物体被部分遮挡。
# train.py 节选 import torch from torch.optim import Adam from torch.optim.lr_scheduler import StepLR model = DenseFusionNet().cuda() optimizer = Adam(model.parameters(), lr=1e-4) scheduler = StepLR(optimizer, step_size=10, gamma=0.5) for epoch in range(60): for batch in train_loader: color, points, uv_feat, rot_gt, trans_gt = [x.cuda() for x in batch] optimizer.zero_grad() rot_pred, trans_pred = model(color, points, uv_feat) loss = pose_loss(rot_pred, trans_pred, rot_gt, trans_gt, points) loss.backward() optimizer.step() scheduler.step() print(f"epoch {epoch:02d}, loss = {loss.item():.4f}")代码逻辑说明
学习率从 1e-4 开始,每 10 个 epoch 减半,跑 60 个 epoch 作为基线。这个配置在 YCB-Video 单类物体上通常能稳定收敛;如果 loss 震荡,优先检查数据对齐而不是调学习率,我踩过几回,最后发现都是数据问题。数据增强里的旋转要和平移真值同步左乘同一个旋转矩阵,否则几何关系就被破坏,模型在验证集上会出现诡异的视角偏差。
4. 数据侧的准备:YCB-Video 目录、数据加载器与掩膜采样
4.1 YCB-Video 数据集结构:文件与标签在哪
做 6D 位姿估计绕不开 YCB-Video 数据集,它包含 21 类常见物体,覆盖了不同纹理、材质和形状,还有大量堆叠遮挡场景,是这一领域事实上的基准。它的目录结构我整理在下面:
| 目录/文件 | 内容 |
|---|---|
| data/000001/-depth.png | 未配准的原始深度图 |
| data/000001/000001-color.png | 彩色图 |
| data/000001/000001-depth.png | 配准后的深度图 |
| data/000001/000001-label.png | 逐像素物体标签 |
| data/000001/000001-meta.mat | 相机内参、旋转矩阵、平移向量、物体类别 |
| models/ | 物体的 3D 模型(obj 与纹理贴图) |
最常见的坑就藏在两个 depth 文件里:带横杠的-depth.png是没和彩色图对齐的原始深度,直接拿来做训练,点云投影到彩色图上会有肉眼可见的错位。训练和评估一律用000001-depth.png这种带帧号的配准深度图。
4.2 数据加载器:从彩色图到归一化点云
下面这个 Dataset 实现把加载和预处理写在一起。注意深度图的像素值单位是毫米,要先除以 1000 转成米,再结合相机内参反投影成三维点。
# dataset.py import os import cv2 import numpy as np import scipy.io as sio import torch from torch.utils.data import Dataset class YCBVideoDataset(Dataset): def __init__(self, root, list_file, num_points=1024): self.root = root self.num_points = num_points with open(list_file) as f: self.samples = [line.strip() for line in f if line.strip()] def __len__(self): return len(self.samples) def __getitem__(self, idx): sample = self.samples[idx] # 形如 data/000001/000001 seq, frame = sample.split('/')[-2], sample.split('/')[-1] data_dir = os.path.join(self.root, seq) color = cv2.imread(os.path.join(data_dir, f'{frame}-color.png')) depth = cv2.imread(os.path.join(data_dir, f'{frame}-depth.png'), cv2.IMREAD_UNCHANGED) label = cv2.imread(os.path.join(data_dir, f'{frame}-label.png'), cv2.IMREAD_UNCHANGED) meta = sio.loadmat(os.path.join(data_dir, f'{frame}-meta.mat')) # 取这一帧第一个物体的类别 ID,并生成掩膜 obj_idx = 0 cls_idx = int(meta['cls_indexes'][0, obj_idx]) mask = label == cls_idx ys, xs = np.nonzero(mask) if len(xs) < 16: # 掩膜区域太小,直接跳过或换一帧 return self.__getitem__((idx + 1) % len(self.samples)) # 随机采样固定数量的像素,避免不同帧点数不一致 if len(xs) < self.num_points: pick = np.random.choice(len(xs), self.num_points, replace=True) else: pick = np.random.choice(len(xs), self.num_points, replace=False) px, py = xs[pick], ys[pick] # 深度图单位是毫米,转成米并反投影 intrinsic = meta['intrinsic_matrix'] fx, fy = intrinsic[0, 0], intrinsic[1, 1] cx, cy = intrinsic[0, 2], intrinsic[1, 2] z = depth[py, px] / 1000.0 x = (px - cx) * z / fx y = (py - cy) * z / fy points = np.stack([x, y, z], axis=1) # [N, 3] # 去中心化:让网络学相对几何,而不是绝对坐标 centroid = points.mean(axis=0) points = points - centroid # 真值旋转/平移也同步调整 pose = meta['poses'][:, :, obj_idx] rot_gt = pose[:3, :3].astype(np.float32) trans_gt = (pose[:3, 3] - centroid).astype(np.float32) # 计算每个点在颜色特征图上的采样坐标(特征图分辨率是 1/16) u_feat = px.astype(np.float32) / 16.0 v_feat = py.astype(np.float32) / 16.0 uv_feat = np.stack([u_feat, v_feat], axis=1) # [N, 2] # 颜色 patch 取包含物体的外接框,统一缩放到 256x256 x1, x2 = px.min(), px.max() y1, y2 = py.min(), py.max() patch = color[max(y1-5, 0):y2+6, max(x1-5, 0):x2+6] patch = cv2.resize(patch, (256, 256)) return ( torch.from_numpy(patch.transpose(2, 0, 1)).float() / 255.0, torch.from_numpy(points.T).float(), # [3, N] torch.from_numpy(uv_feat).float(), torch.from_numpy(rot_gt).float(), torch.from_numpy(trans_gt).float(), )代码逻辑说明
这段代码最关键的一步是“去中心化”。把点云减去质心后,网络学的是物体相对形状的几何,而不是它在相机坐标系里的绝对位置,训练稳定性会好很多。但代价是平移真值必须同步减去质心,否则损失函数里几何距离完全对不上,训练直接崩掉。
uv_feat是为颜色特征图准备的采样坐标,因为模型里的color_encoder把输入缩小了 16 倍,所以原始像素坐标除以 16 才是特征图坐标。颜色 patch 先把物体区域连同 5 像素边距裁出来,再缩放到 256×256。如果掩膜像素太少,len(xs) < 16时递归换一帧,这样比硬撑着训练稳定得多。
4.3 训练/验证列表与类别标签的映射
YCB-Video 官方会提供 train 和 test 的列表文件,每行格式是data/000001/000001这种“目录/帧号”。如果你自己划分数据,注意保持训练和验证的序列不重叠,否则同一段视频的相邻帧会同时出现在两边,验证指标会虚高得离谱。
类别标签的映射也容易绕晕:label.png里的像素值是物体类别 ID(1-based),meta['cls_indexes']同样是类别 ID,两者直接相等。而模型目录里的003_cracker_box这种序号,是 YCB 官方的额外编号,和 label 里的类别 ID 不是一套体系,需要单独维护一张映射表。我在前期吃过这个亏,把cls_indexes当成模型目录序号去读 obj 文件,结果模型和训练数据根本对不上。
5. 避坑指南:DenseFusion 训练与落地时的 5 个常见问题
5.1 深度图和彩色图“视觉对齐但数值上没对齐”
现象:训练 loss 能降,但可视化发现点云投影到彩色图上偏移了 3~5 个像素,验证集 ADD 指标一直上不去。
原因:YCB-Video 里的-depth.png是没有和彩色图配准的原始深度,两个传感器各自有自己的视差,直接用就是错位。这类问题在训练时不容易发现,因为网络会“记住”这种偏移模式。
解决:统一读000001-depth.png,且内参用meta['intrinsic_matrix']里带帧号的那个矩阵。如果是从别的数据集自采数据,务必先用标定板或 APRILTAG 做一次 RGB-D 对齐,再谈姿态估计。
5.2 四元数符号模糊:同一个姿态两种表达
现象:用四元数作为旋转输出时,训练 loss 降得很快,但验证集准确率很低,而且有时候同一个物体相隔几帧预测的四元数符号不断跳变。
原因:四元数 q 和 -q 表示同一个三维旋转,网络如果对两个样本输出符号相反的等价四元数,直接做 L1 损失就会把它们当作完全不同的姿态来惩罚。
解决:要么在训练前对真值四元数做符号统一(比如保证实部大于 0),要么干脆学 DenseFusion 的输出 9 维旋转矩阵再正交化。我后来一直在用旋转矩阵路线,符号问题一次性消失。
5.3 点云去中心化后平移真值没跟着改
现象:loss 在初期降到 0.1 左右就不再下降,甚至反弹,平移误差始终在 5~10 厘米的范围抖动。
原因:数据加载器里对点云做了减去质心的操作,但平移真值仍保持相机坐标系下的原始值。模型学到的点云是局部坐标,真值却是全局坐标,两者根本对不上。
解决:去中心化时同步更新平移真值,trans_gt = pose[:3, 3] - centroid。如果你还会对点云做缩放归一化,平移真值也要除以同样的缩放系数。这个坑我在第三章数据加载器里已经处理过,但每次换数据集都会有人再踩一次。
5.4 掩膜区域太小,点云采样崩了
现象:某些物体在画面里只有几十个像素,强行采样 1024 个点时重复点太多,模型对这块物体的预测姿态在几个候选之间漂移,看起来像“抖”。
原因:点云下采样用了replace=False,当掩膜点数不足时要么报错,要么退化成大量重复点,局部几何被重复点放大,破坏了真实的形状分布。
解决:训练时设置最小点数阈值(我一般用 64),低于阈值直接换一帧;推理时如果掩膜太小,宁可不上机械臂,也不要给一个抖动姿态去执行抓取。还有一招是训练时随机丢弃点云中的 10%~20% 的点,让模型对这种稀疏输入更鲁棒。
5.5 迭代精化“越修越歪”
现象:加了迭代精化网络后,第一次迭代精度提升,第二次迭代 ADD 误差反而变大,损失在验证集上不降反升。
原因:精化网络训练时的初始姿态分布和测试不一致。如果训练时初值都是从真值附近采样的,测试时上游检测网络给的初值误差一大,精化网络就会把姿态推离正确位置,因为它从没见过这么差的输入。
解决:训练精化网络时,用比较大的高斯扰动来造初值,旋转扰动可以到 ±20°,平移扰动到 ±2cm;推理时控制迭代次数,一般 1~2 次足够,不要贪多。这个超参数要在验证集上单独调,不能和姿态头训练共用一套配置。
6. 进阶:迭代精化与 ADD 评估的实战技巧
6.1 迭代精化网络:用残差旋转估计更稳
迭代精化的思路是:把当前估计姿态作用到点云上,让精化网络预测一个残余的旋转和平移,叠加回去。这里有个工程小技巧,残余旋转不要直接回归旋转矩阵或欧拉角,而是回归 so(3) 李代数的 3 维向量,更新时用指数映射转回旋转矩阵,数值稳定性好很多。
# refiner.py 核心逻辑 import torch import torch.nn as nn class PoseRefiner(nn.Module): def __init__(self): super().__init__() self.encoder = nn.Sequential( nn.Conv1d(3, 128, 1), nn.ReLU(), nn.Conv1d(128, 256, 1), nn.ReLU(), ) self.residual_head = nn.Sequential( nn.Conv1d(256, 128, 1), nn.ReLU(), nn.Conv1d(128, 6, 1), # 3 维旋转残差 + 3 维平移残差 ) def forward(self, points_cur): # points_cur: 当前姿态变换后的三维点 [B, 3, N] feat = self.encoder(points_cur).mean(dim=2, keepdim=True) residual = self.residual_head(feat) # [B, 6, 1] return residual[:, :3], residual[:, 3:]精化网络最终输出的是残差量,迭代两次时点云会被更新两次。我在项目里的习惯是先固定姿态主干,单独训练精化网络 20 个 epoch,再联合微调,这样收敛稳定不少。
6.2 ADD 阈值怎么定:对称物体必须换指标
评估 DenseFusion 这类 6D 姿态估计模型,绕不开 ADD 和 ADD-S 两个指标。ADD 是模型点云在预测姿态下变换后的平均距离,而 ADD-S 针对对称物体做了放宽:每个点取预测点到真值最近点的距离,再求平均。
| 指标 | 计算方式 | 推荐场景 |
|---|---|---|
| ADD | 点到点平均距离 | 非对称物体 |
| ADD-S | 每个点取最近距离 | 对称物体(杯子、易拉罐、碗) |
工程上一般以 ADD 小于物体直径的 10% 作为正确阈值,比如一个直径 8cm 的易拉罐,误差要小于 8mm 才算一次成功的姿态估计。别把阈值放宽到 5cm,那在真实抓取场景里已经足够让夹爪打翻旁边的瓶罐了。
现在我做姿态估计项目,一定先跑通一个类别、一台相机、一个抓取场景,再去铺数据量。DenseFusion 的坑主要集中在数据对齐和姿态表示上,这两点想清楚了,后面的路就顺了。希望这些经验帮到你。
本文还有配套的精品资源,点击获取