#!/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)