☰
PyTorch版DenseFusion:RGB-D六自由度位姿估计
2026/10/29 2:27:08 网站建设 项目流程

简介:面向计算机视觉、机器人与AR/VR方向开发者的DenseFusion 6D物体姿态估计工程包,重点解决基于RGB-D图像精确估计物体三轴平移与旋转的问题,适合具备一定深度学习基础、希望复现或改进端到端姿态估计流程的中高级读者。压缩包共55个文件,约3.51MB,以20个Python源码为核心,配合sh下载与运行脚本、m评估绘图文件、cu/h/CUDA扩展、png结果图及md/license/yml等配置说明,覆盖训练、评估和实验组织。目录中可找到网络结构、损失与姿态精化、YCB/LINEMOD数据加载、预训练模型与结果可视化等模块,便于读者梳理像素级RGB-D特征融合、迭代优化和跨数据集评测的实现脉络,并据此搭建自己的训练与验证流程。已有1872人学习,可作为6D姿态估计项目复现与二次开发的参考。

1. 从一张 RGB-D 图到六自由度位姿:机械臂抓取前最难处理的那一步

料框里堆着十几个一模一样的金属件,机械臂要伸进去抓,第一步不是规划路径,而是算清楚目标零件此刻在相机坐标系里绕三个轴各转了多少、中心偏了多少毫米——这就是 6D 物体姿态估计。这份 Python 版 DenseFusion 工程做的正是这一步:喂一张已经对齐的彩色图加深度图,再给一个实例掩码,它吐出旋转矩阵、平移向量,外加一张逐像素的置信度图。它没有采用“先检测框、再回归位姿”的两段式做法,而是把每个像素的颜色特征和它反投影出的三维点特征拼在一起,让上千个像素各出一个位姿,最后挑最靠谱的那个。手里有 RGB-D 相机、想把姿态估计接进抓取或装配流程的开发者是它的主要用户;只做 2D 检测的话,看到点云坐标系那一段会比较吃力。

2. 环境与数据准备:依赖版本、目录结构和 meta 文件的三个对齐点

这套代码属于那种“跑不起来八成不是代码的锅,是数据没对齐”的项目。模型本身只有几百行,真正耗时间的是把相机内参、深度单位、标注坐标系这三样东西捋顺。我一般拿到这类工程,先不急着python train.py,而是写个十几行的小脚本随机抽三帧出来,把点云投回去看一眼叠加是否吻合,吻合了再谈训练。

2.1 依赖版本与显存核算

框架是 PyTorch,Python 3.8 上下最稳。版本不要盲目追新,torch.gather、grid_sample这类算子在老版本上的行为更一致,而且颜色分支要用 torchvision 的 ResNet 预训练权重,新版本里pretrained=True这个参数已经被废弃,直接抄会报错。

conda create -n df6d python=3.8 -y conda activate df6d # 按本机 CUDA 版本选对应的 torch 轮子,不要直接装最新版 pip install torch==1.10.1 torchvision==0.11.2 pip install numpy scipy opencv-python open3d plyfile pyyaml tqdm tensorboardX

open3d只用于可视化和点云采样,plyfile读 CAD 转出来的模型点云,tensorboardX记录 loss 曲线和验证集 ADD,这几个缺了会在 import 阶段就挂掉。

显存方面按经验估:输入 480×640,每帧采 1000 个像素点,batch size 8,光颜色分支的 ResNet18 加逐像素位姿头大概占 6~8G;打开精修网络再叠 1.5G 左右。8G 卡就把 batch 压到 4,再靠梯度累积凑等效 batch,别硬扛。

2.2 数据集目录与自采数据对齐

工程默认吃两类公开数据集的目录约定,一个是单物体多实例的 LineMOD 风格,一个是多物体堆叠的 YCB 风格。目录必须严格按下面的层级来,脚本里的路径是硬编码拼字符串的,少一层就FileNotFoundError:

datasets/linemod/ LINEMOD/ ape/ RGB/0000.png # 彩色图 depth/0000.png # 16 位深度图,单位毫米 mask/0000.png # 该物体实例的二值掩码 meta/0000.pkl # 内参 + 物体到相机的 4x4 位姿 benchvise/ ... datasets/ycb/ YCB_Video_Dataset/ data/0000/0000-color.png data/0000/0000-depth.png data/0000/0000-meta.pkl models/ # 每个物体的模型点云,ply 格式

自采数据往这个格式上靠,最容易出问题的是 meta。它至少要有三样东西:3×3 内参、物体坐标系到相机坐标系的 4×4 变换矩阵、深度缩放系数。我一般写个小转换脚本,从标注工具的 json 直接生成:

import numpy as np, cv2 def make_meta(depth_path, K, T_obj2cam, out_path, depth_scale=1000.0): """把自采的一帧转成工程能读的 meta K: 3x3 内参;T_obj2cam: 4x4,物体坐标系 -> 相机坐标系 depth_scale: 深度图原始单位换算到米的除数,毫米图填 1000 """ depth = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED).astype(np.float32) depth = depth / depth_scale # 工程内部统一按米算,别混着来 meta = { 'intrinsics': {'fx': float(K[0, 0]), 'fy': float(K[1, 1]), 'cx': float(K[0, 2]), 'cy': float(K[1, 2])}, 'poses': T_obj2cam.tolist(), # 注意方向:物体 -> 相机,反了整个训练就废了 } np.save(out_path, meta, allow_pickle=True)

T_obj2cam的方向是第一个大坑。标注工具经常给的是相机到物体,取逆之后平移分量会完全不对,loss 却照样能降——因为网络在学一个“看起来合理”的错位姿,最后评测 ADD 死活上不去。深度图还有一条必须先做:align_depth_to_color,深度和彩色没做硬件对齐的话,反投影出来的点云和像素位置错开几个像素,位姿整体就会偏。

2.3 配置文件里的关键参数

参数集中在一个 yaml 里,改之前先搞清楚哪几个是牵一发动全身的:

dataset: linemod n_classes: 13 # 数据集里物体实例总数,换数据集必须改 batch_size: 8 n_epochs: 30 lr: 1e-4 # 主网络学习率 lr_refine: 1e-4 # 精修网络单独一个学习率 n_sample_points: 1000 # 每个实例采多少像素参与 loss,直接决定显存 w_conf: 0.01 # 置信度正则项权重,调这个比调学习率有用 bbox_expand: 1.2 # ROI 外扩比例,太小会切掉物体边缘 use_refine: true refine_epoch: 25 # 前 25 轮只训主网络,之后才开精修

n_classes不是分类数,是位姿头索引用的物体 id 数量,加物体时忘了改会在取 one-hot 时维度不匹配。n_sample_points是性价比最高的一个旋钮,从 1000 降到 500,显存差不多减半,精度掉零点几个点,粗调阶段完全可以先用 500 快速验证数据通不通。refine_epoch的意思是精修网络必须在主网络收敛到一定程度之后再开,两边一起训的话精修会去拟合噪声,越修越歪。

3. 网络结构拆解:颜色分支、点云分支和逐像素稠密融合的三段链路

整个前向过程可以拆成“两路编码 → 逐像素拼接 → 位姿头 + 置信度头 → 精修”四步。理解这张图的关键在于:它不是对整张图出一个位姿,而是对 roi 里的每个点各出一个位姿,最后用置信度做加权选择。这个设计的好处是对遮挡天然鲁棒——被挡住一半的物体,剩下的可见像素照样能投票出正确位姿。

3.1 颜色分支与点云分支各自的输入输出

颜色分支用 ResNet18 的骨干,但砍掉最后的全局池化和全连接,只保留到 layer4,再用 1×1 卷积把 512 通道压到 128 通道。之所以要保留空间分辨率,是因为后面要按像素索引取特征,一旦做了全局池化,空间信息就没了。

import torch, torch.nn as nn, torchvision class ColorEncoder(nn.Module): """RGB 分支:480x640 彩图 -> (B, 128, 60, 80) 的逐像素特征图""" def __init__(self, out_dim=128): super().__init__() resnet = torchvision.models.resnet18(pretrained=True) # conv1 ~ layer3,输出 stride=16 self.base = nn.Sequential(*list(resnet.children())[:7]) self.layer4 = nn.Sequential(*list(resnet.children())[7]) # 用 1x1 卷积压通道,不改变空间尺寸 self.c2 = nn.Sequential( nn.Conv2d(512, 256, 1), nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.Conv2d(256, out_dim, 1), nn.BatchNorm2d(out_dim), nn.ReLU(inplace=True), ) def forward(self, x): return self.c2(self.layer4(self.base(x))) # (B, 128, 60, 80)

点云分支就轻得多,本质是个 PointNet 的简化版:每个点只有 xyz 三个通道,用一串 1×1 卷积(等价于对每个点独立做 MLP)升到 128 维,不做任何池化,因为每个点都要保留自己的特征。

class DepthEncoder(nn.Module): """点云分支:(B, N, 3) 的相机坐标系点 -> (B, N, 128) 的逐点特征""" def __init__(self, out_dim=128): super().__init__() self.mlp = nn.Sequential( nn.Conv1d(3, 64, 1), nn.BatchNorm1d(64), nn.ReLU(inplace=True), nn.Conv1d(64, 128, 1), nn.BatchNorm1d(128), nn.ReLU(inplace=True), nn.Conv1d(128, out_dim, 1), ) def forward(self, pts): # Conv1d 要求 (B, C, N),所以先转置 return self.mlp(pts.transpose(1, 2)).transpose(1, 2) # (B, N, 128)

两个分支的输入必须严格对应:第 i 个点在图像上的像素位置,就是颜色特征图上要取的那个位置。这个对应关系由 dataset 里的choose字段提供,是像素坐标除以 8 之后的结果(因为特征图是 1/8 分辨率)。这里一旦出现 off-by-one,比如用了x // 8却按(x - 1) // 8算,特征就整体错开一列,位姿会稳定地偏一点点,非常难查。

3.2 稠密融合层:把逐像素特征拼起来

融合这一步是整个方法的核心,说穿了很简单:把颜色特征图按索引 gather 成“每个点对应的颜色特征”,然后和点特征在通道维拼接,再过几层 1×1 卷积融合。

class DenseFusion(nn.Module): def __init__(self, out_dim=128): super().__init__() self.color = ColorEncoder() self.depth = DepthEncoder() self.fuse = nn.Sequential( nn.Conv1d(256, 256, 1), nn.BatchNorm1d(256), nn.ReLU(inplace=True), nn.Conv1d(256, out_dim, 1), ) def forward(self, img, pts, choose): """img: (B,3,H,W) pts: (B,3,N) choose: (B,N,2) 每个点对应的 (h, w)""" b, c, h, w = img.shape[1], 128, img.shape[2] // 8, img.shape[3] // 8 cfeat = self.color(img).view(b, c, h * w) # (B,128,H*W) # 把 (h,w) 压成一维下标,再按通道复制,供 gather 使用 idx = (choose[..., 0] * w + choose[..., 1]) # (B,N) idx = idx.unsqueeze(1).repeat(1, c, 1) # (B,128,N) cfeat = torch.gather(cfeat, 2, idx) # (B,128,N) dfeat = self.depth(pts).transpose(1, 2) # (B,128,N) fused = torch.cat([cfeat, dfeat], dim=1) # (B,256,N) return self.fuse(fused) # (B,128,N)

torch.gather那两行是性能热点,也是内存大户。idx的构造别写成循环,一帧 1000 个点循环 1000 次会让 dataloader 成为瓶颈。还有一种更省内存的写法是先把颜色特征裁到 roi 范围内,用相对坐标索引,特征图能小一个数量级,代价是要多维护一套坐标偏移。

融合之后每个点拿到 128 维特征,这个特征同时包含了“这块像素长什么样”和“这个点在三维空间里在哪”,比单独用任何一路都强。

3.3 位姿回归头与置信度分支

位姿头和置信度头共用一个 1×1 卷积塔,最后输出 8 个通道:前 4 个是四元数,中间 3 个是平移,最后 1 个是置信度。

class PoseHead(nn.Module): def __init__(self, in_dim=128): super().__init__() self.conv = nn.Sequential( nn.Conv1d(in_dim, 128, 1), nn.ReLU(inplace=True), nn.Conv1d(128, 128, 1), nn.ReLU(inplace=True), nn.Conv1d(128, 8, 1), # 4 + 3 + 1 ) def forward(self, feat): out = self.conv(feat) # (B, 8, N) quat = out[:, :4].transpose(1, 2) # (B, N, 4) # 四元数必须归一化,否则转旋转矩阵时会出现尺度漂移 quat = quat / (quat.norm(dim=2, keepdim=True) + 1e-8) trans = out[:, 4:7].transpose(1, 2) # (B, N, 3) conf = torch.sigmoid(out[:, 7]).transpose(1, 2) # (B, N, 1),约束到 0~1 return quat, trans, conf

sigmoid那一步容易漏。置信度参与的是-log(c)这种正则项,如果不约束到 (0,1) 区间,log 直接出 NaN,训练第一个 epoch 就崩。归一化也不能省,四元数长度不为 1 时用公式转旋转矩阵得到的不是正交矩阵,变换后的点云会被缩放。

3.4 迭代精修网络 RefineNet

初次预测的位姿精度通常在 ADD 阈值边缘,精修网络负责把结果再推一把。它的输入是点云坐标加上每个点的当前位姿估计拼成 10 维,输出是一个残差位姿,迭代 2~3 次。

class RefineNet(nn.Module): def __init__(self, n_pts=1000): super().__init__() # 输入 3(xyz) + 4(quat) + 3(trans) = 10 维 self.mlp = nn.Sequential( nn.Conv1d(10, 128, 1), nn.BatchNorm1d(128), nn.ReLU(inplace=True), nn.Conv1d(128, 256, 1), nn.BatchNorm1d(256), nn.ReLU(inplace=True), nn.Conv1d(256, 512, 1), nn.BatchNorm1d(512), nn.ReLU(inplace=True), ) # 全局池化后回归一个 7 维残差 self.head = nn.Sequential(nn.Linear(512, 256), nn.ReLU(inplace=True), nn.Linear(256, 7)) def forward(self, pts, quat, trans): feat = torch.cat([pts.transpose(1, 2), quat.transpose(1, 2), trans.transpose(1, 2)], dim=1) # (B,10,N) feat = self.mlp(feat).max(dim=2).values # (B,512) 全局特征 delta = self.head(feat) # (B,7) 残差 return delta[:, :4], delta[:, 4:7]

全局池化意味着精修网络只输出一个位姿,不管前面有多少个像素提了意见。所以在进精修之前,得先用置信度把逐像素的预测聚合成一个初始值——常见做法是取置信度最高的那个像素,或者用置信度加权平均后归一化。精修网络训练时输入的初始位姿必须来自主网络的真实输出并叠加噪声,如果直接用真值当输入,训练时 loss 掉得飞快,一上测试集就崩。

4. 训练与评测:位姿损失、置信度自监督和 ADD 指标的算法细节

训练这块最值得抄的是置信度那套自监督机制。传统做法是给每个像素的损失等权重求平均,但被遮挡的像素本来就估不准,硬拉低它们的损失反而会把整体带偏。DenseFusion 让网络自己学一个置信度,高置信的像素损失权重大,低置信的允许它躺平。

4.1 位姿损失怎么算才不会爆炸

因为预测点云和真值点云用的是同一组模型点,只是变换矩阵不同,所以不需要做最近邻搜索,直接逐点求距离就行,这是这个实现比 ICP 类方法快很多的原因。

def pose_loss(pred_q, pred_t, gt_R, gt_t, model_pts, sym=False): """pred_q: (B,N,4) pred_t: (B,N,3) model_pts: (M,3) 返回逐像素损失 (B,N) """ pred_R = quat2mat(pred_q) # (B,N,3,3) mp = model_pts.view(1, 1, -1, 3) # (1,1,M,3) # 两组变换作用在同一批模型点上 pred_pts = torch.einsum('bnij,bmj->bnmi', pred_R, mp) + pred_t.unsqueeze(2) gt_pts = torch.einsum('ij,mj->mi', gt_R, mp[0, 0]) + gt_t # (M,3) gt_pts = gt_pts.view(1, 1, -1, 3) dis = torch.norm(pred_pts - gt_pts, dim=-1) # (B,N,M) if sym: # 对称物体:把模型点绕某个轴转 180 度,取两组里的最小值 mp_flip = torch.stack([-mp[..., 0], -mp[..., 1], mp[..., 2]], dim=-1) pred_flip = torch.einsum('bnij,bmj->bnmi', pred_R, mp_flip) + pred_t.unsqueeze(2) dis = torch.min(dis, torch.norm(pred_flip - gt_pts, dim=-1)) return dis.mean(dim=-1) # (B,N)

sym这个分支非常关键。杯子、螺丝、垫片这类旋转对称物体,同一个姿态换个角度转 180 度看起来一模一样,如果按非对称算损失,网络会被互相矛盾的监督信号撕扯,最后收敛到一个中间态,谁都对不上。哪些物体算对称得手工维护一张表,这个表一般放在数据集类里,加新物体时别忘了填。

dis的形状是(B, N, M),1000 个点 × 1000 个模型点 = 一百万个浮点数,再乘 batch 就是显存炸弹。常见的省法是模型点只采样 500 个,或者损失不回传到全部像素、只对置信度高的前 K 个点求梯度。

4.2 置信度损失和它为什么能收敛

置信度损失的形式是L_i * c_i - w * log(c_i),前一项让网络倾向于把置信度压低来减小损失,后一项惩罚置信度太小,两项拉扯出来的平衡点就是“这个像素的位姿有多靠谱”。

def confidence_loss(per_pixel_loss, conf, w=0.01): """per_pixel_loss: (B,N) conf: (B,N,1)""" conf = conf.squeeze(-1) return (per_pixel_loss * conf - w * torch.log(conf + 1e-8)).mean()

w是唯一的超参,我一般从 0.01 起调。w太大会让所有像素的置信度都趋近 1,退化回等权平均;太小则所有置信度趋近 0,正则项直接主导,主损失根本没法学。判断方法很直接:训练几个 epoch 后把 conf 的直方图打出来,如果全挤在 0.99 以上或者 0.01 以下,就是这个参数不对。

+1e-8那个 eps 别省,sigmoid 在极端输入下会输出精确的 0 或 1,log(0) 直接给你一个 inf,然后整个 batch 的梯度全废。

4.3 训练循环里的采样与数据增强

dataset 的__getitem__是做实验时改动最频繁的地方,核心逻辑其实就三步:取掩码、采点、反投影。

def __getitem__(self, idx): item = self.items[idx] # 一条样本 = 一帧 + 一个物体实例 color = read_rgb(item['color_path']) # (3,H,W),已归一化 depth = read_depth(item['depth_path']) # (H,W),单位米 mask = read_mask(item['mask_path']) # (H,W),0/1 K, T = item['K'], item['RT'] # 内参、物体到相机的 4x4 ys, xs = np.nonzero(mask) if len(xs) > self.n_points: # 掩码里的点数远超需求,随机下采样 sel = np.random.choice(len(xs), self.n_points, replace=False) xs, ys = xs[sel], ys[sel] zs = depth[ys, xs] keep = zs > 0 # 深度空洞直接丢,别用 0 或均值补 xs, ys, zs = xs[keep], ys[keep], zs[keep] # 反投影:像素坐标 + 深度 -> 相机坐标系三维点 x3 = (xs - K[0, 2]) * zs / K[0, 0] y3 = (ys - K[1, 2]) * zs / K[1, 1] pts = np.stack([x3, y3, zs], axis=1).astype(np.float32) # (N,3) # 对应到 1/8 分辨率特征图上的下标 choose = np.stack([ys // 8, xs // 8], axis=1).astype(np.int64) return {'color': color, 'pts': pts, 'choose': choose, 'gt_R': T[:3, :3].astype(np.float32), 'gt_t': T[:3, 3].astype(np.float32)}

keep = zs > 0这一行是有实际影响的:深度图里的空洞点如果补零,反投影出来就是相机原点附近的一堆点,等于往训练集里掺了噪声标签,网络会学出一个朝相机方向塌缩的偏移。至于增强,常用的就三样:随机换背景(从别的图里抠一块贴上去)、颜色抖动、深度上加高斯噪声。第三种最贴近真实场景,因为深度相机的噪声随距离变化,远距离误差能到几厘米。

训练循环本身没太多花样,注意两点:一是主网络和精修网络的优化器要分开,学习率各管各的;二是每个 epoch 结束存一次 checkpoint,这个项目的训练动辄十几个小时,没有断点续训的话一次断电就是灾难。

4.4 推理链路与 ADD、ADD-S 指标

推理时先用检测或分割模型出实例掩码,再对每个掩码走一遍前向,得到 N 个候选位姿和置信度,聚合后送精修网络,最终输出一个位姿。

评测用的 ADD 指标是衡量“变换后的模型点平均偏了多少”,对称物体换用 ADD-S(最近点距离):

def add_metric(pred_R, pred_t, gt_R, gt_t, model_pts, sym=False): """返回 ADD / ADD-S,单位与点云一致(米)""" p = model_pts @ pred_R.T + pred_t # (M,3) g = model_pts @ gt_R.T + gt_t if sym: d = torch.cdist(p, g).min(dim=1).values # 每个预测点找最近的 gt 点 else: d = torch.norm(p - g, dim=1) # 同下标直接算距离 return d.mean().item()

判定标准一般是ADD < 0.1 * 模型直径算预测正确。这里的“模型直径”得用物体的最大外接尺寸算,如果拿包围盒对角线凑合,阈值会虚高。还有一点:对称物体一定要切到 ADD-S,我用 ADD 评过一个圆柱类零件,分数只有 30% 出头,换成 ADD-S 直接到 90%,纯粹是评价方式的问题,不是模型不行。

5. 避坑与排查:跑崩、指标不动、置信度塌缩的五类现场

下面这几条都是这类工程里反复出现的“血泪经验”,每条按现象、原因、解决顺序写,对着查能省不少时间。

5.1 跑不起来的那一类

现象:训练第一个 epoch 中期,loss 突然变成nan,之后再也回不来。

原因:四元数没有归一化,转旋转矩阵时出现数值爆炸;或者置信度损失里没有加 eps,log(0)直接产出inf。

解决:在位姿头里强制归一化(q / q.norm()),给所有 log 加1e-8,再开torch.autograd.set_detect_anomaly(True)跑一个 batch,会直接指出是哪一层的反传出问题。

现象:CUDA out of memory,而且把 batch 降到 1 也没救。

原因:位姿损失的中间张量是(B, N, M)量级,batch 降下来只解决了 B 这一维,N 和 M 还是 1000 × 1000。

解决:把模型点从 1000 采到 500,或者把损失函数里对像素维求均值改成先算(B,N)的逐点损失再 mean,避免在中途保留三维张量。实在不够就用梯度累积,batch 4 累积 4 次等效 batch 16。

现象:KeyError或者 shape mismatch,提示维度是 14 而不是 13。

原因:n_classes只跟数据集大小有关,加了新物体忘了改这个值,或者把背景类也算进去了。

解决:把n_classes直接改成len(物体列表),别手写数字。

5.2 能跑但指标不动的那一类

现象:训练 loss 稳定下降到 0.02,验证集 ADD 一直在 40% 上下不动,怎么看都像欠拟合。

原因:十有八九是 meta 里的位姿方向反了,或者是深度单位没统一(图像是毫米,模型点云是米)。网络能学会在这个错位空间里做回归,但预测出来和真实世界差一个常数变换。

解决:抽三帧出来,用真值位姿把模型点云变换过去,投回图像看是否和彩色图上的物体重合。这一步能查出 90% 的坐标系问题。

现象:置信度图全是一片高亮,或者全是一片黑。

原因:w_conf调偏了,或者主损失和正则项量级差太多。

解决:打印conf的均值,健康值大概在 0.3~0.7 之间呈双峰分布。全高就把w_conf调小,全低就调大。

现象:开了精修之后指标反而下降。

原因:精修网络训练时喂的是真值加噪声,而推理时喂的是主网络的预测,两者分布不一致。或者refine_epoch设得太早,主网络还没收敛就开了精修。

解决:精修必须先等主网络训练结束;喂入必须用主网络实际的输出,噪声幅度按主网络的实测误差来定,别随便设个 0.01。

5.3 部署阶段的那一类

现象:模型导出 ONNX 后结果和 PyTorch 对不上。

原因:torch.gather的动态索引在部分 opset 下会退化,四元数归一化里那个+1e-8也可能被算子融合吃掉。

解决:导出时把 opset 固定到 11 以上,导完用同一帧输入对比两边输出,误差应该在 1e-5 量级;如果 gather 出问题,就把它改成按 roi 裁剪后再索引的写法,兼容性更好。

现象:单帧推理 80ms,达不到产线节拍。

原因:耗时大头不在网络,在数据预处理——反投影、点云下采样、掩码处理都是 CPU 单线程。

解决:把预处理搬到 GPU 上用张量算,或者提前把点云缓存下来。网络侧把n_sample_points从 1000 降到 300 对精度影响很小,但速度快一倍多。

6. 二次开发:换成自己的物体、自己的相机和自己的验证闭环

真要落地,第一步是换物体。这套工程加一个新物体要改四处:数据集里的物体 id 映射表、n_classes、对称性配置表、以及models/下那个 ply 点云模型。点云从 CAD 导,注意单位要跟深度图一致,CAD 里是毫米就除 1000,忘了这一步,训练出来平移分量会差三个数量级,loss 曲线却看着挺正常。

标注量上别一上来就冲几千帧。我的经验是每个物体先标 150~200 帧,覆盖不同的视角和光照,跑一轮看 ADD 能不能过 70%;过了再补数据,不过就先查坐标系。合成数据是个性价比很高的补充手段:用 CAD 渲染一批图,随机化背景、光照、相机位姿,再叠加真实深度图的噪声模型,跟真实标注按 3:1 混着训,泛化性能提升明显。

# 新增物体的三处改动,漏一处都会静默出错 OBJ_ID = {'part_a': 0, 'part_b': 1, 'part_c': 2} # 1. id 映射 N_CLASSES = len(OBJ_ID) # 2. 跟着映射走,别手写 SYMMETRIC = {'part_a': False, 'part_b': True, 'part_c': False} # 3. 对称性

换相机则是另一条线。内参必须重新标定,而且要注意深度图和彩色图的分辨率是否一致——很多相机深度是 640×480,彩色是 1280×720,直接把彩色图 resize 到深度尺寸会改变内参里的 fx、fy 和主点,必须按同一比例缩放后再填进 meta。上机械臂的话还要做手眼标定,得到相机到基座的变换,这一步的误差会直接叠加到最终抓取位姿上,标定板拍少了误差能到厘米级。

验证闭环我一般这样搭:留 20 帧完全不参与训练的数据,跑推理后把预测位姿变换后的模型点云和真值点云同时可视化,看两者是否贴合。ADD 数字只是参考,肉眼看贴合度才能发现“整体偏一个方向”这种系统性问题。精修前后各存一份对比图,如果精修后反而更歪,说明精修网络的训练数据分布不对,回去查refine_epoch。

有个细节我踩

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

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询