import argparse import imagesize import numpy as np import os parser = argparse.ArgumentParser(description="MegaDepth preprocessing script") parser.add_argument("--base_path", type=str, required=True, help="path to MegaDepth") parser.add_argument("--scene_id", type=str, required=True, help="scene ID") parser.add_argument( "--output_path", type=str, required=True, help="path to the output directory" ) args = parser.parse_args() base_path = args.base_path # Remove the trailing / if need be. if base_path[-1] in ["/", "\\"]: base_path = base_path[:-1] scene_id = args.scene_id base_depth_path = os.path.join(base_path, "phoenix/S6/zl548/MegaDepth_v1") base_undistorted_sfm_path = os.path.join(base_path, "Undistorted_SfM") undistorted_sparse_path = os.path.join( base_undistorted_sfm_path, scene_id, "sparse-txt" ) if not os.path.exists(undistorted_sparse_path): exit() depths_path = os.path.join(base_depth_path, scene_id, "dense0", "depths") if not os.path.exists(depths_path): exit() images_path = os.path.join(base_undistorted_sfm_path, scene_id, "images") if not os.path.exists(images_path): exit() # Process cameras.txt with open(os.path.join(undistorted_sparse_path, "cameras.txt"), "r") as f: raw = f.readlines()[3:] # skip the header camera_intrinsics = {} for camera in raw: camera = camera.split(" ") camera_intrinsics[int(camera[0])] = [float(elem) for elem in camera[2:]] # Process points3D.txt with open(os.path.join(undistorted_sparse_path, "points3D.txt"), "r") as f: raw = f.readlines()[3:] # skip the header points3D = {} for point3D in raw: point3D = point3D.split(" ") points3D[int(point3D[0])] = np.array( [float(point3D[1]), float(point3D[2]), float(point3D[3])] ) # Process images.txt with open(os.path.join(undistorted_sparse_path, "images.txt"), "r") as f: raw = f.readlines()[4:] # skip the header image_id_to_idx = {} image_names = [] raw_pose = [] camera = [] points3D_id_to_2D = [] n_points3D = [] for idx, (image, points) in enumerate(zip(raw[::2], raw[1::2])): image = image.split(" ") points = points.split(" ") image_id_to_idx[int(image[0])] = idx image_name = image[-1].strip("\n") image_names.append(image_name) raw_pose.append([float(elem) for elem in image[1:-2]]) camera.append(int(image[-2])) current_points3D_id_to_2D = {} for x, y, point3D_id in zip(points[::3], points[1::3], points[2::3]): if int(point3D_id) == -1: continue current_points3D_id_to_2D[int(point3D_id)] = [float(x), float(y)] points3D_id_to_2D.append(current_points3D_id_to_2D) n_points3D.append(len(current_points3D_id_to_2D)) n_images = len(image_names) # Image and depthmaps paths image_paths = [] depth_paths = [] for image_name in image_names: image_path = os.path.join(images_path, image_name) # Path to the depth file depth_path = os.path.join(depths_path, "%s.h5" % os.path.splitext(image_name)[0]) if os.path.exists(depth_path): # Check if depth map or background / foreground mask file_size = os.stat(depth_path).st_size # Rough estimate - 75KB might work as well if file_size < 100 * 1024: depth_paths.append(None) image_paths.append(None) else: depth_paths.append(depth_path[len(base_path) + 1 :]) image_paths.append(image_path[len(base_path) + 1 :]) else: depth_paths.append(None) image_paths.append(None) # Camera configuration intrinsics = [] poses = [] principal_axis = [] points3D_id_to_ndepth = [] for idx, image_name in enumerate(image_names): if image_paths[idx] is None: intrinsics.append(None) poses.append(None) principal_axis.append([0, 0, 0]) points3D_id_to_ndepth.append({}) continue image_intrinsics = camera_intrinsics[camera[idx]] K = np.zeros([3, 3]) K[0, 0] = image_intrinsics[2] K[0, 2] = image_intrinsics[4] K[1, 1] = image_intrinsics[3] K[1, 2] = image_intrinsics[5] K[2, 2] = 1 intrinsics.append(K) image_pose = raw_pose[idx] qvec = image_pose[:4] qvec = qvec / np.linalg.norm(qvec) w, x, y, z = qvec R = np.array( [ [1 - 2 * y * y - 2 * z * z, 2 * x * y - 2 * z * w, 2 * x * z + 2 * y * w], [2 * x * y + 2 * z * w, 1 - 2 * x * x - 2 * z * z, 2 * y * z - 2 * x * w], [2 * x * z - 2 * y * w, 2 * y * z + 2 * x * w, 1 - 2 * x * x - 2 * y * y], ] ) principal_axis.append(R[2, :]) t = image_pose[4:7] # World-to-Camera pose current_pose = np.zeros([4, 4]) current_pose[:3, :3] = R current_pose[:3, 3] = t current_pose[3, 3] = 1 # Camera-to-World pose # pose = np.zeros([4, 4]) # pose[: 3, : 3] = np.transpose(R) # pose[: 3, 3] = -np.matmul(np.transpose(R), t) # pose[3, 3] = 1 poses.append(current_pose) current_points3D_id_to_ndepth = {} for point3D_id in points3D_id_to_2D[idx].keys(): p3d = points3D[point3D_id] current_points3D_id_to_ndepth[point3D_id] = (np.dot(R[2, :], p3d) + t[2]) / ( 0.5 * (K[0, 0] + K[1, 1]) ) points3D_id_to_ndepth.append(current_points3D_id_to_ndepth) principal_axis = np.array(principal_axis) angles = np.rad2deg( np.arccos(np.clip(np.dot(principal_axis, np.transpose(principal_axis)), -1, 1)) ) # Compute overlap score overlap_matrix = np.full([n_images, n_images], -1.0) scale_ratio_matrix = np.full([n_images, n_images], -1.0) for idx1 in range(n_images): if image_paths[idx1] is None or depth_paths[idx1] is None: continue for idx2 in range(idx1 + 1, n_images): if image_paths[idx2] is None or depth_paths[idx2] is None: continue matches = points3D_id_to_2D[idx1].keys() & points3D_id_to_2D[idx2].keys() min_num_points3D = min( len(points3D_id_to_2D[idx1]), len(points3D_id_to_2D[idx2]) ) overlap_matrix[idx1, idx2] = len(matches) / len( points3D_id_to_2D[idx1] ) # min_num_points3D overlap_matrix[idx2, idx1] = len(matches) / len( points3D_id_to_2D[idx2] ) # min_num_points3D if len(matches) == 0: continue points3D_id_to_ndepth1 = points3D_id_to_ndepth[idx1] points3D_id_to_ndepth2 = points3D_id_to_ndepth[idx2] nd1 = np.array([points3D_id_to_ndepth1[match] for match in matches]) nd2 = np.array([points3D_id_to_ndepth2[match] for match in matches]) min_scale_ratio = np.min(np.maximum(nd1 / nd2, nd2 / nd1)) scale_ratio_matrix[idx1, idx2] = min_scale_ratio scale_ratio_matrix[idx2, idx1] = min_scale_ratio np.savez( os.path.join(args.output_path, "%s.npz" % scene_id), image_paths=image_paths, depth_paths=depth_paths, intrinsics=intrinsics, poses=poses, overlap_matrix=overlap_matrix, scale_ratio_matrix=scale_ratio_matrix, angles=angles, n_points3D=n_points3D, points3D_id_to_2D=points3D_id_to_2D, points3D_id_to_ndepth=points3D_id_to_ndepth, )