| 129 | return len(self.entrydata_list) |
| 130 | |
| 131 | class MegaDepth_CAPS(PoseDataset): |
| 132 | def __init__(self, args, split='train'): |
| 133 | super(MegaDepth_CAPS, self).__init__(args) |
| 134 | root = os.path.join(args.training_data_dir, args.dataset) |
| 135 | if not os.path.exists(root): |
| 136 | raise colored('[Error: ]', 'red') + 'file path {:s} is not exist!'.format(root) |
| 137 | |
| 138 | scenes_list = sorted(os.listdir(osp.join(root, split))) |
| 139 | scenes_list = [scene for scene in scenes_list if osp.isdir(osp.join(root, split, scene))] |
| 140 | |
| 141 | dense_list = [] |
| 142 | for scene in scenes_list: |
| 143 | denses = os.listdir(osp.join(root, split, scene)) |
| 144 | dense_list += [osp.join(root, split, scene, dense) for dense in denses if osp.isdir(osp.join(root, split, scene, dense))] |
| 145 | |
| 146 | if self.args.debug: |
| 147 | dense_list = dense_list[:1] |
| 148 | |
| 149 | start_time = time.time() |
| 150 | for dense_path in tqdm(dense_list, ncols=70): |
| 151 | im_path = osp.join(dense_path, 'aligned', 'images') |
| 152 | |
| 153 | # ignore empty file's warnings |
| 154 | with warnings.catch_warnings(): |
| 155 | warnings.simplefilter("ignore") |
| 156 | pairs_list = np.loadtxt(osp.join(dense_path, 'aligned', 'pairs.txt'), dtype=str) |
| 157 | pose_list = np.loadtxt(osp.join(dense_path, 'aligned', 'img_cam.txt'), dtype=str) |
| 158 | |
| 159 | pose_map = {} |
| 160 | for item in pose_list: |
| 161 | pose_map.update({item[0]: item[1:].astype(np.float32)}) |
| 162 | |
| 163 | if len(pairs_list) > self.args.sample_maxlen: |
| 164 | index = np.arange(len(pairs_list)) |
| 165 | np.random.shuffle(index) |
| 166 | pairs_list = np.array(pairs_list)[index[:self.args.sample_maxlen]].tolist() |
| 167 | |
| 168 | for im1_name, im2_name in pairs_list: |
| 169 | im1_path = osp.join(im_path, im1_name) |
| 170 | im2_path = osp.join(im_path, im2_name) |
| 171 | Tcw1, K1 = entry_convert(pose_map[im1_name]) |
| 172 | Tcw2, K2 = entry_convert(pose_map[im2_name]) |
| 173 | |
| 174 | # Avoid the different path map to the same image and pose |
| 175 | if np.linalg.norm(Tcw1 - Tcw2) < self.args.epsilon: |
| 176 | continue |
| 177 | |
| 178 | im1 = Image(im1_path, Tcw1.astype(np.float32), K1.astype(np.float32)) |
| 179 | im2 = Image(im2_path, Tcw2.astype(np.float32), K2.astype(np.float32)) |
| 180 | |
| 181 | self.entrydata_list += [[im1, im2]] |
| 182 | end_time = time.time() |
| 183 | |
| 184 | print('\nLoading dataset cost {:.3f}s, data size {} pairs'.format(end_time - start_time, len(self.entrydata_list))) |
| 185 | |
| 186 | |
| 187 | class SequenceDataset(data.Dataset): |