Dataloader Wrapper for Pittsburgh Dataset
| 22 | |
| 23 | |
| 24 | class PittsburghDataset(TripletDataLoader): |
| 25 | """Dataloader Wrapper for Pittsburgh Dataset""" |
| 26 | def __init__(self, config, is_train): |
| 27 | super().__init__(config, is_train) |
| 28 | self.dataset_dir = os.path.join(config.DATA.BENCHMARK_DIR, str(config.DATA.DATASET_NAME)) |
| 29 | if config.MODEL.TYPE not in ["Lidar", "LiSPH"]: |
| 30 | raise ValueError("Pittsburgh dataset only provides Lidar data!") |
| 31 | if is_train: |
| 32 | self.generate_pickles() |
| 33 | self.queries = self.get_queries_dict(os.path.join(self.dataset_dir, config.DATA.TRAIN_PICKLE)) |
| 34 | log_print("Number of training tuples: %d" % len(self.queries), "y") |
| 35 | else: |
| 36 | self.queries = self.get_queries_dict(os.path.join(self.dataset_dir, config.DATA.VAL_PICKLE)) |
| 37 | log_print("Number of val tuples: %d" % len(self.queries), "y") |
| 38 | self.file_idxs = np.arange(0, len(self.queries.keys())) |
| 39 | |
| 40 | def load_file_func(self, filename): |
| 41 | if self.config.MODEL.TYPE == "Lidar" or self.branch == "Lidar": |
| 42 | pcd = np.asarray(o3d.io.read_point_cloud(filename + ".pcd").points) |
| 43 | if pcd.shape[0] != 4096: |
| 44 | raise ValueError('{}.pcd does not have sufficient points'.format(filename)) |
| 45 | pcd = pc_normalize(pcd) |
| 46 | output = self.lidar_data_aug(pcd) |
| 47 | elif self.config.MODEL.TYPE == "LiSPH" or self.branch == "LiSPH": |
| 48 | sph_img = Image.open(filename + "_sph.png") |
| 49 | output = self.lisph_data_aug(sph_img) |
| 50 | return output |
| 51 | |
| 52 | def generate_pickles(self): |
| 53 | if self.config.TRAINING.IS_TRAIN: |
| 54 | pair_train = os.path.join( |
| 55 | self.dataset_dir, self.config.DATA.TRAIN_PICKLE) |
| 56 | pair_test = os.path.join( |
| 57 | self.dataset_dir, self.config.DATA.VAL_PICKLE) |
| 58 | #! Load pickles if exist |
| 59 | if os.path.exists(pair_train) and os.path.exists(pair_test): |
| 60 | log_print("Load previous pickles", "b") |
| 61 | return |
| 62 | # ANCHOR generate train/val query for |
| 63 | pair_train = self.get_df(self.dataset_dir, 'train', is_shuffle=False) |
| 64 | pair_train = pair_train.sample(frac=1).reset_index(drop=True) |
| 65 | pair_val = self.get_df(self.dataset_dir, 'val', is_shuffle=False) |
| 66 | pair_val = pair_val.sample(frac=1).reset_index(drop=True) |
| 67 | |
| 68 | self.construct_query_dict(pair_train, self.config.DATA.TRAIN_PICKLE) |
| 69 | self.construct_query_dict(pair_val, self.config.DATA.VAL_PICKLE) |
| 70 | log_print("Generated pickles!\n", "g") |
| 71 | |
| 72 | def get_from_folder(self, data_path, index): |
| 73 | |
| 74 | all_file_id = [] |
| 75 | pose_data = glob(data_path+'/*_pose.npy') |
| 76 | for file_name in pose_data: |
| 77 | all_file_id.append(file_name.split('_pose')[-2]) |
| 78 | all_file_id.sort() |
| 79 | all_data_df = pd.DataFrame(all_file_id, columns=["file"]) |
| 80 | all_data_df["pcd_position_x"] = all_data_df["file"].apply( |
| 81 | lambda x: np.load(x + '_pose.npy')[0]) |