| 16 | |
| 17 | |
| 18 | class Matcher: |
| 19 | image_mean = torch.tensor([0.5, 0.5, 0.5]).reshape(1, 3, 1, 1) |
| 20 | image_std = torch.tensor([0.5, 0.5, 0.5]).reshape(1, 3, 1, 1) |
| 21 | |
| 22 | def __init__(self, device, model_name=None): |
| 23 | self.device = device |
| 24 | self.fast_nn_params = dict(device=device, dist='dot', block_size=2**13) |
| 25 | self.model = AsymmetricMASt3R.from_pretrained(model_name or "naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric").to(device).eval() |
| 26 | self.maxdim = max(self.model.patch_embed.img_size) |
| 27 | self.cmap = plt.get_cmap('jet') |
| 28 | |
| 29 | def match_coarse_to_fine(self, query, reference, pts3d, intrinsics, max_match=100000, conf_thr=-1, vis_fig=None): |
| 30 | coarse_matches_im0, coarse_matches_im1 = self.match(query, reference, conf_thr) |
| 31 | |
| 32 | query_pil = Image.fromarray(query.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()) |
| 33 | query_rgb_tensor, query_K, query_to_orig_max, query_to_resize_max, (HQ, WQ) = resize_image_to_max( |
| 34 | None, query_pil, intrinsics) |
| 35 | reference_pil = Image.fromarray(reference.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()) |
| 36 | map_rgb_tensor, map_K, map_to_orig_max, map_to_resize_max, (HM, WM) = resize_image_to_max( |
| 37 | None, reference_pil, intrinsics) |
| 38 | WM_full, HM_full = reference_pil.size |
| 39 | valid_all = torch.ones((HM_full, WM_full), dtype=torch.bool, device=self.device) |
| 40 | if WM_full != WM or HM_full != HM: |
| 41 | height, width, _ = pts3d.shape |
| 42 | y_full, x_full = torch.meshgrid(torch.arange(height), torch.arange(width)) |
| 43 | pos2d_cv2 = torch.stack([x_full, y_full], dim=-1).numpy().astype(np.float64) |
| 44 | _, _, pts3d_max, valid_max = rescale_points3d(pos2d_cv2, pts3d, map_to_resize_max, HM, WM) |
| 45 | pts3d = torch.from_numpy(pts3d_max) |
| 46 | valid_all = torch.from_numpy(valid_max) |
| 47 | |
| 48 | coarse_matches_im0 = geotrf(query_to_resize_max, coarse_matches_im0, norm=True) |
| 49 | coarse_matches_im1 = geotrf(map_to_resize_max, coarse_matches_im1, norm=True) |
| 50 | |
| 51 | crops1, crops2 = [], [] |
| 52 | crops_v1, crops_p1 = [], [] |
| 53 | to_orig1, to_orig2 = [], [] |
| 54 | map_resolution = get_HW_resolution(HM, WM, maxdim=self.maxdim, patchsize=self.model.patch_embed.patch_size) |
| 55 | query_resolution = get_HW_resolution(HQ, WQ, maxdim=self.maxdim, patchsize=self.model.patch_embed.patch_size) |
| 56 | for crop_q, crop_b, pair_tag in select_pairs_of_crops(map_rgb_tensor, |
| 57 | query_rgb_tensor, |
| 58 | coarse_matches_im1, |
| 59 | coarse_matches_im0, |
| 60 | maxdim=self.maxdim, |
| 61 | overlap=.5, |
| 62 | forced_resolution=[map_resolution, |
| 63 | query_resolution]): |
| 64 | # Per crop processing |
| 65 | c1, v1, p1, trf1 = crop(map_rgb_tensor, valid_all, pts3d, crop_q, None) |
| 66 | c2, _, _, trf2 = crop(query_rgb_tensor, None, None, crop_b, None) |
| 67 | crops1.append(c1) |
| 68 | crops2.append(c2) |
| 69 | crops_v1.append(v1) |
| 70 | crops_p1.append(p1) |
| 71 | to_orig1.append(trf1) |
| 72 | to_orig2.append(trf2) |
| 73 | |
| 74 | if len(crops1) == 0 or len(crops2) == 0: |
| 75 | valid_pts3d, matches_im_query, matches_im_map, matches_conf = [], [], [], [] |