简介：这份Python-DenseFusion 6D物体姿态估计项目源码包，面向计算机视觉、机器人抓取与AR/VR方向的学习者和开发者，用于理解和复现基于RGB-D图像的端到端6D位姿估计流程。压缩包共55个文件，约3.51MB，以20个Python源码为核心，覆盖网络结构、数据加载、训练与评估脚本；辅以Shell运行脚本、PNG结果图、TXT日志说明，以及CUDA/C++扩展和预训练模型相关文件，便于按模块查阅与搭建实验。已有1872人学习下载，说明其在姿态估计入门与项目复现中具有一定参考价值。包内DenseFusion-master目录包含Linemod与YCB数据配置、训练与评估入口、损失与位姿优化实现、可视化对比图等，可帮助读者梳理从特征提取、像素级融合到姿态精化的完整链路，并作为课程设计、科研复现或二次开发的代码基础。

1. 从一张 RGB-D 图到六自由度位姿：这套 Python 版 DenseFusion 能解决什么

拿到一张 RGB-D 图，要回答的其实只有两件事：物体在哪，物体朝哪个方向。机械臂抓取、AR 装配引导、无序分拣这些场景里，后者的要求往往比前者更严格：位置差两厘米，吸盘勉强还能处理；姿态差二十度，就会直接抓空。这套 Python 版的 DenseFusion 六自由度位姿估计实现，走的是机器学习里看起来朴素但稳的路子——把颜色和深度分成两路提特征，再在每个像素上融合，让每个像素都投出一份自己的位姿估计，最后用置信度打分选出最可信的那份。它适合已经跑通目标检测、想继续接入位姿估计的开发者，也适合需要一个能改、能训、能换成自有数据集的基线的人。

2. 环境和目录：把跑通成本压进半小时

先说明：这个仓库不是那种pip install就完事的现代工程。它更像一份研究代码的整理版，能跑，但需要你把版本、目录、权重三件事对齐。我见过不少人卡在第一天，不是模型难，而是环境装完跑出个undefined symbol。这一节的目标很单纯——让你在半小时内看到第一帧的输出位姿，精度的事放到后面。

2.1 版本组合与硬件门槛的取舍

这类 6D 位姿代码大多沿用的是几年前那套写法：Variable包一层、tensor.data[0]取标量、align_corners不显式传。你在 PyTorch 2.x 上直接跑，最常见的报错是grid_sample行为变了、torch.symeig已经改名为torch.linalg.eigh。我一般不去追最新，而是提前固定一套旧版本，这样一晚上能省下来。

提示：如果你手上这份代码带了需要编译的点云算子（最远点采样、KNN 之类），而编译一直失败，先用纯 PyTorch 的暴力实现顶着。慢三到五倍，但能把整条链路跑通，把问题从“装不上”降级成“有点慢”。

还有一点新手容易忽略：DenseFusion 那套思路是要在每个有效像素上回归一个位姿的，一个 640×480 的裁剪块里有效点可能有几万个。这不是分类网络那种“batch 拉到 32”的玩法，训练时 batch=1 到 4 是常态。所以别拿一张 4GB 的卡去试训练，会浪费你半天时间在调 batch 上。

2.2 目录骨架与权重的放置

我习惯先把目录建死，再往里拷东西。原因是这类代码里到处是相对路径，datasets/、weights/这种目录名写死在 config 里，你随便换个名字就得全局搜一遍。

# 代码和数据分离，权重单独放，避免 git 把几百兆的 pth 一起带进去 mkdir -p df6d/{configs,models,utils,datasets,scripts} mkdir -p df6d/weights mkdir -p data/{rgb,depth,mask,meta} mkdir -p df6d/logs # 确认目录结构，注意 depth 和 rgb 必须是同一套命名，靠文件名对齐 ls data/rgb | head -3 ls data/depth | head -3

几个目录的职责要分清楚：

• configs/放 yaml 或 py 配置文件，物体类别、点采样数、对称物体列表、数据集路径全在这里改，不要散落在代码里。
• models/是网络定义，改结构只动这里。
• datasets/放 Dataset 类和它的注册逻辑，接自有数据时基本只改这个目录。
• weights/放预训练权重。文件名通常是作者起的，别改，改了就去找引用它的那一行。
• logs/放 tensorboard 和 checkpoint，训练跑起来之后这个目录会涨得很快，注意盘。

rgb和depth的对应关系靠文件名，这点很关键。深度图常见的格式是 16 位 PNG（单位毫米）或 32 位浮点 npy（单位米），下面读取的时候我会强制转成米。如果你的深度图和 RGB 尺寸不一致，说明它们没对齐，需要先用内参和深度图尺度做一次重投影，这个坑在第 5 章会展开。

2.3 一帧推理：先确认链路是通的

不要一上来就训练。先写一个最小推理脚本，喂一张图进去，看输出是不是个合理的位姿。这一步花十分钟，能帮你把“环境问题”和“模型问题”分开。

import numpy as np import torch import cv2 # 模块名按你手上这份代码的实际命名替换，不要照抄 from models.pose_net import build_model from utils.geometry import depth_to_xyz CKPT = "df6d/weights/你的权重文件.pth" def load_intrinsic(meta_path): """内参我一定从数据集的 meta 里读，不用估算值""" K = np.loadtxt(meta_path) # 3x3 return torch.tensor(K, dtype=torch.float32) def main(): model = build_model(num_points=8000, num_obj=8) state = torch.load(CKPT, map_location="cpu") # 常见坑：保存时外面套了一层 "model" 或 "state_dict" state = state.get("model", state.get("state_dict", state)) model.load_state_dict(state, strict=False) model.eval().cuda() rgb = cv2.imread("data/rgb/0001.png")[:, :, ::-1] / 255.0 depth = cv2.imread("data/depth/0001.png", cv2.IMREAD_UNCHANGED).astype(np.float32) / 1000.0 mask = cv2.imread("data/mask/0001.png", cv2.IMREAD_GRAYSCALE) > 0 K = load_intrinsic("data/meta/0001.txt") pts = depth_to_xyz(depth, K.numpy(), mask) # N x 3，单位米 # 训练和推理的点数必须一致，不够就重复采样，多了就下采样 idx = np.random.choice(len(pts), 8000, replace=len(pts) < 8000) pts = torch.tensor(pts[idx], dtype=torch.float32).cuda().unsqueeze(0) with torch.no_grad(): pred_t, pred_q, conf = model(rgb_tensor, pts) # rgb_tensor 的预处理见仓库 best = conf.argmax() print("平移(米):", pred_t[0, best].cpu().numpy()) print("四元数(wxyz):", pred_q[0, best].cpu().numpy()) if __name__ == "__main__": main()

这里的参数只有几个，但每个都能让你出错：

• num_points必须和训练时一致。这类网络里有个逐像素的拼接操作，点数变了通道维度就对不上，报错信息通常很含糊。
• 深度除以 1000 是因为 16 位 PNG 存的是毫米。如果你的深度图已经是浮点米，这一行去掉，否则平移会小一千倍。
• strict=False是故意的：权重可能不含某些辅助分支。加载后要打印一下missing_keys，如果没有输出为空才叫干净，缺了主干层说明权重和代码不配套。
• 输出的四元数要归一化后再转旋转矩阵，预测值模长不一定等于 1。

看到平移量的量级在 0.3~1.5 米、四元数四个分量都在 [-1,1] 之间，说明链路通了。如果平移是几千或者零点零零几，先回去查单位，不用怀疑模型。

3. 数据管线：内参、标签、裁剪三个约定

在 6D 位姿任务里，模型只是其中一半，另一半是数据。从我的经验看，精度上不去的原因里，坐标约定错了占一半以上。这一章把三个最容易出错的约定单独拆开讲。

3.1 内参与反投影：点云是从深度图算出来的

网络吃的不是原始深度图，而是由深度和内参反投影出来的点云（有的实现是点云加深度图一起喂）。反投影的公式本身很简单，但参数单位错了会很隐蔽。

import numpy as np def depth_to_xyz(depth, K, mask=None, min_z=0.05, max_z=3.0): """ depth: HxW float32，单位米 K: 3x3 相机内参 mask: HxW bool，只保留物体区域 """ h, w = depth.shape fx, fy = K[0, 0], K[1, 1] cx, cy = K[0, 2], K[1, 2] u, v = np.meshgrid(np.arange(w), np.arange(h)) z = depth # 反投影：像素坐标 -> 相机坐标系 x = (u - cx) * z / fx y = (v - cy) * z / fy pts = np.stack([x, y, z], axis=-1).reshape(-1, 3) # 深度为 0 的像素是无效值，进了网络会产生 NaN valid = (pts[:, 2] > min_z) & (pts[:, 2] < max_z) if mask is not None: valid &= mask.reshape(-1) return pts[valid]

逐行说一下为什么这么写。fx, fy分开取是因为有些相机不是方形像素；cx, cy一定用真实值，不要拿w/2, h/2凑，工业相机里偏心几十个像素很常见。min_z/max_z是个保险丝：深度图里的 0 和饱和值都属于无效数据，不滤掉的话，后面的四元数归一化会直接算出 NaN，而这个 NaN 会污染整个 batch 的梯度。mask的作用是裁掉背景，背景点参与回归只会干扰置信度打分。

注意：这套代码里y轴是向下的，和 OpenCV 一致，不是数学教材里向上的。你写可视化的时候如果发现点云上下颠倒，就是这里的问题，别去改模型。

3.2 位姿标签的四种约定，错了不会报错

标签错误最麻烦的地方在于它不报错。loss 照样降，曲线照样好看，就是评估指标上不去。下面四项是我每次接新数据集都一定会确认的。

判断方向对不对有个很土但有效的办法：把物体的 CAD 模型顶点用标签位姿变换到相机坐标系，再和反投影出的点云叠在一起可视化。如果两片点云大致重合，说明方向是对的；如果它们是镜像或者反向的，说明你用成了逆矩阵。这一步我一般会在训练前强制做一遍，比看 loss 曲线靠谱得多。

单位的问题更好查：打印一批标签平移的均值和最大值。抓取场景的工作距离一般在 0.3~1.5 米，如果你的标签均值是 800，那单位就是毫米，写 config 的时候除以 1000。

3.3 接自有数据集：必须实现的两个接口

不管仓库里原本接了几个数据集，你要加自己的数据，本质上就是实现__len__和__getitem__。难点不在写法，在于裁剪时内参也要跟着变。

import torch from torch.utils.data import Dataset class MyPoseDataset(Dataset): def __init__(self, root, obj_names, img_size=224, enlarge=1.2): self.root = root self.obj_names = obj_names self.img_size = img_size self.enlarge = enlarge # bbox 四周留白比例 self.samples = self._scan(root) # 返回 [(rgb路径, depth路径, mask路径, 类别id, R, t)] def __len__(self): return len(self.samples) def __getitem__(self, i): rgb_p, dep_p, msk_p, cls_id, R, t = self.samples[i] rgb, depth, mask, K = self._read(rgb_p, dep_p, msk_p) # 1) 用 mask 的外接框裁一个正方形，并留出余量 ys, xs = mask.nonzero(as_tuple=True) y0, y1 = ys.min().item(), ys.max().item() x0, x1 = xs.min().item(), xs.max().item() cw = int(max(y1 - y0, x1 - x0) * self.enlarge) cx0, cy0 = x0 + (x1 - x0) // 2, y0 + (y1 - y0) // 2 left, top = cx0 - cw // 2, cy0 - cw // 2 rgb = self._crop_pad(rgb, left, top, cw) depth = self._crop_pad(depth, left, top, cw) mask = self._crop_pad(mask, left, top, cw) # 2) 裁剪会平移主点，缩放会改焦距，两步都不能省 K[0, 2] -= left K[1, 2] -= top s = self.img_size / cw K[0, :] *= s K[1, :] *= s rgb = torch.from_numpy(rgb).permute(2, 0, 1).float() / 255.0 return { "rgb": rgb, "depth": depth, "mask": mask, "K": torch.from_numpy(K).float(), "R": torch.from_numpy(R).float(), "t": torch.from_numpy(t).float(), "cls": torch.tensor(cls_id), "sym": torch.tensor(cls_id in SYMMETRIC_IDS), # 对称物体标记，损失函数要用 }

逻辑说明：enlarge控制留白，取值 1.1~1.4 之间比较常见，太小会把物体边缘切掉，太大则背景点占比升高，注意力被稀释。裁剪成正方形是为了后面 resize 不变形，变形的话内参的fx、fy缩放系数就不一样，公式要改。主点平移这一步是最容易漏的，漏了以后的表现是“预测的平移总是偏几十厘米”，而且偏的方向和裁剪位置相关——如果发现误差和物体在图像里的位置有关，八成就是这里。

cls和sym两个字段看着不起眼，但它们是损失函数的输入。对称物体（圆柱、方块这类）如果按普通物体算 ADD，模型会在两个等价姿态之间反复切换，loss 卡在一个固定值下不去。

4. 网络与损失：逐像素融合、残差位姿头、置信度打分

这一章讲的是这套方法区别于“先检测再回归”的地方。传统两阶段做法是：检测框里抠出特征，全局池化成一个向量，再回归位姿。问题是全局池化把物体不同部位的信息混在一堆，无纹理的金属件上更是没什么可混的。DenseFusion 的做法是保留像素级差异，让每个点自己投一票。

4.1 颜色分支和几何分支的逐像素融合

颜色分支是常规的 CNN，把 RGB 提成 H/8 × W/8 × 128 的特征图；几何分支处理点云，输出每个点的 128 维特征。融合发生在“对齐到同一个像素位置”之后——把几何特征贴回图像网格，用相机内参做投影，然后再拼接。

import torch import torch.nn as nn class PixelFusion(nn.Module): """逐像素融合：RGB 特征 + 几何特征 + 全局特征，压回统一维度""" def __init__(self, c_rgb=128, c_geo=128, c_global=128, c_out=128): super().__init__() # 注意是 c_rgb + c_geo + c_global，三份拼一起 self.conv = nn.Sequential( nn.Conv1d(c_rgb + c_geo + c_global, 256, 1), nn.ReLU(inplace=True), nn.Conv1d(256, c_out, 1), nn.ReLU(inplace=True), ) def forward(self, f_rgb, f_geo, f_global): # f_rgb: B x C x N，f_geo: B x C x N，f_global: B x C x 1 g = f_global.expand(-1, -1, f_rgb.shape[-1]) x = torch.cat([f_rgb, f_geo, g], dim=1) return self.conv(x) # B x c_out x N，每个点一份特征

参数上只有三个要留意：c_out决定后面位姿头的宽度，太大显存涨得快，太小表达不够，128 是个稳妥的起点；1x1卷积是为了让每个点独立做通道压缩，不引入空间混叠；f_global那一列是额外拼进去的，它的作用是给每个点提供一点“整体上下文”，避免模型只看到局部。这个expand会复制 N 份，是显存的主要开销之一。

提示：融合时如果几何点的顺序和 RGB 特征图的采样顺序对不上，网络也能训，但精度会掉一截而且不容易发现。我一般会写一个断言，检查两边的点数是否一致，并且用第 3.2 节的可视化方法抽查几个点的位置。

4.2 位姿回归头为什么用残差

如果让网络直接回归绝对平移t，它的输出量级是 1 米左右，而且每个点的输入坐标本身就已经包含了这个量级的信息。残差的做法是：先回归“从这个点到物体中心的偏移量”，再加回该点自身坐标。这样网络只需要学一个小的修正量，数值范围小、收敛快，也更容易泛化到不同工作距离。

class PoseHead(nn.Module): def __init__(self, c_in=128, c_mid=256): super().__init__() self.trunk = nn.Sequential( nn.Conv1d(c_in, c_mid, 1), nn.ReLU(inplace=True), nn.Conv1d(c_mid, c_mid, 1), nn.ReLU(inplace=True), ) self.head_t = nn.Conv1d(c_mid, 3, 1) # 平移残差 self.head_q = nn.Conv1d(c_mid, 4, 1) # 四元数 self.head_c = nn.Conv1d(c_mid, 1, 1) # 置信度 def forward(self, feat, pts): x = self.trunk(feat) delta_t = self.head_t(x).transpose(1, 2) # B x N x 3 t = pts + delta_t # 绝对值 = 点坐标 + 残差 q = torch.nn.functional.normalize(self.head_q(x).transpose(1, 2), dim=-1) c = self.head_c(x).transpose(1, 2) # B x N x 1 return t, q, c

normalize那一行不能删。四元数不归一化的话，转旋转矩阵时会出现缩放，评估误差看着不大但姿态是错的。置信度c后面会作为权重进入损失，推理时用它挑出最好的那一个点，这也是这套方法比“直接回归一个位姿”更稳的原因——它其实是在做一次带权重的集成。

4.3 损失函数：三个权重决定你能不能收敛

损失由两部分组成：位姿损失（平移 + 旋转）和置信度损失。置信度那部分本质上是让网络学会“哪些点的预测更可信”，用预测误差去监督置信度，误差大的点置信度就应该低。

def batch_pose_loss(pred_t, pred_q, pred_c, gt_t, gt_q, pts, sym, w=0.5): """ pred_*: B x N x D，每个点一份预测 gt_*: B x D，整张图一个真值 sym: B，是否对称物体 """ # 1) 平移损失：对称物体用 ADD-S 的最近点近似 if sym.any(): # 对称物体：预测点绕对称轴有多种等价姿态，用最近点距离代替直接相减 dist = torch.norm(pred_t.unsqueeze(2) - gt_t.unsqueeze(1), dim=-1) loss_t = dist.min(dim=-1).values.mean() else: loss_t = torch.norm(pred_t - gt_t.unsqueeze(1), dim=-1).mean() # 2) 旋转损失：四元数取最短弧，q 和 -q 表示同一个旋转 dot = (pred_q * gt_q.unsqueeze(1)).sum(-1).abs().clamp(max=1.0) loss_q = (1 - dot).mean() # 3) 置信度损失：误差越大的点，置信度应该越低 err = (loss_t.detach() + loss_q.detach()) loss_c = (pred_c.squeeze(-1) * err).mean() - w * pred_c.mean() return loss_t + loss_q + loss_c

三个参数是关键。abs()处理的是四元数的双重覆盖，不加的后果是网络在q和-q之间来回跳，loss 曲线呈现规律性震荡。w控制置信度的正则强度，太小网络会把所有点的置信度拉平，太大则会过度自信，取值 0.5 到 2.0 之间调。detach()很重要，置信度分支不应该反过来影响位姿分支的梯度，否则两边互相拉扯，训练会不稳。

对称物体的处理是这套流程里我最建议提前做的一步。抓取场景里圆柱形零件特别多，如果不区分对称和非对称，同一个模型在两类物体上的表现会差出十倍，而且你会以为是网络容量不够。

5. 避坑与排查：六条踩过的记录

下面这些是我在不同数据集上反复遇到的，按现象、原因、解决来写，你可以当作排查清单用。

5.1 数据与坐标类

现象：训练 loss 一路平稳下降，评估时 ADD 一直在几十厘米。

原因：深度图单位是毫米，标签平移单位是米。网络学到的是一个“看起来合理但量纲错”的映射。

解决：在 Dataset 里加一行断言，assert gt_t.max() < 10，超过 10 就说明单位没统一。这件事写进代码里，比记在文档里可靠。

现象：点云可视化和 RGB 对不上，物体整体“偏移”了半个屏幕。

原因：内参用了估算值（拿W/2, H/2当主点），而相机实际主点有偏移。

解决：内参必须从数据集自带的 meta 里读，读不到就做一次标定。主点偏移几十像素在工业相机上很正常，省这一步会让后面所有工作白做。

现象：换了输入分辨率之后，位置全对但姿态全错。

原因：resize 图像时没有同步缩放内参，反投影出来的点云形状被拉伸了。

解决：fx' = fx * s，fy' = fy * s，cx' = (cx - left) * s，三个一起改。只改前两个是最常见的错误。

5.2 训练与显存类

现象：第一个 epoch 没跑完就 CUDA out of memory。

原因：逐像素回归意味着点数就是 batch 维度。一个裁剪块里几万个有效点，中间张量是“点数 × 特征维度 × 层数”。

解决：训练时随机下采样到固定点数（3000~8000），推理时再全量或分批。采样数固定在 config 里，不要每次随机，否则多卡训练时各卡点数不等会卡住。

现象：loss 跑了几百步之后变成 NaN。

原因：深度图里的 0 值点没滤掉，x/z直接除出 inf；或者四元数没归一化，误差项出现负值开根号。

解决：反投影时加min_z过滤，损失里对四元数先归一化再算，clamp(max=1.0)防止1 - dot出现负数。

现象：多卡训练比单卡还慢。

原因：每个 batch 里样本点数不一致，跨卡同步时反复等待。

解决：固定采样点数，让每张卡的张量形状完全一致。这件事在写 Dataset 的时候就该定下来。

5.3 评估类

现象：离线评估 ADD 很好，实际部署抓取成功率和评估对不上。

原因：评估时用的是真值 mask，部署时用的是检测器输出的 mask，两者边缘和漏检情况完全不同。

解决：用真实检测器的 mask 做一次端到端评估。哪怕指标掉几个点，那也是你实际会遇到的数字。我现在的习惯是训练用真值 mask、评估用检测 mask，两套数字都留着。

6. 最后三公里：指标验收、ICP 后处理和吞吐

模型训出来只是起点。真正的差距往往在“怎么判断它行不行”和“怎么让它跑得动”上。

6.1 别盯 loss，用 ADD(-S) 和 5cm5° 验收

loss 是给优化器看的，不是给你看的。我一般只看两个数：ADD(-S) 的平均距离，以及 5cm5° 的通过率。前者衡量模型顶点误差，后者更贴近工程判断——平移误差小于 5 厘米且旋转误差小于 5 度，才算这次抓取是安全的。

import numpy as np def add_s(pred_t, pred_R, gt_t, gt_R, model_pts, sym=False): """ model_pts: M x 3，物体的 CAD 模型顶点 sym=True 时用最近点距离，处理对称物体 """ p = (pred_R @ model_pts.T).T + pred_t g = (gt_R @ model_pts.T).T + gt_t if sym: # 每个预测点到真值点集的最近距离，取平均 d = np.linalg.norm(p[:, None, :] - g[None, :, :], axis=-1) return d.min(axis=1).mean() return np.linalg.norm(p - g, axis=1).mean() def is_5cm5deg(pred_t, pred_R, gt_t, gt_R): dt = np.linalg.norm(pred_t - gt_t) cos = (np.trace(pred_R @ gt_R.T) - 1) / 2 ddeg = np.degrees(np.arccos(np.clip(cos, -1, 1))) return dt < 0.05 and ddeg < 5.0

model_pts一定要用和训练时同一份 CAD 模型，用错了坐标系或者缩放，ADD 会大得离谱而你还以为是模型问题。对称物体的add_s那个双重循环在顶点多的时候很慢，实际跑的时候我一般会先降采样到 500 个点，误差影响很小。

6.2 ICP 后处理和吞吐：精度和速度的交换

预测出来的位姿通常还差一点点，加一步 ICP 精配准往往能把这部分补回来。但这一步的代价很实在，下面是几种常见手段的对比。

我的做法是先固定一个置信度阈值，低于它的点直接丢弃，这一条几乎没有成本；然后再看节拍够不够，不够就减点数，够了再加 ICP。

说个自己的教训。早期我调这类模型，习惯盯着 loss 曲线找问题，一晚上能跑十几组超参，结果精度就是不动。后来发现是深度图单位错了——loss 下降得很漂亮，因为网络学会了一个缩放关系，而 ADD 指标根本没法看。从那以后我每次换数据集，都强制先做一遍可视化：把反投影点云和 CAD 模型顶点按标签位姿叠在一起，用眼睛确认它们重合，再去调任何参数。这个习惯帮我省下的时间，比任何调参技巧都多。

希望帮到你。

本文还有配套的精品资源，点击获取