【6PACK代码注解】数据集载入与预处理

momo_vv

已于 2022-08-06 13:14:26 修改

阅读量908

点赞数 1

分类专栏： 6pack 文章标签：计算机视觉人工智能深度学习

于 2022-08-04 10:43:56 首次发布

本文链接：https://blog.csdn.net/weixin_44695308/article/details/125936072

版权

6pack 专栏收录该内容

9 篇文章 2 订阅

订阅专栏

文章目录

前言
一、生成真值数据dataset_preprocess.py
二、数据载入dataset_nocs.py

前言

【6PACK全记录】6-PACK论文学习及复现记录

一、生成真值数据dataset_preprocess.py

1.1 已知量：

model_scales：每个文件为某个图片中对应实例的边框范围，对训练集体现为2个对角顶点的坐标(已归一化)：

对验证集，体现为8个顶点的坐标：
data_pose：每个文件为对应图片中每个实例的真实位姿，每个实例id后有4行，其中红框是3*3的R矩阵，蓝色是T
My_NOCS/data/train中color.png、depth.png、mask.png分别为彩色图、深度图和带标号的蒙版（每个黑色区域像素的取值不为[0,0,0]而是对应其它meta中的实例id，例如[1,1,1]、[2,2,2]）
meta.txt：每个文件存储对应图片的相关内容，出去最前列Linux额外添加的序号，第一列为实例id，第二列为其所述的类别id，最后一列为该实例对应的model_scales文件名

1.2 目标

生成bbox.txt文件：保存对应图片中目标实例的8个边框顶点的坐标（世界坐标系）
在这里插入图片描述

1.3 代码逻辑

遍历图片掩码（蒙版）中所有像素值，对应了图片中实例的id和背景亮度255，在meta.txt中找出对应实例的scale文件名（这一步即过滤掉了meta中无关的背景和可能存在的多余实例）;通过读取初始位姿R/T和相机系点云获得物体在世界坐标系下的点云。再通过scale找到包围点云的最小边界框，写入bbox.txt。
详细代码及注释见dataset_preprocess.py文件

二、数据载入dataset_nocs.py

为处理自定义的数据集，继承torch.utils.data.Dataset基类来定义Dataset类，该类对数据进行包装，使得DataLoader类操作更便捷。在继承后，重写

_len_()函数---返回数据集样本数;
_getitem_()---使得dataset支持用索引取出特定样本，即dataset[i]会调用getitem(i)来取出第i个样本。

2.1 初始化init()

相关参数及含义：

mode----train/val
root----数据集路径
add_noise----是否加随机噪声
num_pt------点云的点数
num_cates-----类别数=6
count-----数据集所取的样本数
cate_id----当前训练的类别

训练集（验证集与之基本一致）
生成两个字典：

self.obj_name_list：从路径My_NOCS/data_list获得，是各类别实例的相关信息，每个索引下存其子文件夹的名称(实际是一个实例)，每个子文件中有一个list.txt文件。
list中为每个实例出现在不同的图片的地址：
(训练过程中取这个实例在不同图中的位姿与一个虚拟背景结合，看作随时间变化，实例的位姿再变化)

其中print(tmp_cate_id, item) 的输出结果如下，为类别id+类别中每个子文件夹名称：
在这里插入图片描述

self.obj_list：二维字典，self.obj_list[类别id][实例名]中存储list.txt的内容，为每个实例出现在不同图片的绝对地址

#以下两个字典都是从1开始索引（类别1：6）
self.obj_list = {}#{1:{'文件夹1':[],'文件夹2':[]}} 2维字典
self.obj_name_list = {}
#各类别中所有实例信息，每个文件夹中有1个list文件，字典{1:['文件夹1','文件夹2'],2:....}

if self.mode == 'train':
	for tmp_cate_id in range(1, self.num_cates+1):
		self.obj_name_list[tmp_cate_id] = os.listdir('{0}/data_list/train/{1}/'.format(self.root, tmp_cate_id))
        self.obj_list[tmp_cate_id] = {}

        for item in self.obj_name_list[tmp_cate_id]:
            print(tmp_cate_id, item)  #输出每个类别的子文件夹名称
            self.obj_list[tmp_cate_id][item] = []

            input_file = open('{0}/data_list/train/{1}/{2}/list.txt'.format(self.root, tmp_cate_id, item), 'r')
            while 1:
                input_line = input_file.readline()
                if not input_line:
                      break
                if input_line[-1:] == '\n':
                      input_line = input_line[:-1]#不读'\n'
                self.obj_list[tmp_cate_id][item].append('{0}/data/{1}'.format(self.root, input_line))
            input_file.close()

读取dataset/sphere.xyz文件中的坐标，dataset/sphere.xyz文件内容如下：
在这里插入图片描述
其他：

self.back_list---为了数据增强采用的背景（train2017提供数据）
self.cam_cx/y、self.cam_fx/y---相机内参
self.xmap、ymap---2D图像网格
self.trancolor----图像色彩修正，亮度、对比度、饱和度、色调
self.norm----对数据归一化，image=(image-mean)/std

2.2 getitem

以下简述相关函数及其作用，函数具体实现可见代码文件中的注释

函数名	作用
divide_scale（scale, pts）	输入点的坐标pts，归一化范围scale；输出点的归一化坐标，使得各维度距离中心最远点的坐标为1或-1
get_anchor_box（ori_bbox）	生成锚点网格，输入为实例的边界框，输出ans, scale，ans为每个维度有5个锚点的归一化锚点网格坐标，scale为他的归一化参数
change_to_scale（scale, cloud_fr, cloud_to）	根据scale将点云cloud_fr, cloud_to归一化再输出
enlarge_bbox(target)	将边框放大；输入为实例边框的2个对角顶点或者8个顶点，输出为放大后的边框
load_depth(depth_path)	加载目标深度信息，注意我们的深度图中，背景是蓝色通道，目标深度用G+R
re_scale(target_fr, target_to)	输出target_fr和ans_scale，用于将fr、to范围统一，将target_to的相关数据*ans_scale即得到与target_fr范围一致的数据
search_fit(points)	找到points各个维度的最值，例如xmax、xmin
get_2dbbox(cloud, cam_cx, cam_cy, cam_fx, cam_fy, cam_scale)	获取像素平面坐标系(r，c)下图像中实例的2维边框；输入中cloud是真实世界坐标系(x,y,z)下的点云坐标，其余为相机内参，变换关系见6D位姿估计学习
get_pose( choose_frame, choose_obj)	获取给定实例的真实位姿ans_r, ans_t以及该实例的索引ans_idx；输入分别为目标实例所在图片的路径和目标实例名称
get_frame( choose_frame, choose_obj, syn_or_real)	获取画面帧choose_frame的相关信息。输出内容：img：color_crop，即图片(可能有虚拟背景)中处于2维边框的部分；choose：所选的点云对应在2D图中的像素索引；cloud：相机坐标系中的点云；r、t：当前位姿；target：世界坐标系3维边框边界；mesh_pts：数据集自带的3维边框的网格点；mesh_bbox：数据集自带的3维边框；mask_target：实例的蒙版，并只取处于2维边框内的部分

gettiem(index)

在train.py中执行：

for i, data in enumerate(dataloader, 0)

代码段时会调用，相当于data[index]，即取出给定索引idx的数据。

随机参数syn_or_real决定是采用真实世界的实例图还是建模的实例图，对于建模图，将在get_frame函数中随机决定是否给该实例更换train2017的虚拟背景来数据增强。

def __getitem__(self, index):
        #随机决定是否采用虚拟背景来数据增强（只针对训练集）
        syn_or_real = (random.randint(1, 20) < 15)
        if self.mode == 'val':
            syn_or_real = False

对于选定的建模实例，随机找两张含有这个实例的图片choose_frame[0]、choose_frame[1]，把这两张图看作视频中连续的两个画面帧。通过get_frame()函数得到前一帧的相关数据“_fr”（见上述函数表格）和后一帧数据“_to”；
并通过re_scale()函数将两帧的点云、框架数据调整在一个范围中。
真实实例类似。

if syn_or_real:
            while 1:
                try:
                    #随机取该类别中的一个建模实例
                    choose_obj = random.sample(self.obj_name_list[self.cate_id], 1)[0]
                    #取出包含该实例的任意2个图片的地址，把这两个看成位姿随时间的变化
                    choose_frame = random.sample(self.obj_list[self.cate_id][choose_obj], 2)
                    #_fr：位姿变化前（上一时刻）
                    img_fr, choose_fr, cloud_fr, r_fr, t_fr, target_fr, mesh_pts_fr, mesh_bbox_fr, mask_target = self.get_frame(choose_frame[0], choose_obj, syn_or_real)
                    if np.max(abs(target_fr)) > 1.0:#正常情况下边框已做归一化处理，最远边界为±1
                        continue
                    #_to：位姿变化前（下一时刻）
                    img_to, choose_to, cloud_to, r_to, t_to, target_to, _, _, _, = self.get_frame(choose_frame[1], choose_obj, syn_or_real)
                    if np.max(abs(target_to)) > 1.0:
                        continue

                    #将_to的点云、位姿调整为与_fr同范围
                    target, scale_factor = self.re_scale(target_fr, target_to)
                    target_mesh_fr, scale_factor_mesh_fr = self.re_scale(target_fr, mesh_bbox_fr)

                    cloud_to = cloud_to * scale_factor
                    mesh = mesh_pts_fr * scale_factor_mesh_fr#是边框网格中的点？
                    t_to = t_to * scale_factor
                    break
                except:
                    continue

        else:
            while 1:
                try:
                    #real_obj_name_list中是真实的实例，且不需要添加背景
                    choose_obj = random.sample(self.real_obj_name_list[self.cate_id], 1)[0]
                    choose_frame = random.sample(self.real_obj_list[self.cate_id][choose_obj], 2)

                    img_fr, choose_fr, cloud_fr, r_fr, t_fr, target, _ = self.get_frame(choose_frame[0], choose_obj, syn_or_real)
                    img_to, choose_to, cloud_to, r_to, t_to, target, _ = self.get_frame(choose_frame[1], choose_obj, syn_or_real)
                    if np.max(abs(target)) > 1.0:
                        continue
                    break
                except:
                    continue

生成锚点网格anchor_box并将点云按照网格归一化

class_gt = np.array([self.cate_id-1])
#根据边框生成锚点网格，每维度5个点，scale是归一化范围（限制在1内）
anchor_box, scale = self.get_anchor_box(target)
#归一化两帧的点云
cloud_fr, cloud_to = self.change_to_scale(scale, cloud_fr, cloud_to)

mesh = self.mesh * scale

输出如下：

return self.norm(torch.from_numpy(img_fr.astype(np.float32))), \
               torch.LongTensor(choose_fr.astype(np.int32)), \
               torch.from_numpy(cloud_fr.astype(np.float32)), \
               torch.from_numpy(r_fr.astype(np.float32)), \
               torch.from_numpy(t_fr.astype(np.float32)), \
               self.norm(torch.from_numpy(img_to.astype(np.float32))), \
               torch.LongTensor(choose_to.astype(np.int32)), \
               torch.from_numpy(cloud_to.astype(np.float32)), \
               torch.from_numpy(r_to.astype(np.float32)), \
               torch.from_numpy(t_to.astype(np.float32)), \
               torch.from_numpy(mesh.astype(np.float32)), \
               torch.from_numpy(anchor_box.astype(np.float32)), \
               torch.from_numpy(scale.astype(np.float32)), \
               torch.LongTensor(class_gt.astype(np.int32))

img：color_crop
choose:所选的点云对应在2D图中的像素索引
cloud：选择的500个点云
r：当前位姿中的R
t:当前位姿中的T
mesh：自带的3维边框的网格点
anchor：锚点网格的坐标（5*5*5个锚点）
scale：归一化范围

至此完成论文中步骤1、步骤2
在这里插入图片描述
以某一个实例为例，该实例的color_crop大小为[3,160,160]，输出大小如下：