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.0, | |
| "img1": img1 / 255.0, | |
| "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 | |