| 110 | return file_df |
| 111 | |
| 112 | def construct_query_dict(self, data_df, filename): |
| 113 | data_df.reset_index(drop=True, inplace=True) |
| 114 | |
| 115 | tree = KDTree( |
| 116 | data_df[["pcd_position_x", "pcd_position_y", "pcd_position_z"]]) |
| 117 | ind_nn = tree.query_radius(data_df[["pcd_position_x", "pcd_position_y", "pcd_position_z"]], |
| 118 | r=self.config.DATA.POSITIVES_RADIUS) |
| 119 | ind_r = tree.query_radius(data_df[["pcd_position_x", "pcd_position_y", "pcd_position_z"]], |
| 120 | r=self.config.DATA.NEGATIVES_RADIUS) |
| 121 | ind_traj = tree.query_radius(data_df[["pcd_position_x", "pcd_position_y", "pcd_position_z"]], |
| 122 | r=self.config.DATA.TRAJ_RADIUS) |
| 123 | |
| 124 | queries = {} |
| 125 | for i in tqdm(range(len(ind_nn)), total=len(ind_nn), desc='construct queries', leave=False): |
| 126 | query = data_df.iloc[i]["file"] |
| 127 | positives = np.setdiff1d(ind_nn[i], [i]).tolist() |
| 128 | negatives = np.setdiff1d(ind_traj[i], ind_r[i]).tolist() |
| 129 | |
| 130 | random.shuffle(negatives) |
| 131 | random.shuffle(positives) |
| 132 | |
| 133 | queries[i] = {"query": query, "positives": positives, "negatives": negatives} |
| 134 | |
| 135 | with open(os.path.join(self.dataset_dir, filename), 'wb') as handle: |
| 136 | pickle.dump(queries, handle, protocol=pickle.HIGHEST_PROTOCOL) |