Spaces:
Sleeping
Sleeping
File size: 3,039 Bytes
6ba5875 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
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") |