基于多头注意力机制和迭代细化的类别级位姿估计方法

专利主权项内容

1.一种基于多头注意力机制和迭代细化的类别级位姿估计方法，其特征在于，包括如下步骤：步骤1、获取现有公开数据集的RGB-D图像；步骤2、构建基于多头注意力机制和迭代细化的类别级位姿估计模型；步骤3、构建整体损失函数来约束估计模型，得到粗粒度的初始位姿估计结果，在此基础上进行不断细化，得到训练完成的类别级位姿估计模型；步骤4、获取当前待预测物体的RGB-D图像，基于已经训练完成的位姿估计模型直接预测当前图像中目标物体的6D位姿；所述步骤1中，公开数据集包括合成数据集CAMERA25和真实数据集REAL275；其中，合成数据集CAMERA25包含300K张由真实背景和虚拟对象合成的RGB-D图像；真实数据集REAL275包含8K张18个不同的真实场景下的RGB-D图像；两个数据集都包含瓶子、碗、相机、罐头、笔记本电脑和马克杯六类物品；RGB-D图像为RGB颜色与深度图像；所述步骤2中，类别级位姿估计模型包括特征编码器模块、坐标变换注意力模块、初始位姿估计模块和位姿迭代精化模块；类别级位姿估计模型的工作过程为：步骤2.1、采用特征编码模块对RGB-D图像进行特征提取；步骤2.2、基于多头注意力机制的坐标变换注意力模块，获取世界坐标系下的RGB图像特征；步骤2.3、基于初始位姿估计模块，预测获得粗粒度的位姿估计结果；步骤2.4、基于位姿迭代精化模块完成对位姿的迭代精化；所述步骤2.1的具体过程为：步骤2.1.1、采用Mask-Rcnn算法对RGB图像中的物体进行实例分割，获得物体的掩码图，并通过金字塔场景解析网络进行特征提取，获得相机坐标系下的RGB图像特征S；Mask-Rcnn算法为一种实例分割算法；0步骤2.1.2、将深度图像映射为三维点云，并通过PointNet++对三维点云进行特征提取，得到相机坐标系下的位置编码和几何特征G；PointNet++是一个点云特征提取与分类网络；0步骤2.1.3、将提取的相机坐标系下的RGB图像特征S、位置编码和几何特征G进行拼接，并利用神经网络对点云进行隐式坐标变换，得到世界坐标系下的点云及其点云特征G；00w所述步骤2.2的具体过程为：步骤2.2.1、将G、G和S分别设为多头注意力的查询Q、键K与值V，使用点积运算的方式计算查询Q与键K之间的相似度，并与值V相乘获得RGB图像在世界坐标系下的RGB图像特征S；多头注意力计算中第i个头的输出公式为：w00w其中，H为第i个头的输出；softmax(·)为softmax归一化操作；T为转置符号；d为头的数量；Q为第i个头的查询，K为第i个头的键，V为第i个头的值，计算公式分别如下：iiii其中，分别为第i个头的查询、键、值可学习的投影矩阵；步骤2.2.2、每个头分别对点云在不同世界坐标系中的特征关联度进行单独计算；将每个头的输出结果进行拼接，最终结果表示为：H＝Concat(H, …, H) (5)；1m其中，H为多头注意力机制最终的输出结果，H为第1个头的输出，H为第m个头的输出，Concat(·)为拼接操作；1m步骤2.2.3、将H传入基于位置的前馈神经网络，获取世界坐标系下的RGB图像特征S；前馈神经网络由多个全连接层和激活函数组成，全连接层进行线性变换，具体计算公式为：wS＝FFN(H) (6)；w其中，FFN(·)为前馈神经网络计算；所述步骤2.3的具体过程为：将相机坐标系下的RGB图像特征S、位置编码和几何特征G以及世界坐标系下的RGB图像特征S进行特征拼接，通过初始位姿解码器进行预测，得到物体粗粒度的初始位姿；将初始位姿解码器与多头位姿残差解码器的参数进行共享，初始位姿解码器与位姿迭代精化模块中多头位姿残差解码器的计算方式相同；00w所述步骤2.4的具体过程如下：步骤2.4.1、将步骤2.3得到的物体粗粒度的初始位姿设置为4×4的坐标变换矩阵，通过矩阵相乘将三维点云从相机坐标系显式坐标变换到世界坐标系下，并使用pointnet++再次提取世界坐标系下的三维点云特征来优化点云的几何特征；步骤2.4.2、将不同坐标系下的RGB图像特征、点云特征和位置编码进行特征拼接，传入多头位姿残差解码器；步骤2.4.3、多头位姿残差解码器部分采用多层感知机与转置操作利用物体的局部几何特征预测旋转残差，采用多层感知机与全局平均池化操作提取全局特征预测位移残差，并将预测的初始位姿的旋转残差和位移残差添加到粗粒度的初始位姿中对其进行优化，将优化后的位姿重新设为初始位姿，循环步骤2.4.3，直到达到设定的最大迭代次数，循环结束，此时得到最终的位姿估计结果；所述步骤2.4.3中，多头位姿残差解码器中采用旋转和位移两个独立的分支输出头分别进行预测；在位移输出头中，使用全连接层将输入特征映射到256维，并使用全局平均池化操作获取全局特征，将其与池化前的特征向量进行拼接，最后预测位移，位移包括位置与大小；位移输出头表示为：f＝MLP([S, G, G]) (7)；transw0wt, s＝MLP([Tile(AvgPool(f)), f]) (8)；Mtranstrans其中，f为位移预测特征向量；t为位置；s为大小；[·]表示特征拼接；MLP(·)为多层感知器；AvgPool(·)为平均池化操作；Tile表示将特征向量复制M次；transM在旋转输出头中，使用全连接层对物体旋转进行预测，旋转输出头表示为：R＝MLP([S, G, G]) (9)；00w其中，R为旋转向量；所述步骤3中，整体损失函数由初始位姿估计损失与迭代位姿估计损失两部分组成；整体损失函数L如下：其中，L为初始位姿估计损失；k为当前迭代次数；n为最大迭代次数；为第k次的特征损失超参数；/>为第k次的特征估计损失；/>为第k次的点损失超参数；/>为第k次的点损失；/>为第k次的位姿估计损失。ini