Spaces:
Running
Running
| import os | |
| import numpy as np | |
| import torch | |
| from torch.utils.data import Dataset | |
| from random import shuffle, seed | |
| from .gl3d.io import read_list, _parse_img, _parse_depth, _parse_kpts | |
| from .utils.common import Notify | |
| from .utils.photaug import photaug | |
| class GL3DDataset(Dataset): | |
| def __init__(self, dataset_dir, config, data_split, is_training): | |
| self.dataset_dir = dataset_dir | |
| self.config = config | |
| self.is_training = is_training | |
| self.data_split = data_split | |
| self.match_set_list, self.global_img_list, \ | |
| self.global_depth_list = self.prepare_match_sets() | |
| pass | |
| def __len__(self): | |
| return len(self.match_set_list) | |
| def __getitem__(self, idx): | |
| match_set_path = self.match_set_list[idx] | |
| decoded = np.fromfile(match_set_path, dtype=np.float32) | |
| idx0, idx1 = int(decoded[0]), int(decoded[1]) | |
| inlier_num = int(decoded[2]) | |
| ori_img_size0 = np.reshape(decoded[3:5], (2,)) | |
| ori_img_size1 = np.reshape(decoded[5:7], (2,)) | |
| K0 = np.reshape(decoded[7:16], (3, 3)) | |
| K1 = np.reshape(decoded[16:25], (3, 3)) | |
| rel_pose = np.reshape(decoded[34:46], (3, 4)) | |
| # parse images. | |
| img0 = _parse_img(self.global_img_list, idx0, self.config) | |
| img1 = _parse_img(self.global_img_list, idx1, self.config) | |
| # parse depths | |
| depth0 = _parse_depth(self.global_depth_list, idx0, self.config) | |
| depth1 = _parse_depth(self.global_depth_list, idx1, self.config) | |
| # photometric augmentation | |
| img0 = photaug(img0) | |
| img1 = photaug(img1) | |
| return { | |
| 'img0': img0 / 255., | |
| 'img1': img1 / 255., | |
| 'depth0': depth0, | |
| 'depth1': depth1, | |
| 'ori_img_size0': ori_img_size0, | |
| 'ori_img_size1': ori_img_size1, | |
| 'K0': K0, | |
| 'K1': K1, | |
| 'rel_pose': rel_pose, | |
| 'inlier_num': inlier_num | |
| } | |
| def points_to_2D(self, pnts, H, W): | |
| labels = np.zeros((H, W)) | |
| pnts = pnts.astype(int) | |
| labels[pnts[:, 1], pnts[:, 0]] = 1 | |
| return labels | |
| def prepare_match_sets(self, q_diff_thld=3, rot_diff_thld=60): | |
| """Get match sets. | |
| Args: | |
| is_training: Use training imageset or testing imageset. | |
| data_split: Data split name. | |
| Returns: | |
| match_set_list: List of match sets path. | |
| global_img_list: List of global image path. | |
| global_context_feat_list: | |
| """ | |
| # get necessary lists. | |
| gl3d_list_folder = os.path.join(self.dataset_dir, 'list', self.data_split) | |
| global_info = read_list(os.path.join( | |
| gl3d_list_folder, 'image_index_offset.txt')) | |
| global_img_list = [os.path.join(self.dataset_dir, i) for i in read_list( | |
| os.path.join(gl3d_list_folder, 'image_list.txt'))] | |
| global_depth_list = [os.path.join(self.dataset_dir, i) for i in read_list( | |
| os.path.join(gl3d_list_folder, 'depth_list.txt'))] | |
| imageset_list_name = 'imageset_train.txt' if self.is_training else 'imageset_test.txt' | |
| match_set_list = self.get_match_set_list(os.path.join( | |
| gl3d_list_folder, imageset_list_name), q_diff_thld, rot_diff_thld) | |
| return match_set_list, global_img_list, global_depth_list | |
| def get_match_set_list(self, imageset_list_path, q_diff_thld, rot_diff_thld): | |
| """Get the path list of match sets. | |
| Args: | |
| imageset_list_path: Path to imageset list. | |
| q_diff_thld: Threshold of image pair sampling regarding camera orientation. | |
| Returns: | |
| match_set_list: List of match set path. | |
| """ | |
| imageset_list = [os.path.join(self.dataset_dir, 'data', i) | |
| for i in read_list(imageset_list_path)] | |
| print(Notify.INFO, 'Use # imageset', len(imageset_list), Notify.ENDC) | |
| match_set_list = [] | |
| # discard image pairs whose image simiarity is beyond the threshold. | |
| for i in imageset_list: | |
| match_set_folder = os.path.join(i, 'match_sets') | |
| if os.path.exists(match_set_folder): | |
| match_set_files = os.listdir(match_set_folder) | |
| for val in match_set_files: | |
| name, ext = os.path.splitext(val) | |
| if ext == '.match_set': | |
| splits = name.split('_') | |
| q_diff = int(splits[2]) | |
| rot_diff = int(splits[3]) | |
| if q_diff >= q_diff_thld and rot_diff <= rot_diff_thld: | |
| match_set_list.append( | |
| os.path.join(match_set_folder, val)) | |
| print(Notify.INFO, 'Get # match sets', len(match_set_list), Notify.ENDC) | |
| return match_set_list | |