forked from naver/dust3r
-
Notifications
You must be signed in to change notification settings - Fork 0
/
visloc.py
193 lines (170 loc) · 8.98 KB
/
visloc.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
190
191
192
193
#!/usr/bin/env python3
# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# Simple visloc script
# --------------------------------------------------------
import numpy as np
import random
import argparse
from tqdm import tqdm
import math
from dust3r.inference import inference
from dust3r.model import AsymmetricCroCo3DStereo
from dust3r.utils.geometry import find_reciprocal_matches, xy_grid, geotrf
from dust3r_visloc.datasets import *
from dust3r_visloc.localization import run_pnp
from dust3r_visloc.evaluation import get_pose_error, aggregate_stats, export_results
def get_args_parser():
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", type=str, required=True, help="visloc dataset to eval")
parser_weights = parser.add_mutually_exclusive_group(required=True)
parser_weights.add_argument("--weights", type=str, help="path to the model weights", default=None)
parser_weights.add_argument("--model_name", type=str, help="name of the model weights",
choices=["DUSt3R_ViTLarge_BaseDecoder_512_dpt",
"DUSt3R_ViTLarge_BaseDecoder_512_linear",
"DUSt3R_ViTLarge_BaseDecoder_224_linear"])
parser.add_argument("--confidence_threshold", type=float, default=3.0,
help="confidence values higher than threshold are invalid")
parser.add_argument("--device", type=str, default='cuda', help="pytorch device")
parser.add_argument("--pnp_mode", type=str, default="cv2", choices=['cv2', 'poselib', 'pycolmap'],
help="pnp lib to use")
parser_reproj = parser.add_mutually_exclusive_group()
parser_reproj.add_argument("--reprojection_error", type=float, default=5.0, help="pnp reprojection error")
parser_reproj.add_argument("--reprojection_error_diag_ratio", type=float, default=None,
help="pnp reprojection error as a ratio of the diagonal of the image")
parser.add_argument("--pnp_max_points", type=int, default=100_000, help="pnp maximum number of points kept")
parser.add_argument("--viz_matches", type=int, default=0, help="debug matches")
parser.add_argument("--output_dir", type=str, default=None, help="output path")
parser.add_argument("--output_label", type=str, default='', help="prefix for results files")
return parser
if __name__ == '__main__':
parser = get_args_parser()
args = parser.parse_args()
conf_thr = args.confidence_threshold
device = args.device
pnp_mode = args.pnp_mode
reprojection_error = args.reprojection_error
reprojection_error_diag_ratio = args.reprojection_error_diag_ratio
pnp_max_points = args.pnp_max_points
viz_matches = args.viz_matches
if args.weights is not None:
weights_path = args.weights
else:
weights_path = "naver/" + args.model_name
model = AsymmetricCroCo3DStereo.from_pretrained(weights_path).to(args.device)
dataset = eval(args.dataset)
dataset.set_resolution(model)
query_names = []
poses_pred = []
pose_errors = []
angular_errors = []
for idx in tqdm(range(len(dataset))):
views = dataset[(idx)] # 0 is the query
query_view = views[0]
map_views = views[1:]
query_names.append(query_view['image_name'])
query_pts2d = []
query_pts3d = []
for map_view in map_views:
# prepare batch
imgs = []
for idx, img in enumerate([query_view['rgb_rescaled'], map_view['rgb_rescaled']]):
imgs.append(dict(img=img.unsqueeze(0), true_shape=np.int32([img.shape[1:]]),
idx=idx, instance=str(idx)))
output = inference([tuple(imgs)], model, device, batch_size=1, verbose=False)
pred1, pred2 = output['pred1'], output['pred2']
confidence_masks = [pred1['conf'].squeeze(0) >= conf_thr,
(pred2['conf'].squeeze(0) >= conf_thr) & map_view['valid_rescaled']]
pts3d = [pred1['pts3d'].squeeze(0), pred2['pts3d_in_other_view'].squeeze(0)]
# find 2D-2D matches between the two images
pts2d_list, pts3d_list = [], []
for i in range(2):
conf_i = confidence_masks[i].cpu().numpy()
true_shape_i = imgs[i]['true_shape'][0]
pts2d_list.append(xy_grid(true_shape_i[1], true_shape_i[0])[conf_i])
pts3d_list.append(pts3d[i].detach().cpu().numpy()[conf_i])
PQ, PM = pts3d_list[0], pts3d_list[1]
if len(PQ) == 0 or len(PM) == 0:
continue
reciprocal_in_PM, nnM_in_PQ, num_matches = find_reciprocal_matches(PQ, PM)
if viz_matches > 0:
print(f'found {num_matches} matches')
matches_im1 = pts2d_list[1][reciprocal_in_PM]
matches_im0 = pts2d_list[0][nnM_in_PQ][reciprocal_in_PM]
valid_pts3d = map_view['pts3d_rescaled'][matches_im1[:, 1], matches_im1[:, 0]]
# from cv2 to colmap
matches_im0 = matches_im0.astype(np.float64)
matches_im1 = matches_im1.astype(np.float64)
matches_im0[:, 0] += 0.5
matches_im0[:, 1] += 0.5
matches_im1[:, 0] += 0.5
matches_im1[:, 1] += 0.5
# rescale coordinates
matches_im0 = geotrf(query_view['to_orig'], matches_im0, norm=True)
matches_im1 = geotrf(query_view['to_orig'], matches_im1, norm=True)
# from colmap back to cv2
matches_im0[:, 0] -= 0.5
matches_im0[:, 1] -= 0.5
matches_im1[:, 0] -= 0.5
matches_im1[:, 1] -= 0.5
# visualize a few matches
if viz_matches > 0:
viz_imgs = [np.array(query_view['rgb']), np.array(map_view['rgb'])]
from matplotlib import pyplot as pl
n_viz = viz_matches
match_idx_to_viz = np.round(np.linspace(0, num_matches - 1, n_viz)).astype(int)
viz_matches_im0, viz_matches_im1 = matches_im0[match_idx_to_viz], matches_im1[match_idx_to_viz]
H0, W0, H1, W1 = *viz_imgs[0].shape[:2], *viz_imgs[1].shape[:2]
img0 = np.pad(viz_imgs[0], ((0, max(H1 - H0, 0)), (0, 0), (0, 0)), 'constant', constant_values=0)
img1 = np.pad(viz_imgs[1], ((0, max(H0 - H1, 0)), (0, 0), (0, 0)), 'constant', constant_values=0)
img = np.concatenate((img0, img1), axis=1)
pl.figure()
pl.imshow(img)
cmap = pl.get_cmap('jet')
for i in range(n_viz):
(x0, y0), (x1, y1) = viz_matches_im0[i].T, viz_matches_im1[i].T
pl.plot([x0, x1 + W0], [y0, y1], '-+', color=cmap(i / (n_viz - 1)), scalex=False, scaley=False)
pl.show(block=True)
if len(valid_pts3d) == 0:
pass
else:
query_pts3d.append(valid_pts3d.cpu().numpy())
query_pts2d.append(matches_im0)
if len(query_pts2d) == 0:
success = False
pr_querycam_to_world = None
else:
query_pts2d = np.concatenate(query_pts2d, axis=0).astype(np.float32)
query_pts3d = np.concatenate(query_pts3d, axis=0)
if len(query_pts2d) > pnp_max_points:
idxs = random.sample(range(len(query_pts2d)), pnp_max_points)
query_pts3d = query_pts3d[idxs]
query_pts2d = query_pts2d[idxs]
W, H = query_view['rgb'].size
if reprojection_error_diag_ratio is not None:
reprojection_error_img = reprojection_error_diag_ratio * math.sqrt(W**2 + H**2)
else:
reprojection_error_img = reprojection_error
success, pr_querycam_to_world = run_pnp(query_pts2d, query_pts3d,
query_view['intrinsics'], query_view['distortion'],
pnp_mode, reprojection_error_img, img_size=[W, H])
if not success:
abs_transl_error = float('inf')
abs_angular_error = float('inf')
else:
abs_transl_error, abs_angular_error = get_pose_error(pr_querycam_to_world, query_view['cam_to_world'])
pose_errors.append(abs_transl_error)
angular_errors.append(abs_angular_error)
poses_pred.append(pr_querycam_to_world)
xp_label = f'tol_conf_{conf_thr}'
if args.output_label:
xp_label = args.output_label + '_' + xp_label
if reprojection_error_diag_ratio is not None:
xp_label = xp_label + f'_reproj_diag_{reprojection_error_diag_ratio}'
else:
xp_label = xp_label + f'_reproj_err_{reprojection_error}'
export_results(args.output_dir, xp_label, query_names, poses_pred)
out_string = aggregate_stats(f'{args.dataset}', pose_errors, angular_errors)
print(out_string)