def __getitem__(self, idx): file0 = os.path.join(self.root, self.files[idx][0]) file1 = os.path.join(self.root, self.files[idx][1]) data0 = np.load(file0) data1 = np.load(file1) xyz0 = data0["pcd"] xyz1 = data1["pcd"] color0 = data0["color"] color1 = data1["color"] matching_search_voxel_size = self.matching_search_voxel_size if self.random_scale and random.random() < 0.95: scale = self.min_scale + \ (self.max_scale - self.min_scale) * random.random() matching_search_voxel_size *= scale xyz0 = scale * xyz0 xyz1 = scale * xyz1 if self.random_rotation: T0 = sample_random_trans(xyz0, self.randg, self.rotation_range) T1 = sample_random_trans(xyz1, self.randg, self.rotation_range) trans = T1 @ np.linalg.inv(T0) xyz0 = self.apply_transform(xyz0, T0) xyz1 = self.apply_transform(xyz1, T1) else: trans = np.identity(4) # Voxelization _, sel0 = ME.utils.sparse_quantize(xyz0 / self.voxel_size, return_index=True) _, sel1 = ME.utils.sparse_quantize(xyz1 / self.voxel_size, return_index=True) # Make point clouds using voxelized points pcd0 = make_open3d_point_cloud(xyz0) pcd1 = make_open3d_point_cloud(xyz1) # Select features and points using the returned voxelized indices pcd0.colors = o3d.utility.Vector3dVector(color0[sel0]) pcd1.colors = o3d.utility.Vector3dVector(color1[sel1]) pcd0.points = o3d.utility.Vector3dVector(np.array(pcd0.points)[sel0]) pcd1.points = o3d.utility.Vector3dVector(np.array(pcd1.points)[sel1]) # Get matches matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size) # Get features npts0 = len(pcd0.colors) npts1 = len(pcd1.colors) feats_train0, feats_train1 = [], [] feats_train0.append(np.ones((npts0, 1))) feats_train1.append(np.ones((npts1, 1))) feats0 = np.hstack(feats_train0) feats1 = np.hstack(feats_train1) # Get coords xyz0 = np.array(pcd0.points) xyz1 = np.array(pcd1.points) coords0 = np.floor(xyz0 / self.voxel_size) coords1 = np.floor(xyz1 / self.voxel_size) if self.transform: coords0, feats0 = self.transform(coords0, feats0) coords1, feats1 = self.transform(coords1, feats1) return (xyz0, xyz1, coords0, coords1, feats0, feats1, matches, trans)
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] key = '%d_%d_%d' % (drive, t0, t1) filename = self.icp_path + '/' + key + '.npy' if key not in kitti_icp_cache: if not os.path.exists(filename): # work on the downsampled xyzs, 0.05m == 5cm _, sel0 = ME.utils.sparse_quantize(xyz0 / 0.05, return_index=True) _, sel1 = ME.utils.sparse_quantize(xyz1 / 0.05, return_index=True) M = (self.velo2cam @ positions[0].T @ np.linalg.inv(positions[1].T) @ np.linalg.inv(self.velo2cam)).T xyz0_t = self.apply_transform(xyz0[sel0], M) pcd0 = make_open3d_point_cloud(xyz0_t) pcd1 = make_open3d_point_cloud(xyz1[sel1]) reg = o3d.registration.registration_icp( pcd0, pcd1, 0.2, np.eye(4), o3d.registration.TransformationEstimationPointToPoint(), o3d.registration.ICPConvergenceCriteria(max_iteration=200)) pcd0.transform(reg.transformation) # pcd0.transform(M2) or self.apply_transform(xyz0, M2) M2 = M @ reg.transformation # o3d.draw_geometries([pcd0, pcd1]) # write to a file np.save(filename, M2) else: M2 = np.load(filename) kitti_icp_cache[key] = M2 else: M2 = kitti_icp_cache[key] if self.random_rotation: T0 = sample_random_trans(xyz0, self.randg, np.pi / 4) T1 = sample_random_trans(xyz1, self.randg, np.pi / 4) trans = T1 @ M2 @ np.linalg.inv(T0) xyz0 = self.apply_transform(xyz0, T0) xyz1 = self.apply_transform(xyz1, T1) else: trans = M2 matching_search_voxel_size = self.matching_search_voxel_size if self.random_scale and random.random() < 0.95: scale = self.min_scale + \ (self.max_scale - self.min_scale) * random.random() matching_search_voxel_size *= scale xyz0 = scale * xyz0 xyz1 = scale * xyz1 # Voxelization xyz0_th = torch.from_numpy(xyz0) xyz1_th = torch.from_numpy(xyz1) _, sel0 = ME.utils.sparse_quantize(xyz0_th / self.voxel_size, return_index=True) _, sel1 = ME.utils.sparse_quantize(xyz1_th / self.voxel_size, return_index=True) # Make point clouds using voxelized points pcd0 = make_open3d_point_cloud(xyz0[sel0]) pcd1 = make_open3d_point_cloud(xyz1[sel1]) # Get matches matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size) if len(matches) < 1000: raise ValueError(f"{drive}, {t0}, {t1}") # Get features npts0 = len(sel0) npts1 = len(sel1) feats_train0, feats_train1 = [], [] unique_xyz0_th = xyz0_th[sel0] unique_xyz1_th = xyz1_th[sel1] feats_train0.append(torch.ones((npts0, 1))) feats_train1.append(torch.ones((npts1, 1))) feats0 = torch.cat(feats_train0, 1) feats1 = torch.cat(feats_train1, 1) coords0 = torch.floor(unique_xyz0_th / self.voxel_size) coords1 = torch.floor(unique_xyz1_th / self.voxel_size) if self.transform: coords0, feats0 = self.transform(coords0, feats0) coords1, feats1 = self.transform(coords1, feats1) return (unique_xyz0_th.float(), unique_xyz1_th.float(), coords0.int(), coords1.int(), feats0.float(), feats1.float(), matches, trans)
def __getitem__(self, idx): file0 = os.path.join(self.root, self.files[idx][0]) file1 = os.path.join(self.root, self.files[idx][1]) data0 = np.load(file0) data1 = np.load(file1) # 这俩点云读进来是可以拼在一起的,只不过都只是完整点云的一部分 xyz0 = data0["pcd"]#50万左右个点 xyz1 = data1["pcd"] color0 = data0["color"]#每个点都有它自己的颜色,且不一样 color1 = data1["color"] matching_search_voxel_size = self.matching_search_voxel_size #0.0375 if self.random_scale and random.random() < 0.95: scale = self.min_scale + \ (self.max_scale - self.min_scale) * random.random() matching_search_voxel_size *= scale xyz0 = scale * xyz0 xyz1 = scale * xyz1 if self.random_rotation: T0 = sample_random_trans(xyz0, self.randg, self.rotation_range) #rotation_range:360 且旋转平移之后的点云中心坐标为0 T1 = sample_random_trans(xyz1, self.randg, self.rotation_range) # xyz0到xyz1的相对位姿 trans = T1 @ np.linalg.inv(T0) xyz0 = self.apply_transform(xyz0, T0) xyz1 = self.apply_transform(xyz1, T1) else: trans = np.identity(4) # Voxelization voxel_size 0.025 _, sel0 = ME.utils.sparse_quantize(xyz0 / self.voxel_size, return_index=True) _, sel1 = ME.utils.sparse_quantize(xyz1 / self.voxel_size, return_index=True) # 用稀疏点云生成pcd pcd0 = make_open3d_point_cloud(xyz0) pcd1 = make_open3d_point_cloud(xyz1) # Select features and points using the returned voxelized indices pcd0.colors = o3d.utility.Vector3dVector(color0[sel0]) pcd1.colors = o3d.utility.Vector3dVector(color1[sel1]) pcd0.points = o3d.utility.Vector3dVector(np.array(pcd0.points)[sel0]) pcd1.points = o3d.utility.Vector3dVector(np.array(pcd1.points)[sel1]) # Get matches # matches包含了点云1的每一点在点云二的(一定范围内的 #0.0375)匹配点 matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size) # Get features,特征都为1 npts0 = len(pcd0.colors) npts1 = len(pcd1.colors) feats_train0, feats_train1 = [], [] feats_train0.append(np.ones((npts0, 1))) feats_train1.append(np.ones((npts1, 1))) feats0 = np.hstack(feats_train0) feats1 = np.hstack(feats_train1) # Get coords xyz0 = np.array(pcd0.points) xyz1 = np.array(pcd1.points) coords0 = np.floor(xyz0 / self.voxel_size) coords1 = np.floor(xyz1 / self.voxel_size) if self.transform: coords0, feats0 = self.transform(coords0, feats0) coords1, feats1 = self.transform(coords1, feats1) # 返回体素化之后的np点云,稀疏点云坐标、特征,匹配关系,真实旋转平移 return (xyz0, xyz1, coords0, coords1, feats0, feats1, matches, trans)
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] coords0 = (xyz0 - xyz0.min(0)) / 0.05 coords1 = (xyz1 - xyz1.min(0)) / 0.05 sel0 = ME.utils.sparse_quantize(coords0, return_index=True) sel1 = ME.utils.sparse_quantize(coords1, return_index=True) xyz0 = xyz0[sel0] xyz1 = xyz1[sel1] # r0 = xyzr0[:, -1].reshape(-1, 1) # r1 = xyzr1[:, -1].reshape(-1, 1) # pcd0 = make_open3d_point_cloud(xyz0_t, 0.7 * np.ones((len(xyz0), 3))) # pcd1 = make_open3d_point_cloud(xyz1, 0.3 * np.ones((len(xyz1), 3))) key = '%d_%d_%d' % (drive, t0, t1) filename = self.icp_path + '/' + key + '.npy' if key not in kitti_icp_cache: if not os.path.exists(filename): if self.IS_ODOMETRY: M = (self.velo2cam @ positions[0].T @ np.linalg.inv(positions[1].T) @ np.linalg.inv(self.velo2cam)).T else: M = self.get_position_transform(positions[0], positions[1], invert=True).T xyz0_t = self.apply_transform(xyz0, M) pcd0 = make_open3d_point_cloud(xyz0_t) pcd1 = make_open3d_point_cloud(xyz1) reg = o3d.registration_icp(pcd0, pcd1, 0.2, np.eye(4), o3d.TransformationEstimationPointToPoint(), o3d.ICPConvergenceCriteria(max_iteration=200)) pcd0.transform(reg.transformation) # pcd0.transform(M2) or self.apply_transform(xyz0, M2) M2 = M @ reg.transformation # o3d.draw_geometries([pcd0, pcd1]) # write to a file np.save(filename, M2) else: M2 = np.load(filename) kitti_icp_cache[key] = M2 else: M2 = kitti_icp_cache[key] if self.random_rotation: T0 = sample_random_trans(xyz0, self.randg, np.pi / 4) T1 = sample_random_trans(xyz1, self.randg, np.pi / 4) trans = T1 @ M2 @ np.linalg.inv(T0) xyz0 = self.apply_transform(xyz0, T0) xyz1 = self.apply_transform(xyz1, T1) else: trans = M2 matching_search_voxel_size = self.matching_search_voxel_size if self.random_scale and random.random() < 0.95: scale = self.min_scale + \ (self.max_scale - self.min_scale) * random.random() matching_search_voxel_size *= scale xyz0 = scale * xyz0 xyz1 = scale * xyz1 # Voxelization coords0 = np.floor(xyz0 / self.voxel_size) coords1 = np.floor(xyz1 / self.voxel_size) sel0 = ME.utils.sparse_quantize(coords0, return_index=True) sel1 = ME.utils.sparse_quantize(coords1, return_index=True) coords0, coords1 = coords0[sel0], coords1[sel1] # r0, r1 = r0[sel0], r1[sel1] pcd0 = make_open3d_point_cloud(xyz0) pcd1 = make_open3d_point_cloud(xyz1) # Select features and points using the returned voxelized indices pcd0.points = o3d.utility.Vector3dVector(np.array(pcd0.points)[sel0]) pcd1.points = o3d.utility.Vector3dVector(np.array(pcd1.points)[sel1]) # Get matches matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size) if len(matches) < 1000: raise ValueError(f"{drive}, {t0}, {t1}") feats_train0, feats_train1 = [], [] feats_train0.append(np.ones((len(sel0), 1))) feats_train1.append(np.ones((len(sel1), 1))) feats0 = np.hstack(feats_train0) feats1 = np.hstack(feats_train1) # Get coords xyz0 = np.array(pcd0.points) xyz1 = np.array(pcd1.points) if self.transform: coords0, feats0 = self.transform(coords0, feats0) coords1, feats1 = self.transform(coords1, feats1) return (xyz0, xyz1, coords0, coords1, feats0, feats1, matches, trans)
def __getitem__(self, idx): file0 = os.path.join(self.root, self.files[idx][0]) file1 = os.path.join(self.root, self.files[idx][1]) data0 = np.load(file0) data1 = np.load(file1) xyz0 = data0["pcd"] xyz1 = data1["pcd"] matching_search_voxel_size = self.matching_search_voxel_size if self.random_scale and random.random() < 0.95: scale = self.min_scale + \ (self.max_scale - self.min_scale) * random.random() matching_search_voxel_size *= scale xyz0 = scale * xyz0 xyz1 = scale * xyz1 if self.random_rotation: T0 = sample_random_trans(xyz0, self.randg, self.rotation_range) T1 = sample_random_trans(xyz1, self.randg, self.rotation_range) trans = T1 @ np.linalg.inv(T0) xyz0 = self.apply_transform(xyz0, T0) xyz1 = self.apply_transform(xyz1, T1) else: trans = np.identity(4) # Voxelization xyz0_th = torch.from_numpy(xyz0) xyz1_th = torch.from_numpy(xyz1) _, sel0 = ME.utils.sparse_quantize(xyz0_th / self.voxel_size, return_index=True) _, sel1 = ME.utils.sparse_quantize(xyz1_th / self.voxel_size, return_index=True) # Make point clouds using voxelized points pcd0 = make_open3d_point_cloud(xyz0[sel0]) pcd1 = make_open3d_point_cloud(xyz1[sel1]) # Select features and points using the returned voxelized indices # 3DMatch color is not helpful # pcd0.colors = o3d.utility.Vector3dVector(color0[sel0]) # pcd1.colors = o3d.utility.Vector3dVector(color1[sel1]) # Get matches matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size) # Get features npts0 = len(sel0) npts1 = len(sel1) feats_train0, feats_train1 = [], [] unique_xyz0_th = xyz0_th[sel0] unique_xyz1_th = xyz1_th[sel1] # xyz as feats if self.use_xyz_feature: feats_train0.append(unique_xyz0_th - unique_xyz0_th.mean(0)) feats_train1.append(unique_xyz1_th - unique_xyz1_th.mean(0)) else: feats_train0.append(torch.ones((npts0, 1))) feats_train1.append(torch.ones((npts1, 1))) feats0 = torch.cat(feats_train0, 1) feats1 = torch.cat(feats_train1, 1) coords0 = torch.floor(unique_xyz0_th / self.voxel_size) coords1 = torch.floor(unique_xyz1_th / self.voxel_size) if self.transform: coords0, feats0 = self.transform(coords0, feats0) coords1, feats1 = self.transform(coords1, feats1) extra_package = {'idx': idx, 'file0': file0, 'file1': file1} return (unique_xyz0_th.float(), unique_xyz1_th.float(), coords0.int(), coords1.int(), feats0.float(), feats1.float(), matches, trans, extra_package)
def __getitem__(self, idx): drive = self.U.pairs[idx, 0] t0 = self.U.pairs[idx, 1] t1 = self.U.pairs[idx, 2] M2, xyz0, xyz1 = self.U.get_pair(idx) if self.random_rotation: T0 = sample_almost_planar_rotation(self.randg) T1 = sample_almost_planar_rotation(self.randg) trans = T1 @ M2 @ np.linalg.inv(T0) xyz0 = self.apply_transform(xyz0, T0) xyz1 = self.apply_transform(xyz1, T1) else: trans = M2 matching_search_voxel_size = self.matching_search_voxel_size if self.random_scale and random.random() < 0.95: scale = self.min_scale + \ (self.max_scale - self.min_scale) * random.random() matching_search_voxel_size *= scale xyz0 = scale * xyz0 xyz1 = scale * xyz1 M2[:3, 3] *= scale # Voxelization xyz0_th = torch.from_numpy(xyz0) xyz1_th = torch.from_numpy(xyz1) _, sel0 = ME.utils.sparse_quantize(xyz0_th / self.voxel_size, return_index=True) _, sel1 = ME.utils.sparse_quantize(xyz1_th / self.voxel_size, return_index=True) # Make point clouds using voxelized points pcd0 = make_open3d_point_cloud(xyz0[sel0]) pcd1 = make_open3d_point_cloud(xyz1[sel1]) # Get matches matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size) # Get features npts0 = len(sel0) npts1 = len(sel1) feats_train0, feats_train1 = [], [] unique_xyz0_th = xyz0_th[sel0] unique_xyz1_th = xyz1_th[sel1] feats_train0.append(torch.ones((npts0, 1))) feats_train1.append(torch.ones((npts1, 1))) feats0 = torch.cat(feats_train0, 1) feats1 = torch.cat(feats_train1, 1) coords0 = torch.floor(unique_xyz0_th / self.voxel_size) coords1 = torch.floor(unique_xyz1_th / self.voxel_size) if self.transform: coords0, feats0 = self.transform(coords0, feats0) coords1, feats1 = self.transform(coords1, feats1) extra_package = {'drive': drive, 't0': t0, 't1': t1} return (unique_xyz0_th.float(), unique_xyz1_th.float(), coords0.int(), coords1.int(), feats0.float(), feats1.float(), matches, trans, extra_package)
def __getitem__(self, idx): # split, idx): if self.split == "test": sample = self.list_test_sample[idx] else: T_noise = self.generate_noise_T() sample = self.get_sample(idx, T_noise) xyz0 = sample["source"].points xyz1 = sample["target"].points trans = torch.Tensor(self.data["T_map"][idx]) matching_search_voxel_size = self.matching_search_voxel_size import copy # sel0 = ME.utils.sparse_quantize( xyz0.contiguous() / self.voxel_size, return_index=True)[1] sel1 = ME.utils.sparse_quantize( xyz1.contiguous() / self.voxel_size, return_index=True)[1] unique_xyz0_th = xyz0[sel0] # [ind_0] unique_xyz1_th = xyz1[sel1] # [ind_1] pcd0 = make_open3d_point_cloud(unique_xyz0_th) pcd1 = make_open3d_point_cloud(unique_xyz1_th) # Get matches matches = get_matching_indices( pcd0, pcd1, trans, matching_search_voxel_size) # Get features feats_train0, feats_train1 = [], [] npts0 = unique_xyz0_th.shape[0] npts1 = unique_xyz1_th.shape[0] feats_train0.append(torch.ones((npts0, 1))) feats_train1.append(torch.ones((npts1, 1))) feats0 = torch.cat(feats_train0, 1) feats1 = torch.cat(feats_train1, 1) coords0 = np.floor(unique_xyz0_th / self.voxel_size) coords1 = np.floor(unique_xyz1_th / self.voxel_size) #coords0_mean = coords0.min(axis=0).int() #coords0 -=coords0_mean if False:#len(matches) < 300: # idx == 113:#len(matches) < print(self.split, "num matches = ", len(matches)) print(coords0) pcd_target = o3d.geometry.PointCloud() pcd_target.points = o3d.utility.Vector3dVector(coords0) o3d.io.write_point_cloud("coords0_before_%d.ply" % idx, pcd_target) # pcd_target = o3d.geometry.PointCloud() pcd_target.points = o3d.utility.Vector3dVector(coords1) o3d.io.write_point_cloud("coords1_before_%d.ply" % idx, pcd_target) # pcd_target = o3d.geometry.PointCloud() pcd_target.points = o3d.utility.Vector3dVector(unique_xyz1_th) o3d.io.write_point_cloud("unique_xyz1_th_%d.ply" % idx, pcd_target) # pcd_target = o3d.geometry.PointCloud() pcd_target.points = o3d.utility.Vector3dVector(unique_xyz0_th) o3d.io.write_point_cloud("unique_xyz0_th_%d.ply" % idx, pcd_target) # pcd_target = o3d.geometry.PointCloud() pcd_target.points = o3d.utility.Vector3dVector(unique_xyz0_th) pcd_target.transform(trans) o3d.io.write_point_cloud( "unique_xyz0_th_trans_%d.ply" % idx, pcd_target) # import pdb pdb.set_trace() #if self.transform: # add noises to the point clouds # coords0, feats0 = self.transform(coords0, feats0) # coords1, feats1 = self.transform(coords1, feats1) return (unique_xyz0_th, unique_xyz1_th, coords0, coords1, feats0.float(), feats1.float(), matches, trans)