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)