Spaces:
Running
Running
| import argparse | |
| import imagesize | |
| import numpy as np | |
| import os | |
| base_path = "data/megadepth" | |
| # Remove the trailing / if need be. | |
| if base_path[-1] in ["/", "\\"]: | |
| base_path = base_path[:-1] | |
| base_depth_path = os.path.join(base_path, "phoenix/S6/zl548/MegaDepth_v1") | |
| base_undistorted_sfm_path = os.path.join(base_path, "Undistorted_SfM") | |
| scene_ids = os.listdir(base_undistorted_sfm_path) | |
| for scene_id in scene_ids: | |
| if os.path.exists( | |
| f"{base_path}/prep_scene_info/detections/detections_{scene_id}.npy" | |
| ): | |
| print(f"skipping {scene_id} as it exists") | |
| continue | |
| undistorted_sparse_path = os.path.join( | |
| base_undistorted_sfm_path, scene_id, "sparse-txt" | |
| ) | |
| if not os.path.exists(undistorted_sparse_path): | |
| print("sparse path doesnt exist") | |
| continue | |
| depths_path = os.path.join(base_depth_path, scene_id, "dense0", "depths") | |
| if not os.path.exists(depths_path): | |
| print("depths doesnt exist") | |
| continue | |
| images_path = os.path.join(base_undistorted_sfm_path, scene_id, "images") | |
| if not os.path.exists(images_path): | |
| print("images path doesnt exist") | |
| continue | |
| # Process cameras.txt | |
| if not os.path.exists(os.path.join(undistorted_sparse_path, "cameras.txt")): | |
| print("no cameras") | |
| continue | |
| 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 = [] | |
| id_to_detections = {} | |
| 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])) | |
| points_np = np.array(points).astype(np.float32).reshape(len(points) // 3, 3) | |
| visible_points = points_np[points_np[:, 2] != -1] | |
| id_to_detections[idx] = visible_points | |
| np.save( | |
| f"{base_path}/prep_scene_info/detections/detections_{scene_id}.npy", | |
| id_to_detections, | |
| ) | |
| print(f"{scene_id} done") | |