def __getitem__(self, item):
# get transformation
rot = self.infos['rot'][item]
trans = self.infos['trans'][item]
# get pointcloud
src_path = os.path.join(self.base_dir, self.infos['src'][item])
tgt_path = os.path.join(self.base_dir, self.infos['tgt'][item])
src_pcd = torch.load(src_path)
tgt_pcd = torch.load(tgt_path)
# if we get too many points, we do some downsampling
if (src_pcd.shape[0] > self.max_points):
idx = np.random.permutation(src_pcd.shape[0])[:self.max_points]
src_pcd = src_pcd[idx]
if (tgt_pcd.shape[0] > self.max_points):
idx = np.random.permutation(tgt_pcd.shape[0])[:self.max_points]
tgt_pcd = tgt_pcd[idx]
# add gaussian noise
if self.data_augmentation:
# rotate the point cloud
euler_ab = np.random.rand(
3) * np.pi * 2 / self.rot_factor # anglez, angley, anglex
rot_ab = Rotation.from_euler('zyx', euler_ab).as_matrix()
if (np.random.rand(1)[0] > 0.5):
src_pcd = np.matmul(rot_ab, src_pcd.T).T
rot = np.matmul(rot, rot_ab.T)
else:
tgt_pcd = np.matmul(rot_ab, tgt_pcd.T).T
rot = np.matmul(rot_ab, rot)
trans = np.matmul(rot_ab, trans)
src_pcd += (np.random.rand(src_pcd.shape[0], 3) -
0.5) * self.augment_noise
tgt_pcd += (np.random.rand(tgt_pcd.shape[0], 3) -
0.5) * self.augment_noise
if (trans.ndim == 1):
trans = trans[:, None]
# get correspondence at fine level
tsfm = to_tsfm(rot, trans)
correspondences = get_correspondences(to_o3d_pcd(src_pcd),
to_o3d_pcd(tgt_pcd), tsfm,
self.overlap_radius)
src_feats = np.ones_like(src_pcd[:, :1]).astype(np.float32)
tgt_feats = np.ones_like(tgt_pcd[:, :1]).astype(np.float32)
rot = rot.astype(np.float32)
trans = trans.astype(np.float32)
return src_pcd, tgt_pcd, src_feats, tgt_feats, rot, trans, correspondences, src_pcd, tgt_pcd, torch.ones(
1)
def __getitem__(self, item):
sample = {'points': self._data[item, :, :], 'label': self._labels[item], 'idx': np.array(item, dtype=np.int32)}
if self._transform:
sample = self._transform(sample)
# transform to our format
src_pcd = sample['points_src'][:,:3]
tgt_pcd = sample['points_ref'][:,:3]
rot = sample['transform_gt'][:,:3]
trans = sample['transform_gt'][:,3][:,None]
matching_inds = get_correspondences(to_o3d_pcd(src_pcd), to_o3d_pcd(tgt_pcd),to_tsfm(rot,trans),self.overlap_radius)
if(self.n_in_feats == 1):
src_feats=np.ones_like(src_pcd[:,:1]).astype(np.float32)
tgt_feats=np.ones_like(tgt_pcd[:,:1]).astype(np.float32)
elif(self.n_in_feats == 3):
src_feats = src_pcd.astype(np.float32)
tgt_feats = tgt_pcd.astype(np.float32)
for k,v in sample.items():
if k not in ['deterministic','label','idx']:
sample[k] = torch.from_numpy(v).unsqueeze(0)
return src_pcd,tgt_pcd,src_feats,tgt_feats,rot,trans, matching_inds, src_pcd, tgt_pcd, sample
def __getitem__(self, item):
# get transformation
rot = self.infos['rot'][item]
trans = self.infos['trans'][item]
# get pointcloud
src_path = os.path.join(self.base_dir, self.infos['src'][item])
tgt_path = os.path.join(self.base_dir, self.infos['tgt'][item])
src_pcd = torch.load(src_path)
tgt_pcd = torch.load(tgt_path)
# add gaussian noise
if self.data_augmentation:
# rotate the point cloud
euler_ab = np.random.rand(
3) * np.pi * 2 / self.rot_factor # anglez, angley, anglex
rot_ab = Rotation.from_euler('zyx', euler_ab).as_matrix()
if (np.random.rand(1)[0] > 0.5):
src_pcd = np.matmul(rot_ab, src_pcd.T).T
rot = np.matmul(rot, rot_ab.T)
else:
tgt_pcd = np.matmul(rot_ab, tgt_pcd.T).T
rot = np.matmul(rot_ab, rot)
trans = np.matmul(rot_ab, trans)
src_pcd += (np.random.rand(src_pcd.shape[0], 3) -
0.5) * self.augment_noise
tgt_pcd += (np.random.rand(tgt_pcd.shape[0], 3) -
0.5) * self.augment_noise
if (trans.ndim == 1):
trans = trans[:, None]
# build sparse tensor
_, sel_src = ME.utils.sparse_quantize(np.ascontiguousarray(src_pcd) /
self.voxel_size,
return_index=True)
_, sel_tgt = ME.utils.sparse_quantize(np.ascontiguousarray(tgt_pcd) /
self.voxel_size,
return_index=True)
# get correspondence
tsfm = to_tsfm(rot, trans)
src_xyz, tgt_xyz = src_pcd[sel_src], tgt_pcd[
sel_tgt] # raw point clouds
matching_inds = get_correspondences(to_o3d_pcd(src_xyz),
to_o3d_pcd(tgt_xyz), tsfm,
self.search_voxel_size)
# get voxelized coordinates
src_coords, tgt_coords = np.floor(src_xyz / self.voxel_size), np.floor(
tgt_xyz / self.voxel_size)
# get feats
src_feats = np.ones((src_coords.shape[0], 1), dtype=np.float32)
tgt_feats = np.ones((tgt_coords.shape[0], 1), dtype=np.float32)
src_xyz, tgt_xyz = to_tensor(src_xyz).float(), to_tensor(
tgt_xyz).float()
rot, trans = to_tensor(rot), to_tensor(trans)
scale = 1
return src_xyz, tgt_xyz, src_coords, tgt_coords, src_feats, tgt_feats, matching_inds, rot, trans, scale
def __getitem__(self, idx):
drive = self.files[idx][0]
t0, t1 = self.files[idx][1], self.files[idx][2]
all_odometry = self.get_video_odometry(drive, [t0, t1])
positions = [
self.odometry_to_positions(odometry) for odometry in all_odometry
]
fname0 = self._get_velodyne_fn(drive, t0)
fname1 = self._get_velodyne_fn(drive, t1)
# extract xyz
xyz0 = np.fromfile(fname0, dtype=np.float32).reshape(-1, 4)[:, :3]
xyz1 = np.fromfile(fname1, dtype=np.float32).reshape(-1, 4)[:, :3]
# use ICP to refine the ground_truth pose, for ICP we don't voxllize the point clouds
key = '%d_%d_%d' % (drive, t0, t1)
filename = self.icp_path + '/' + key + '.npy'
if key not in self.kitti_icp_cache:
if not os.path.exists(filename):
print('missing ICP files, recompute it')
M = (self.velo2cam @ positions[0].T @ np.linalg.inv(
positions[1].T) @ np.linalg.inv(self.velo2cam)).T
xyz0_t = self.apply_transform(xyz0, M)
pcd0 = to_o3d_pcd(xyz0_t)
pcd1 = to_o3d_pcd(xyz1)
reg = open3d.registration.registration_icp(
pcd0, pcd1, 0.2, np.eye(4),
open3d.registration.TransformationEstimationPointToPoint(),
open3d.registration.ICPConvergenceCriteria(
max_iteration=200))
pcd0.transform(reg.transformation)
M2 = M @ reg.transformation
np.save(filename, M2)
else:
M2 = np.load(filename)
self.kitti_icp_cache[key] = M2
else:
M2 = self.kitti_icp_cache[key]
# refined pose is denoted as tsfm
tsfm = M2
rot = tsfm[:3, :3]
trans = tsfm[:3, 3][:, None]
# add data augmentation
src_pcd_input = copy.deepcopy(xyz0)
tgt_pcd_input = copy.deepcopy(xyz1)
if (self.data_augmentation):
# add gaussian noise
src_pcd_input += (np.random.rand(src_pcd_input.shape[0], 3) -
0.5) * self.augment_noise
tgt_pcd_input += (np.random.rand(tgt_pcd_input.shape[0], 3) -
0.5) * self.augment_noise
# rotate the point cloud
euler_ab = np.random.rand(3) * np.pi * 2 # anglez, angley, anglex
rot_ab = Rotation.from_euler('zyx', euler_ab).as_matrix()
if (np.random.rand(1)[0] > 0.5):
src_pcd_input = np.dot(rot_ab, src_pcd_input.T).T
rot = np.matmul(rot, rot_ab.T)
else:
tgt_pcd_input = np.dot(rot_ab, tgt_pcd_input.T).T
rot = np.matmul(rot_ab, rot)
trans = np.matmul(rot_ab, trans)
# scale the pcd
scale = self.augment_scale_min + (
self.augment_scale_max -
self.augment_scale_min) * random.random()
src_pcd_input = src_pcd_input * scale
tgt_pcd_input = tgt_pcd_input * scale
trans = scale * trans
else:
scale = 1
# voxel down-sample the point clouds here
_, sel_src = ME.utils.sparse_quantize(
np.ascontiguousarray(src_pcd_input) / self.voxel_size,
return_index=True)
_, sel_tgt = ME.utils.sparse_quantize(
np.ascontiguousarray(tgt_pcd_input) / self.voxel_size,
return_index=True)
# get correspondence
tsfm = to_tsfm(rot, trans)
src_xyz, tgt_xyz = src_pcd_input[sel_src], tgt_pcd_input[
sel_tgt] # raw point clouds
matching_inds = get_correspondences(to_o3d_pcd(src_xyz),
to_o3d_pcd(tgt_xyz), tsfm,
self.search_voxel_size * scale)
if (matching_inds.size(0) < self.max_corr and self.split == 'train'):
return self.__getitem__(np.random.choice(len(self.files), 1)[0])
# get voxelized coordinates
src_coords, tgt_coords = np.floor(src_xyz / self.voxel_size), np.floor(
tgt_xyz / self.voxel_size)
# get feats
src_feats = np.ones((src_coords.shape[0], 1), dtype=np.float32)
tgt_feats = np.ones((tgt_coords.shape[0], 1), dtype=np.float32)
src_xyz, tgt_xyz = to_tensor(src_xyz).float(), to_tensor(
tgt_xyz).float()
rot, trans = to_tensor(rot), to_tensor(trans)
return src_xyz, tgt_xyz, src_coords, tgt_coords, src_feats, tgt_feats, matching_inds, rot, trans, scale
def __getitem__(self, idx):
drive = self.files[idx][0]
t0, t1 = self.files[idx][1], self.files[idx][2]
all_odometry = self.get_video_odometry(drive, [t0, t1])
positions = [self.odometry_to_positions(odometry) for odometry in all_odometry]
fname0 = self._get_velodyne_fn(drive, t0)
fname1 = self._get_velodyne_fn(drive, t1)
# XYZ and reflectance
xyzr0 = np.fromfile(fname0, dtype=np.float32).reshape(-1, 4)
xyzr1 = np.fromfile(fname1, dtype=np.float32).reshape(-1, 4)
xyz0 = xyzr0[:, :3]
xyz1 = xyzr1[:, :3]
# use ICP to refine the ground_truth pose, for ICP we don't voxllize the point clouds
key = '%d_%d_%d' % (drive, t0, t1)
filename = self.icp_path + '/' + key + '.npy'
if key not in self.kitti_icp_cache:
if not os.path.exists(filename):
print('missing ICP files, recompute it')
M = (self.velo2cam @ positions[0].T @ np.linalg.inv(positions[1].T)
@ np.linalg.inv(self.velo2cam)).T
xyz0_t = self.apply_transform(xyz0, M)
pcd0 = to_o3d_pcd(xyz0_t)
pcd1 = to_o3d_pcd(xyz1)
reg = open3d.registration.registration_icp(pcd0, pcd1, 0.2, np.eye(4),
open3d.registration.TransformationEstimationPointToPoint(),
open3d.registration.ICPConvergenceCriteria(max_iteration=200))
pcd0.transform(reg.transformation)
M2 = M @ reg.transformation
np.save(filename, M2)
else:
M2 = np.load(filename)
self.kitti_icp_cache[key] = M2
else:
M2 = self.kitti_icp_cache[key]
# refined pose is denoted as trans
tsfm = M2
rot = tsfm[:3,:3]
trans = tsfm[:3,3][:,None]
# voxelize the point clouds here
pcd0 = to_o3d_pcd(xyz0)
pcd1 = to_o3d_pcd(xyz1)
pcd0 = pcd0.voxel_down_sample(self.voxel_size)
pcd1 = pcd1.voxel_down_sample(self.voxel_size)
src_pcd = np.array(pcd0.points)
tgt_pcd = np.array(pcd1.points)
# Get matches
matching_inds = get_correspondences(pcd0, pcd1, tsfm, self.matching_search_voxel_size)
if(matching_inds.size(0) < self.max_corr and self.split == 'train'):
return self.__getitem__(np.random.choice(len(self.files),1)[0])
src_feats=np.ones_like(src_pcd[:,:1]).astype(np.float32)
tgt_feats=np.ones_like(tgt_pcd[:,:1]).astype(np.float32)
rot = rot.astype(np.float32)
trans = trans.astype(np.float32)
# add data augmentation
src_pcd_input = copy.deepcopy(src_pcd)
tgt_pcd_input = copy.deepcopy(tgt_pcd)
if(self.data_augmentation):
# add gaussian noise
src_pcd_input += (np.random.rand(src_pcd_input.shape[0],3) - 0.5) * self.augment_noise
tgt_pcd_input += (np.random.rand(tgt_pcd_input.shape[0],3) - 0.5) * self.augment_noise
# rotate the point cloud
euler_ab=np.random.rand(3)*np.pi*2 # anglez, angley, anglex
rot_ab= Rotation.from_euler('zyx', euler_ab).as_matrix()
if(np.random.rand(1)[0]>0.5):
src_pcd_input = np.dot(rot_ab, src_pcd_input.T).T
else:
tgt_pcd_input = np.dot(rot_ab, tgt_pcd_input.T).T
# scale the pcd
scale = self.augment_scale_min + (self.augment_scale_max - self.augment_scale_min) * random.random()
src_pcd_input = src_pcd_input * scale
tgt_pcd_input = tgt_pcd_input * scale
# shift the pcd
shift_src = np.random.uniform(-self.augment_shift_range, self.augment_shift_range, 3)
shift_tgt = np.random.uniform(-self.augment_shift_range, self.augment_shift_range, 3)
src_pcd_input = src_pcd_input + shift_src
tgt_pcd_input = tgt_pcd_input + shift_tgt
return src_pcd_input, tgt_pcd_input, src_feats, tgt_feats, rot, trans, matching_inds, src_pcd, tgt_pcd, torch.ones(1)
def draw_registration_result(src_raw, tgt_raw, src_overlap, tgt_overlap,
src_saliency, tgt_saliency, tsfm):
########################################
# 1. input point cloud
src_pcd_before = to_o3d_pcd(src_raw)
tgt_pcd_before = to_o3d_pcd(tgt_raw)
src_pcd_before.paint_uniform_color(get_yellow())
tgt_pcd_before.paint_uniform_color(get_blue())
src_pcd_before.estimate_normals(
search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.3,
max_nn=50))
tgt_pcd_before.estimate_normals(
search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.3,
max_nn=50))
########################################
# 2. overlap colors
rot, trans = to_tensor(tsfm[:3, :3]), to_tensor(tsfm[:3, 3][:, None])
src_overlap = src_overlap[:, None].repeat(1, 3).numpy()
tgt_overlap = tgt_overlap[:, None].repeat(1, 3).numpy()
src_overlap_color = lighter(get_yellow(), 1 - src_overlap)
tgt_overlap_color = lighter(get_blue(), 1 - tgt_overlap)
src_pcd_overlap = copy.deepcopy(src_pcd_before)
src_pcd_overlap.transform(tsfm)
tgt_pcd_overlap = copy.deepcopy(tgt_pcd_before)
src_pcd_overlap.colors = o3d.utility.Vector3dVector(src_overlap_color)
tgt_pcd_overlap.colors = o3d.utility.Vector3dVector(tgt_overlap_color)
########################################
# 3. draw registrations
src_pcd_after = copy.deepcopy(src_pcd_before)
src_pcd_after.transform(tsfm)
vis1 = o3d.visualization.Visualizer()
vis1.create_window(window_name='Input',
width=960,
height=540,
left=0,
top=0)
vis1.add_geometry(src_pcd_before)
vis1.add_geometry(tgt_pcd_before)
vis2 = o3d.visualization.Visualizer()
vis2.create_window(window_name='Inferred overlap region',
width=960,
height=540,
left=0,
top=600)
vis2.add_geometry(src_pcd_overlap)
vis2.add_geometry(tgt_pcd_overlap)
vis3 = o3d.visualization.Visualizer()
vis3.create_window(window_name='Our registration',
width=960,
height=540,
left=960,
top=0)
vis3.add_geometry(src_pcd_after)
vis3.add_geometry(tgt_pcd_before)
while True:
vis1.update_geometry(src_pcd_before)
vis3.update_geometry(tgt_pcd_before)
if not vis1.poll_events():
break
vis1.update_renderer()
vis2.update_geometry(src_pcd_overlap)
vis2.update_geometry(tgt_pcd_overlap)
if not vis2.poll_events():
break
vis2.update_renderer()
vis3.update_geometry(src_pcd_after)
vis3.update_geometry(tgt_pcd_before)
if not vis3.poll_events():
break
vis3.update_renderer()
vis1.destroy_window()
vis2.destroy_window()
vis3.destroy_window()