-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathexport_colmap_matches.py
185 lines (144 loc) · 6.02 KB
/
export_colmap_matches.py
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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
from path import Path
import numpy as np
import sqlite3
import os
def pair_id_to_image_ids(pair_id):
image_id2 = pair_id % 2147483647
image_id1 = (pair_id - image_id2) / 2147483647
return image_id1 - 1, image_id2 - 1
def image_ids_to_pair(image_id1, image_id2):
image_id1 += 1
image_id2 += 1
pair_id = image_id2 + 2147483647 * image_id1
return pair_id
def get_keypoints(cursor, image_id):
image_id += 1
cursor.execute("SELECT * FROM keypoints WHERE image_id = ?;", (image_id,))
image_idx, n_rows, n_columns, raw_data = cursor.fetchone()
kypnts = np.frombuffer(raw_data, dtype=np.float32).reshape(n_rows, n_columns).copy()
kypnts = kypnts[:, 0:2]
return kypnts
def process_one_scene(scene_dir):
filename_db = Path(scene_dir)/'database.db'
outdir = scene_dir/'colmap_matches'
image = scene_dir/'image'
print("Opening database: " + filename_db)
if not os.path.exists(filename_db):
print('Error db does not exist!')
exit()
if not os.path.exists(image):
print('Error image does not exist!')
exit()
if not os.path.exists(outdir):
os.mkdir(outdir)
print(f'Clean old matches in {outdir}')
for f in Path(outdir).files('*'):
os.remove(f)
connection = sqlite3.connect(filename_db)
cursor = connection.cursor()
list_image_ids = []
img_ids_to_names_dict = {}
cursor.execute('SELECT image_id, name, cameras.width, cameras.height FROM images LEFT JOIN cameras ON images.camera_id == cameras.camera_id;')
for row in cursor:
image_idx, name, width, height = row
list_image_ids.append(image_idx - 1)
img_ids_to_names_dict[image_idx - 1] = name
num_image_ids = len(list_image_ids)
# print(Path(image).files('*'))
exist = [1] * (num_image_ids)
for i in img_ids_to_names_dict:
if os.path.join(image, img_ids_to_names_dict[i]) not in Path(image).files('*'):
# print(os.path.join(image, img_ids_to_names_dict[i]))
exist[i] = 0
# Iterate over entries in the two-view geometry table
cursor.execute('SELECT pair_id, rows, cols, data FROM two_view_geometries;')
all_matches = {}
for row in cursor:
pair_id = row[0]
rows = row[1]
cols = row[2]
raw_data = row[3]
if (rows < 5):
continue
matches = np.frombuffer(raw_data, dtype=np.uint32).reshape(rows, cols)
if matches.shape[0] < 5:
continue
all_matches[pair_id] = matches
# print(len(all_matches))
minus = [0] * (num_image_ids)
for j, x in enumerate(exist):
if not x:
for i in range(j, num_image_ids):
minus[i] += 1
print(exist, minus)
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='new_images'")
table_exists = cursor.fetchone()
if not table_exists:
cursor.execute('CREATE TABLE new_images AS SELECT * FROM images')
connection.commit()
for image_id, value_to_subtract in enumerate(minus):
if not exist[image_id]:
cursor.execute('''
DELETE FROM new_images
WHERE image_id = ?
''', (image_id + 1,))
connection.commit()
else:
cursor.execute('''
UPDATE new_images
SET image_id = image_id - 1 - ?
WHERE image_id = ?
''', (value_to_subtract, image_id + 1))
connection.commit()
for key in all_matches:
pair_id = key
matches = all_matches[key]
# # skip if too few matches are given
# if matches.shape[0] < 300:
# continue
id1, id2 = pair_id_to_image_ids(pair_id)
image_name1 = img_ids_to_names_dict[id1]
image_name2 = img_ids_to_names_dict[id2]
keys1 = get_keypoints(cursor, id1)
keys2 = get_keypoints(cursor, id2)
match_positions = np.empty([matches.shape[0], 4])
for i in range(0, matches.shape[0]):
match_positions[i, :] = np.array([keys1[matches[i, 0]][0], keys1[matches[i, 0]][1], keys2[matches[i, 1]][0], keys2[matches[i, 1]][1]])
# print(match_positions.shape)
# outfile = os.path.join(outdir, image_name1.split("/")[0].split(".jpg")[0] + "_" + image_name2.split("/")[0].split(".jpg")[0] + ".txt")
if exist[int(id1)] and exist[int(id2)]:
id1 = id1 - minus[int(id1)]
id2 = id2 - minus[int(id2)]
outfile = os.path.join(outdir, '{:06d}_{:06d}.txt'.format(int(id1), int(id2)))
np.savetxt(outfile, match_positions, delimiter=' ')
# reverse and save
match_positions_reverse = np.concatenate([match_positions[:, 2:4], match_positions[:, 0:2]], axis=1)
# outfile = os.path.join(outdir, image_name2.split("/")[0].split(".jpg")[0] + "_" + image_name1.split("/")[0].split(".jpg")[0] + ".txt")
outfile = os.path.join(outdir, '{:06d}_{:06d}.txt'.format(int(id2), int(id1)))
np.savetxt(outfile, match_positions_reverse, delimiter=' ')
cursor.close()
connection.close()
for idx in range(num_image_ids):
if not exist[idx]:
continue
idx2 = idx - minus[idx]
two_view = {}
two_view["src_idx"] = []
two_view["match"] = []
files = sorted(outdir.files('{:06d}_*.txt'.format(idx2)))
for f in files:
j = int(os.path.basename(f)[7:13])
one_pair = np.loadtxt(f)
# one_list = []
# for i in one_pair:
# # print(list(i))
# one_list.append(list(i))
# print("\n")
# print(len(one_list))
two_view["src_idx"].append(j)
# two_view["match"].append(one_list)
two_view["match"].append(one_pair)
# print(two_view)
np.savez(outdir/'{:06d}.npz'.format(idx2), **two_view)
def export_colmap_matches(scene_dir):
process_one_scene(scene_dir)