3D_Reconstruction/matching.py at master · daovietanh190499/3D_Reconstruction · GitHub

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))