以下链接是个人关于DenseFusion(6D姿态估计) 所有见解,如有错误欢迎大家指出,我会第一时间纠正。有兴趣的朋友可以加微信:17575010159 相互讨论技术。若是帮助到了你什么,一定要记得点赞!因为这是对我最大的鼓励。
文末附带
\color{blue}{文末附带}
文末附带
公众号
−
\color{blue}{公众号 -}
公众号−
海量资源。
\color{blue}{ 海量资源}。
海量资源。
姿态估计0-00:DenseFusion(6D姿态估计)-目录-史上最新无死角讲解https://blog.csdn.net/weixin_43013761/article/details/103053585
数据集讲解
通过前面,我们已经了解整个网络框架的训练流程,现在我们先来看看我们训练的数据集Linemod_preprocessed,本人存在如下上个子文件夹:
如果你的和我的长得不一样,那么抱歉,我也无能无力了。首先我们要了解的是,Linemod_preprocessed这个数据集中,一共用15种目标物体,首先我们来看看其中的data目录,如下:
分别对应十五个目标物体,随便既然一个文件夹,本人进入01如下:
其中以及划分了验证和训练数据,保存的都是对应图片的索引,这样固定的划分好,是为了方便论文的复现。下面是对每个文件的总结:
1. depth(目录): 深度图
2. mask(目录):目标物体的掩码,为标准的分割结果
3. rgb(目录):保存为RGB图像
5. gt.yml(文件):保存了拍摄每张图片时,其对应的旋转矩阵和偏移矩阵,以及目标物体的标准box,
和该图像中目标物体所属于的类别
6. info.yml(文件):拍摄每张图像,对应的摄像头的内参,以及深度缩放的比例
这样,我们对于data已经了解完成了,我们再来看看Linemod_preprocessed/segnet_results,进入该文件如下:
其和Linemod_preprocessed/data/mask可以说时一一对应的,大小认真的对比下,发现他们对应的图片时相差无几的,那么这个又什么用?首先我们从名字来看,可以知道segnet_results时通过对目标物体进行语义分割得到的结果。再训练我们模型的时候,我们肯定希望时拿最标准的数据去训练我们的模型,所以训练和验证,都是使用最标准的mask(Linemod_preprocessed/data/mask),作为分割的标签。但是再测试的时候,是和应用场景达搭钩了,我们希望越接近实际场景越好,在实际场景中,DenseFusion网络是对分割出来的目标进行姿态预测,是的,是分割出来的物体,所以实际应用中,还需要一个分割网络。所以segnet_results中,保存的就是语义分割网络分割出来的图片,为了和实际场景更加接近,这也是倒是其和标准的/data/mask有点点区别的原因。
再下来我们就是介绍Linemod_preprocessed/models文件夹,进入如下:
这里的models可以理解为模型,是什么模型呢?那就是我们目标物体对应的模型,这些模型都用点云数据表示,并保存了起来。那么问题来了。在Linemod_preprocessed/data/xx/rgb中,有那么多的图像,这里的models是对应那个图像的模型呢? 进入obj_01.ply文件可以看到如下:
ply
format ascii 1.0
comment VCGLIB generated
element vertex 5841
property float x
property float y
property float z
property float nx
property float ny
property float nz
property uchar red
property uchar green
property uchar blue
property uchar alpha
element face 11678
property list uchar int vertex_indices
end_header
-23.0565 17.4191 -6.756 -0.617876 0.77406 -0.138061 148 57 57 255
-12.8773 -29.5554 -27.7156 -0.4007 -0.894836 0.196745 140 53 51 255
-32.2776 -17.7052 -32.4505 -0.813037 -0.537012 -0.224921 151 61 58 255
这些都是点云数据,简单的介绍可以参考如下连接:
https://blog.csdn.net/qq_36559293/article/details/90949295
可以想象得到,我们在对一个物体构建3D点云数据得时候,我们需要一个基准参考面,如果没有一个基准,我们拍摄多张照片,每个照片构建出来的点云数据都都不一样(那怕他们是同一个物体)。
所以Linemod_preprocessed\models中每个ply文件,其中的点云信息,都是以linemod\Linemod_preprocessed\data\xx\rgb\0000.png作为参考面的。有了这个参考面的点云数据(包含了整个目标的完整点云),我们就能根据其他的拍摄照片时的摄像头参数(旋转和偏移矩阵),计算出其他的视觉对应的点云数据了。
总的来说,我们对第一帧图片进行3D点云数据的建模(当然,整个过程需要多个面的图片-数据的制作过程),建模之后,我们对着目标随便拍摄一张图片,就能根据摄像头的参数,结合第一帧图片的3D点云数据,计算出该图片对应的3D点云数据了。
最后还有就是models_info.yml文件,其保存的是每个目标点云模型的半径,x,y,z轴的起始值和大小范围。
这样,我们对数据集的介绍就完成了。下面我们来看看对数据集的预处理过程
数据预处理
在tools/train.py中,我们可以看到如下:
PoseDataset_ycb()
PoseDataset_warehouse()
PoseDataset_linemod()
其都是对不同数据集的加载过程,对于本人要讲解的当然是PoseDataset_linemod,其代码的实现是在datasets/linemod/dataset.py:
import torch.utils.data as data
from PIL import Image
import os
import os.path
import errno
import torch
import json
import codecs
import numpy as np
import sys
import torchvision.transforms as transforms
import argparse
import json
import time
import random
import numpy
import numpy.ma as ma
import copy
import scipy.misc
import scipy.io as scio
import yaml
import cv2
numpy.set_printoptions(threshold= 1e6)
class PoseDataset(data.Dataset):
def __init__(self, mode, num, add_noise, root, noise_trans, refine):
"""
:param mode: 可以选择train,test,eval
:param num: mesh点的数目
:param add_noise:是否加入噪声
:param root:数据集的根目录
:param noise_trans:噪声增强的相关参数
:param refine:是否需要为refine模型提供相应的数据
"""
self.objlist = [1, 2, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15]
self.mode = mode
self.list_rgb = []
self.list_depth = []
self.list_label = []
self.list_obj = []
self.list_rank = []
self.meta = {}
self.pt = {}
self.root = root
self.noise_trans = noise_trans
self.refine = refine
item_count = 0
for item in self.objlist:
if self.mode == 'train':
input_file = open('{0}/data/{1}/train.txt'.format(self.root, '%02d' % item))
else:
input_file = open('{0}/data/{1}/test.txt'.format(self.root, '%02d' % item))
while 1:
item_count += 1
input_line = input_file.readline()
if self.mode == 'test' and item_count % 10 != 0:
continue
if not input_line:
break
if input_line[-1:] == '\n':
input_line = input_line[:-1]
self.list_rgb.append('{0}/data/{1}/rgb/{2}.png'.format(self.root, '%02d' % item, input_line))
self.list_depth.append('{0}/data/{1}/depth/{2}.png'.format(self.root, '%02d' % item, input_line))
if self.mode == 'eval':
self.list_label.append('{0}/segnet_results/{1}_label/{2}_label.png'.format(self.root, '%02d' % item, input_line))
else:
self.list_label.append('{0}/data/{1}/mask/{2}.png'.format(self.root, '%02d' % item, input_line))
self.list_obj.append(item)
self.list_rank.append(int(input_line))
meta_file = open('{0}/data/{1}/gt.yml'.format(self.root, '%02d' % item), 'r')
self.meta[item] = yaml.load(meta_file)
self.pt[item] = ply_vtx('{0}/models/obj_{1}.ply'.format(self.root, '%02d' % item))
print("Object {0} buffer loaded".format(item))
self.length = len(self.list_rgb)
self.cam_cx = 325.26110
self.cam_cy = 242.04899
self.cam_fx = 572.41140
self.cam_fy = 573.57043
self.xmap = np.array([[j for i in range(640)] for j in range(480)])
self.ymap = np.array([[i for i in range(640)] for j in range(480)])
self.num = num
self.add_noise = add_noise
self.trancolor = transforms.ColorJitter(0.2, 0.2, 0.2, 0.05)
self.norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
self.border_list = [-1, 40, 80, 120, 160, 200, 240, 280, 320, 360, 400, 440, 480, 520, 560, 600, 640, 680]
self.num_pt_mesh_large = 500
self.num_pt_mesh_small = 500
self.symmetry_obj_idx = [7, 8]
def __getitem__(self, index):
img = Image.open(self.list_rgb[index])
depth = np.array(Image.open(self.list_depth[index]))
label = np.array(Image.open(self.list_label[index]))
obj = self.list_obj[index]
rank = self.list_rank[index]
if obj == 2:
for i in range(0, len(self.meta[obj][rank])):
if self.meta[obj][rank][i]['obj_id'] == 2:
meta = self.meta[obj][rank][i]
break
else:
meta = self.meta[obj][rank][0]
mask_depth = ma.getmaskarray(ma.masked_not_equal(depth, 0))
if self.mode == 'eval':
mask_label = ma.getmaskarray(ma.masked_equal(label, np.array(255)))
else:
mask_label = ma.getmaskarray(ma.masked_equal(label, np.array([255, 255, 255])))[:, :, 0]
mask = mask_label * mask_depth
if self.add_noise:
img = self.trancolor(img)
img = np.array(img)[:, :, :3]
img = np.transpose(img, (2, 0, 1))
img_masked = img
if self.mode == 'eval':
rmin, rmax, cmin, cmax = get_bbox(mask_to_bbox(mask_label))
else:
rmin, rmax, cmin, cmax = get_bbox(meta['obj_bb'])
img_masked = img_masked[:, rmin:rmax, cmin:cmax]
target_r = np.resize(np.array(meta['cam_R_m2c']), (3, 3))
target_t = np.array(meta['cam_t_m2c'])
add_t = np.array([random.uniform(-self.noise_trans, self.noise_trans) for i in range(3)])
choose = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0]
if len(choose) == 0:
cc = torch.LongTensor([0])
return(cc, cc, cc, cc, cc, cc)
if len(choose) > self.num:
c_mask = np.zeros(len(choose), dtype=int)
c_mask[:self.num] = 1
np.random.shuffle(c_mask)
choose = choose[c_mask.nonzero()]
else:
choose = np.pad(choose, (0, self.num - len(choose)), 'wrap')
depth_masked = depth[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)
xmap_masked = self.xmap[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)
ymap_masked = self.ymap[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)
choose = np.array([choose])
cam_scale = 1.0
pt2 = depth_masked / cam_scale
pt0 = (ymap_masked - self.cam_cx) * pt2 / self.cam_fx
pt1 = (xmap_masked - self.cam_cy) * pt2 / self.cam_fy
cloud = np.concatenate((pt0, pt1, pt2), axis=1)
cloud = cloud / 1000.0
if self.add_noise:
cloud = np.add(cloud, add_t)
model_points = self.pt[obj] / 1000.0
dellist = [j for j in range(0, len(model_points))]
dellist = random.sample(dellist, len(model_points) - self.num_pt_mesh_small)
model_points = np.delete(model_points, dellist, axis=0)
target = np.dot(model_points, target_r.T)
if self.add_noise:
target = np.add(target, target_t / 1000.0 + add_t)
out_t = target_t / 1000.0 + add_t
else:
target = np.add(target, target_t / 1000.0)
out_t = target_t / 1000.0
return torch.from_numpy(cloud.astype(np.float32)), \
torch.LongTensor(choose.astype(np.int32)), \
self.norm(torch.from_numpy(img_masked.astype(np.float32))), \
torch.from_numpy(target.astype(np.float32)), \
torch.from_numpy(model_points.astype(np.float32)), \
torch.LongTensor([self.objlist.index(obj)])
def __len__(self):
return self.length
def get_sym_list(self):
return self.symmetry_obj_idx
def get_num_points_mesh(self):
if self.refine:
return self.num_pt_mesh_large
else:
return self.num_pt_mesh_small
border_list = [-1, 40, 80, 120, 160, 200, 240, 280, 320, 360, 400, 440, 480, 520, 560, 600, 640, 680]
img_width = 480
img_length = 640
def mask_to_bbox(mask):
mask = mask.astype(np.uint8)
contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
x = 0
y = 0
w = 0
h = 0
for contour in contours:
tmp_x, tmp_y, tmp_w, tmp_h = cv2.boundingRect(contour)
if tmp_w * tmp_h > w * h:
x = tmp_x
y = tmp_y
w = tmp_w
h = tmp_h
return [x, y, w, h]
def get_bbox(bbox):
bbx = [bbox[1], bbox[1] + bbox[3], bbox[0], bbox[0] + bbox[2]]
if bbx[0] < 0:
bbx[0] = 0
if bbx[1] >= 480:
bbx[1] = 479
if bbx[2] < 0:
bbx[2] = 0
if bbx[3] >= 640:
bbx[3] = 639
rmin, rmax, cmin, cmax = bbx[0], bbx[1], bbx[2], bbx[3]
r_b = rmax - rmin
for tt in range(len(border_list)):
if r_b > border_list[tt] and r_b < border_list[tt + 1]:
r_b = border_list[tt + 1]
break
c_b = cmax - cmin
for tt in range(len(border_list)):
if c_b > border_list[tt] and c_b < border_list[tt + 1]:
c_b = border_list[tt + 1]
break
center = [int((rmin + rmax) / 2), int((cmin + cmax) / 2)]
rmin = center[0] - int(r_b / 2)
rmax = center[0] + int(r_b / 2)
cmin = center[1] - int(c_b / 2)
cmax = center[1] + int(c_b / 2)
if rmin < 0:
delt = -rmin
rmin = 0
rmax += delt
if cmin < 0:
delt = -cmin
cmin = 0
cmax += delt
if rmax > 480:
delt = rmax - 480
rmax = 480
rmin -= delt
if cmax > 640:
delt = cmax - 640
cmax = 640
cmin -= delt
return rmin, rmax, cmin, cmax
def ply_vtx(path):
f = open(path)
assert f.readline().strip() == "ply"
f.readline()
f.readline()
N = int(f.readline().split()[-1])
while f.readline().strip() != "end_header":
continue
pts = []
for _ in range(N):
pts.append(np.float32(f.readline().split()[:3]))
return np.array(pts)
我相信,注释已经很详细,我们回到tools/train.py:
points, choose, img, target, model_points, idx = data
其上可以说,就是对整个数据加载,以及预测处理的目的了。
细心的朋友应该发现了,修改batch_size参数,似乎没有什么用,其主要的原因是torch.utils.data.DataLoader()函数的参数batch_size=1设置为1,集每次都是获得一个样本,如果强行修改为其他的整数(如本人修改为2),会报错如下:
taloader.py", line 209, in default_collate
return torch.stack(batch, 0, out=out)
RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 0. Got 160 and 120 in dimension 2 at /pytorch/aten/src/TH/generic/THTensorMoreMath.cpp:1307
大概的原因就是,对于每个我们分割出来的目标图像,其大小都是不一样的,所以我们只能一张一张的去进行训练。多个图片大小不一样的图像,没有办法统一成一个batch。
本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)