Spaces:
Running
Running
| 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]) / (.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.) | |
| scale_ratio_matrix = np.full([n_images, n_images], -1.) | |
| 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 | |
| ) | |