import h5py import numpy as np from PIL import Image import os import torch from torch.utils.data import Dataset import time from tqdm import tqdm from lib.utils import preprocess_image class MegaDepthDataset(Dataset): def __init__( self, scene_list_path='megadepth_utils/train_scenes.txt', scene_info_path='/local/dataset/megadepth/scene_info', base_path='/local/dataset/megadepth', train=True, preprocessing=None, min_overlap_ratio=.5, max_overlap_ratio=1, max_scale_ratio=np.inf, pairs_per_scene=100, image_size=256 ): self.scenes = [] with open(scene_list_path, 'r') as f: lines = f.readlines() for line in lines: self.scenes.append(line.strip('\n')) self.scene_info_path = scene_info_path self.base_path = base_path self.train = train self.preprocessing = preprocessing self.min_overlap_ratio = min_overlap_ratio self.max_overlap_ratio = max_overlap_ratio self.max_scale_ratio = max_scale_ratio self.pairs_per_scene = pairs_per_scene self.image_size = image_size self.dataset = [] def build_dataset(self): self.dataset = [] if not self.train: np_random_state = np.random.get_state() np.random.seed(42) print('Building the validation dataset...') else: print('Building a new training dataset...') for scene in tqdm(self.scenes, total=len(self.scenes)): scene_info_path = os.path.join( self.scene_info_path, '%s.npz' % scene ) if not os.path.exists(scene_info_path): continue scene_info = np.load(scene_info_path, allow_pickle=True) overlap_matrix = scene_info['overlap_matrix'] scale_ratio_matrix = scene_info['scale_ratio_matrix'] valid = np.logical_and( np.logical_and( overlap_matrix >= self.min_overlap_ratio, overlap_matrix <= self.max_overlap_ratio ), scale_ratio_matrix <= self.max_scale_ratio ) pairs = np.vstack(np.where(valid)) try: selected_ids = np.random.choice( pairs.shape[1], self.pairs_per_scene ) except: continue image_paths = scene_info['image_paths'] depth_paths = scene_info['depth_paths'] points3D_id_to_2D = scene_info['points3D_id_to_2D'] points3D_id_to_ndepth = scene_info['points3D_id_to_ndepth'] intrinsics = scene_info['intrinsics'] poses = scene_info['poses'] for pair_idx in selected_ids: idx1 = pairs[0, pair_idx] idx2 = pairs[1, pair_idx] matches = np.array(list( points3D_id_to_2D[idx1].keys() & points3D_id_to_2D[idx2].keys() )) # Scale filtering matches_nd1 = np.array([points3D_id_to_ndepth[idx1][match] for match in matches]) matches_nd2 = np.array([points3D_id_to_ndepth[idx2][match] for match in matches]) scale_ratio = np.maximum(matches_nd1 / matches_nd2, matches_nd2 / matches_nd1) matches = matches[np.where(scale_ratio <= self.max_scale_ratio)[0]] point3D_id = np.random.choice(matches) point2D1 = points3D_id_to_2D[idx1][point3D_id] point2D2 = points3D_id_to_2D[idx2][point3D_id] nd1 = points3D_id_to_ndepth[idx1][point3D_id] nd2 = points3D_id_to_ndepth[idx2][point3D_id] central_match = np.array([ point2D1[1], point2D1[0], point2D2[1], point2D2[0] ]) self.dataset.append({ 'image_path1': image_paths[idx1], 'depth_path1': depth_paths[idx1], 'intrinsics1': intrinsics[idx1], 'pose1': poses[idx1], 'image_path2': image_paths[idx2], 'depth_path2': depth_paths[idx2], 'intrinsics2': intrinsics[idx2], 'pose2': poses[idx2], 'central_match': central_match, 'scale_ratio': max(nd1 / nd2, nd2 / nd1) }) np.random.shuffle(self.dataset) if not self.train: np.random.set_state(np_random_state) def __len__(self): return len(self.dataset) def recover_pair(self, pair_metadata): depth_path1 = os.path.join( self.base_path, pair_metadata['depth_path1'] ) with h5py.File(depth_path1, 'r') as hdf5_file: depth1 = np.array(hdf5_file['/depth']) assert(np.min(depth1) >= 0) image_path1 = os.path.join( self.base_path, pair_metadata['image_path1'] ) image1 = Image.open(image_path1) if image1.mode != 'RGB': image1 = image1.convert('RGB') image1 = np.array(image1) assert(image1.shape[0] == depth1.shape[0] and image1.shape[1] == depth1.shape[1]) intrinsics1 = pair_metadata['intrinsics1'] pose1 = pair_metadata['pose1'] depth_path2 = os.path.join( self.base_path, pair_metadata['depth_path2'] ) with h5py.File(depth_path2, 'r') as hdf5_file: depth2 = np.array(hdf5_file['/depth']) assert(np.min(depth2) >= 0) image_path2 = os.path.join( self.base_path, pair_metadata['image_path2'] ) image2 = Image.open(image_path2) if image2.mode != 'RGB': image2 = image2.convert('RGB') image2 = np.array(image2) assert(image2.shape[0] == depth2.shape[0] and image2.shape[1] == depth2.shape[1]) intrinsics2 = pair_metadata['intrinsics2'] pose2 = pair_metadata['pose2'] central_match = pair_metadata['central_match'] image1, bbox1, image2, bbox2 = self.crop(image1, image2, central_match) depth1 = depth1[ bbox1[0] : bbox1[0] + self.image_size, bbox1[1] : bbox1[1] + self.image_size ] depth2 = depth2[ bbox2[0] : bbox2[0] + self.image_size, bbox2[1] : bbox2[1] + self.image_size ] return ( image1, depth1, intrinsics1, pose1, bbox1, image2, depth2, intrinsics2, pose2, bbox2 ) def crop(self, image1, image2, central_match): bbox1_i = max(int(central_match[0]) - self.image_size // 2, 0) if bbox1_i + self.image_size >= image1.shape[0]: bbox1_i = image1.shape[0] - self.image_size bbox1_j = max(int(central_match[1]) - self.image_size // 2, 0) if bbox1_j + self.image_size >= image1.shape[1]: bbox1_j = image1.shape[1] - self.image_size bbox2_i = max(int(central_match[2]) - self.image_size // 2, 0) if bbox2_i + self.image_size >= image2.shape[0]: bbox2_i = image2.shape[0] - self.image_size bbox2_j = max(int(central_match[3]) - self.image_size // 2, 0) if bbox2_j + self.image_size >= image2.shape[1]: bbox2_j = image2.shape[1] - self.image_size return ( image1[ bbox1_i : bbox1_i + self.image_size, bbox1_j : bbox1_j + self.image_size ], np.array([bbox1_i, bbox1_j]), image2[ bbox2_i : bbox2_i + self.image_size, bbox2_j : bbox2_j + self.image_size ], np.array([bbox2_i, bbox2_j]) ) def __getitem__(self, idx): ( image1, depth1, intrinsics1, pose1, bbox1, image2, depth2, intrinsics2, pose2, bbox2 ) = self.recover_pair(self.dataset[idx]) image1 = preprocess_image(image1, preprocessing=self.preprocessing) image2 = preprocess_image(image2, preprocessing=self.preprocessing) return { 'image1': torch.from_numpy(image1.astype(np.float32)), 'depth1': torch.from_numpy(depth1.astype(np.float32)), 'intrinsics1': torch.from_numpy(intrinsics1.astype(np.float32)), 'pose1': torch.from_numpy(pose1.astype(np.float32)), 'bbox1': torch.from_numpy(bbox1.astype(np.float32)), 'image2': torch.from_numpy(image2.astype(np.float32)), 'depth2': torch.from_numpy(depth2.astype(np.float32)), 'intrinsics2': torch.from_numpy(intrinsics2.astype(np.float32)), 'pose2': torch.from_numpy(pose2.astype(np.float32)), 'bbox2': torch.from_numpy(bbox2.astype(np.float32)) }