一、数据读取

1.1 train.py中的数据加载

数据集定义

# 加载训练数据集和测试数据集
if opt.dataset == 'ycb':
    dataset = PoseDataset_ycb('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start)
elif opt.dataset == 'linemod':
    dataset = PoseDataset_linemod('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start)
# dataloader使用opt.workers个进程加速读取数据
dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers)
if opt.dataset == 'ycb':
    test_dataset = PoseDataset_ycb('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start)
elif opt.dataset == 'linemod':
    test_dataset = PoseDataset_linemod('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start)
testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers)

数据读取

1 2	for i, data in enumerate(dataloader, 0): points, choose, img, target, model_points, idx = data

1.2 dataset.py数据读取

1.2.1 初始化

（1）参数输入

mode：模式，共有三种，train、test、eval
num：输入的点云个数（这里linemod默认500个）
add_noise：是否加入噪声（实验中将train加入噪声，test和eval不加入）
root：数据集的根目录
noise_trans：噪声超参数
refine：是否启用refine过程

（2）加载数据集文件

针对数据集中每一个文件夹（物体类别序号），开始读取。

首先根据mode模式参数不同，获取train.txt或者test.txt，即得到哪些图片模型文件是要被读取来进行训练或测试的。这里train读取train.txt，test和eval模式都读取test.txt。每读取一行train.txt或test.txt：

添加RGB图、深度图、mask标签图的路径到三个列表（list_rgb、list_depth、list_label）中，其中mask标签图如果是eval模式，则添加语义分割后的mask标签，如果是train或test模式，就添加标准的mask。
添加类别号list_obj
添加图片编号list_rank
（这里的几个列表按顺序保存了所有物体的txt对应的图片）

然后加载该物体类别的元数据信息meta，包括每一个图像的旋转矩阵、平移矩阵、包围矩形框的坐标、物体类别号。（每个meta[item]保存了一个物体item的元数据）

最后加载该物体的三维点云文件。（每个pt[item]保存了一个物体item的点云）

（3）初始化参数

初始化设置一些参数：

相机内参
xmap，ymap：xmap 数组的第 i 行和第 j 列的值为j，ymap 数组的第 i 行和第 j 列的值为i
图像调整参数trancolor，通过颜色变换可以使图像变得更随机，从而防止模型对数据集过度拟合，作为噪声项根据需要添加。
图像归一化参数，通过归一化可以是模型在训练的时候更快的收敛，并降低对输入数据的敏感度
图像分割列表，将图像分成多块
点云最大数量、最小数量
对陈物体的序号

1.2.2 数据集读取函数

当train.py使用dataloader读取数据时，每迭代一次，调用数据集类的def getitem(self, index)返回一组数据信息。

cloud：深度图+内参计算的点云
choose：用于选取的choose列表
img_masked：物体所在标准包围框的RGB图
target：模型点云齐次变换后的点云
model_points：模型点云
self.objlist.index(obj)：该物体在objlist的索引

（1）加载图片

使用PIL的Image模块加载当前索引批次的RGB图片、深度图、mask标签图、类别号、图片序号信息。并且要保证mask标签图也是BRG三通道图。

根据物体类别号和图片序号，在meta中找到当前数据批次对应的元数据

（2）获取掩码（得到物体所在的标准大小矩形区域）

获取深度图的掩码，根据深度图上每个像素的值是否为0，来构造一个掩码，以确定深度图上的哪些像素值有效，哪些像素值无效。掩码是一个布尔数组，其中的值为True的元素表示相应位置的值有效，值为False的元素表示相应位置的值无效。

获取mask掩码，将实例分割后的结果(eval)，或者标准mask(train/test)中值为255的像素作为掩码保存在mask_label数组中，并转换为1通道。

取两个掩码图的交集作为mask，其中depth为true代表深度值有效，label为true代表此处有物体

根据超参数确定是否为RGB添加噪声（改变亮度、对比度、饱和度、色调），然后提取RGB图的前三个通道，并将图像的维度从(height, width, channel)转换为(channel, height, width)。很多图像处理的库都把通道数放在第1维，这样可以方便地对每一个通道进行操作。

根据元数据中的包围框坐标，或eval时使用mask_to_bbox获取mask图像中物体的包围框，利用get_bbox将包围框转换为标准大小的包围框（标准大小在border_list中定义）。

利用标准大小的包围框获取RGB图像中物体所在的矩形区域。

（3）获得物体上500个随机点的x、y、深度值点云

将物体所在标准包围框区域的mask展开到1维，选取为true的部分的下标，赋值给choose，此时choose就是所有非零元素的下标的数组。

如果choose长为0，则返回6个0张量
如果choose超过500，则随机选取500个作为choose
如果choose介于0-500个，则使用warp模式填充[1,2,3] -> [1,2,3,1,2,3,…]

根据物体所在标准包围框，获得对应区域的深度图，并展开至1维，然后使用choose进行索引，得到choose的深度值，然后转换为列向量。对xmap和ymap进行同样的操作。

利用相机参数，以及choose点的x、y、深度，计算choose的500个点的相机坐标系下的三维坐标，并将x、y、深度三个列向量(500,1)拼接为(500,3)的点云cloud。

根据超参数判断是否为x、y、depth点云添加噪声

（4）读取模型点云和齐次变换后的模型点云

读取该物体的模型点云（从ply文件中读取），在点云中随机删除，最后只留下500个点云。

从元数据中读取真实的旋转矩阵、平移矩阵。利用齐次变换得到当前姿态的点云target。

根据超参数判断是否为模型变换得到的点云target添加噪声

二、网络模型

2.1 模型定义

# estimator为PoseNet网络，即用于预测姿态的主干网络
estimator = PoseNet(num_points = opt.num_points, num_obj = opt.num_objects)
estimator.cuda()
# refiner为PoseRefineNet网络，用于后续迭代自优化
refiner = PoseRefineNet(num_points = opt.num_points, num_obj = opt.num_objects)
refiner.cuda()

2.2 修改ResNet-18网络

定义一个ModifiedResnet网络，将ResNet-18网络模型转换为nn.DataParallel模型，便于在多GPU下训练。

网络输入一个特征向量，返回一个经修改的ResNet-18网络的特征向量。

2.3 PoseNetFeat特征融合网络

输入

x (batchsize,3,500)，由xydepth生成的三维点云
emb (batchsize,di(32),500)，标准包围框区域的RGB图像，经过ModifiedResnet得到的特征图，经过choose选择后的特征
输出
前向计算得到(batchsize, 128+256+1024,500)的向量，64x2通道特征+128x2通道特征+1024通道全局特征

2.4 PoseNet姿态提取网络

输入

img (batchsize, w, h)，标准包围框区域的RGB图像
x (batchsize,500,3)，由xydepth生成的三维点云
choose (batchsize,1,500)，选择数组
obj (batchsize, 1)，物体索引序号
输出
out_rx (1,500,4)，回归得到的四元数
out_tx (1,500,3)，回归得到的平移变量
out_cx (1,500,1)，回归得到的置信度
emb (batchsize,di(32),500)，标准包围框区域的RGB图像，经过ModifiedResnet得到的特征图，经过choose选择后的特征

2.5 PoseRefineNetFeat优化特征融合网络

与PoseNetFeat结构基本一致。

但是PoseNetFeat先将pointfeat_2进行池化，再转换到1024通道得到全局特征，再和pointfeat_1和pointfeat_2拼接。

而PoseRefineNetFeat则是先将pointfeat_1和pointfeat_2拼接作为pointfeat_3，然后进行池化得到融合特征。

输入：

x (batchsize,3,500)，由xydepth生成的三维点云
emb (batchsize,di(32),500)，标准包围框区域的RGB图像，经过ModifiedResnet得到的特征图，经过choose选择后的特征
输出
前向计算得到(batchsize, 1024)的特征向量

2.6 PoseRefineNet姿态优化网络

输入：

x，预测的三维坐标(batchsize,3,500)
emb，标准包围框区域的RGB图像，经过ModifiedResnet得到的特征图，经过choose选择后的特征(batchsize,di(32),500)
obj，物体索引序号(batchsize, 1)
输出：
out_rx(1,4)，回归得到的四元数，预测的三维坐标相对于真实三维坐标的变换
out_tx(1,4)，回归得到的平移变量，预测的三维坐标相对于真实三维坐标的变换

三、损失函数

3.1 Loss初始化

# 对loss进行初始化
# num_points_mesh=500，sym_list对陈物体列表
criterion = Loss(opt.num_points_mesh, opt.sym_list)
criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list)

3.2 Loss计算

输入：

pred_r，(1, 500, 4)，预测得到的四元数
pred_t，(1, 500, 3)，预测得到的平移变量
pred_c，(1, 500, 1)，预测得到的置信度
target，(bs, 500, 3)，根据model_points点云信息，以及标准旋转偏移矩阵转换过的目标点云
model_points，(bs, 500, 3)，模型点云信息
idx，(1)，物体类别序号
points，(bs, 500, 3)，深度三维坐标
opt.w，权重参数
opt.refine_start，是否开始refine过程
输出：
loss，损失值
dis，置信度最大的点的距离
new_points，由points转换到世界坐标系下的点云坐标
new_target，由target转换到世界坐标系下的点云坐标

首先将预测得到的四元数转换为旋转矩阵。

然后将model_points和target都复制500次，用于对每个点进行计算。

利用pred = torch.add(torch.bmm(model_points, base), points + pred_t)公式求出预测值。

model_points是从点云文件中读取的，其坐标是在世界坐标系下定义的，坐标原点为点云文件中固定的零点。
target是model_points在相机坐标系下的坐标，也就是物体在相机坐标系下的点云坐标，也就是物体的实际坐标。
pred_r是世界坐标系和相机坐标系的变换矩阵。
pred_t我也不知道是什么
pred是预测得到的model_points在相机坐标系下的坐标，也就是预测的物体实际坐标。
其中torch.bmm(model_points, base)利用旋转矩阵将model_points姿态转换到相机坐标系下，然后本来加上target_t就是实际坐标，但是这里加上points+pred_t得到预测的坐标。

再对对称物体进行单独的处理。

计算预测的点云pred和真实的点云target之间的距离，求平均值来计算loss，并保存置信度最大的距离loss。

找到置信度最大的点对应的平移向量和旋转矩阵，利用points(相机坐标系点云)逆变换回世界坐标系得到new_points，将target同样逆变换回世界坐标系得到new_target。