-
Notifications
You must be signed in to change notification settings - Fork 1
/
matching.py
189 lines (151 loc) · 7.11 KB
/
matching.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
186
187
188
189
from tqdm import tqdm
import torch
import numpy as np
from lightglue import LightGlue
from lightglue.utils import rbd
import joblib
from scipy.cluster.vq import vq
from numpy.linalg import norm
import cv2
all_descriptors = np.load("output/all_descriptors.npy", allow_pickle=True)
all_points = np.load("output/all_points.npy", allow_pickle=True)
img_size = np.load("output/img_size.npy", allow_pickle=True)
k, codebook = joblib.load("output/bow_codebook.plk")
torch.set_grad_enabled(False)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
matcher = LightGlue(features='disk').eval().to(device)
print("Build Pairs")
visual_words = []
for desciptors in all_descriptors:
img_visual_words, distance = vq(desciptors.astype("float"), codebook)
visual_words.append(img_visual_words)
frequency_vectors = []
for img_visual_words in visual_words:
img_frequency_vector = np.zeros(k)
for word in img_visual_words:
img_frequency_vector[word] += 1
frequency_vectors.append(img_frequency_vector)
frequency_vectors = np.stack(frequency_vectors)
print(frequency_vectors.shape)
N = frequency_vectors.shape[0]
df = np.sum(frequency_vectors > 0, axis = 0 )
idf = np.log(N/ df)
tfidf = frequency_vectors * idf
all_idx = []
all_score = []
top_k = 10
for i in range(N):
a = tfidf[i]
b = tfidf
cosine_similarity = np.dot(a, b.T)/(norm(a) * norm(b, axis=1))
idx = np.argsort(-cosine_similarity)[1:top_k]
score = np.sort(-cosine_similarity)[1:top_k]
all_idx.append(idx)
all_score.append(score)
connection = [None]*N
for i in range(N):
for j, id in enumerate(all_idx[i]):
if not connection[i]:
connection[i] = []
if not connection[id]:
connection[id] = []
if -all_score[i][j] > 0.75:
if not id in connection[i]:
connection[i].append(id)
if not i in connection[id]:
connection[id].append(i)
print(connection)
max = 0
start = 0
for i, c in enumerate(connection):
if len(c) > max:
max = len(c)
start = i
point3d_index = 0
all_matches = []
all_points3d = [None]*all_points.shape[0]
queue = [(start, start)]
visited = [False]*N
visited[start] = True
i = 0
focal_length = 2378.98305085
with tqdm(total=N) as pbar:
while True:
for id in connection[queue[i][1]]:
if not visited[id]:
reference_id = queue[i][1]
for id_ in connection[id]:
if id_ == queue[i][1]:
break
if visited[id_]:
reference_id = id_
break
feats0 = {
"keypoints": torch.tensor(
np.array([[[p[0] + img_size[reference_id][0]/2, -p[1] + img_size[reference_id][1]/2] for p in all_points[reference_id]]], dtype=float), dtype=torch.float
).to(device),
"descriptors": torch.tensor(np.array([all_descriptors[reference_id]], dtype=float), dtype=torch.float).to(device),
'image_size': torch.tensor(np.array([img_size[reference_id]], dtype=float), dtype=torch.float).to(device)
}
feats1 = {
"keypoints": torch.tensor(
np.array([[[p[0] + img_size[id][0]/2, -p[1] + img_size[id][1]/2] for p in all_points[id]]], dtype=float), dtype=torch.float
).to(device),
"descriptors": torch.tensor(np.array([all_descriptors[id]], dtype=float), dtype=torch.float).to(device),
'image_size': torch.tensor(np.array([img_size[id]], dtype=float), dtype=torch.float).to(device)
}
matches01 = matcher({'image0': feats0, 'image1': feats1})
feats0, feats1, matches01 = [rbd(x) for x in [feats0, feats1, matches01]] # remove batch dimension
kpts0, kpts1, matches = torch.tensor(all_points[reference_id]).to(device), torch.tensor(all_points[id]).to(device), matches01['matches']
m_kpts0, m_kpts1 = kpts0[matches[..., 0]], kpts1[matches[..., 1]]
m_kpts0 = m_kpts0.detach().cpu().numpy().astype(np.float32)
m_kpts1 = m_kpts1.detach().cpu().numpy().astype(np.float32)
idx0, idx1 = matches[..., 0].detach().cpu().numpy(), matches[..., 1].detach().cpu().numpy()
if len(m_kpts0) <= 8:
continue
K = np.array([[focal_length, 0, 0], [0, focal_length, 0], [0, 0, 1]])
E, mask = cv2.findEssentialMat(m_kpts0, m_kpts1, K, method=cv2.RANSAC, prob=0.999, threshold=1)
if mask is None:
continue
_, _, _, mask_inliers = cv2.recoverPose(E, m_kpts0[mask.ravel() > 0], m_kpts1[mask.ravel() > 0], K)
m_kpts0 = m_kpts0[mask.ravel() > 0]
m_kpts0 = m_kpts0[mask_inliers.ravel() > 0]
if len(m_kpts0) > 10:
interlaced_points = 0
for p1, p2 in zip(idx0, idx1):
if not all_points3d[reference_id]:
all_points3d[reference_id] = [-1]*all_points[reference_id].shape[0]
if not all_points3d[id]:
all_points3d[id] = [-1]*all_points[id].shape[0]
if all_points3d[reference_id][p1] == -1 and all_points3d[id][p2] == -1:
continue
elif all_points3d[reference_id][p1] != -1:
interlaced_points += 1
elif all_points3d[id][p1] != -1:
interlaced_points += 1
if len(idx0) >= 500 and (queue[i][1] == start or (queue[i][1] != start and interlaced_points/len(idx0) >= 0.3)):
point3d_indexes = []
for p1, p2 in zip(idx0, idx1):
if all_points3d[reference_id][p1] == -1 and all_points3d[id][p2] == -1:
all_points3d[reference_id][p1] = point3d_index
all_points3d[id][p2] = point3d_index
point3d_index += 1
elif all_points3d[reference_id][p1] != -1:
all_points3d[id][p2] = all_points3d[reference_id][p1]
elif all_points3d[id][p1] != -1:
all_points3d[reference_id][p2] = all_points3d[id][p1]
point3d_indexes.append(all_points3d[reference_id][p1])
all_matches.append([idx0, idx1, np.array(point3d_indexes)])
queue.append((reference_id, id))
visited[id] = True
else:
continue
else:
continue
i += 1
pbar.update(1)
if i >= len(queue):
break
print(queue, len(queue))
np.save('output/img_pairs.npy', queue[1:])
np.save('output/all_matches.npy', np.array(all_matches, dtype=object))