File size: 3,701 Bytes
0a82b18 |
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 |
import os
import torchvision.transforms as tfm
import py3_wget
import numpy as np
import torch
import torch.nn.functional as F
from matching import BaseMatcher, WEIGHTS_DIR, THIRD_PARTY_DIR
from matching.utils import resize_to_divisible, add_to_path
add_to_path(THIRD_PARTY_DIR.joinpath("mast3r"))
from typing import Tuple, Union, List, Optional
from mast3r.model import AsymmetricMASt3R
from mast3r.fast_nn import fast_reciprocal_NNs
from dust3r.inference import inference
class Mast3rMatcher(BaseMatcher):
model_path = WEIGHTS_DIR.joinpath(
"MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth"
)
vit_patch_size = 16
def __init__(self, device="cpu", max_num_keypoints=2048, *args, **kwargs):
super().__init__(device, **kwargs)
self.normalize = tfm.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
self.verbose = False
self.max_keypoints = max_num_keypoints
self.download_weights()
self.model = AsymmetricMASt3R.from_pretrained(self.model_path).to(device)
@staticmethod
def download_weights():
url = "https://download.europe.naverlabs.com/ComputerVision/MASt3R/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth"
if not os.path.isfile(Mast3rMatcher.model_path):
print("Downloading Master(ViT large)... (takes a while)")
py3_wget.download_file(url, Mast3rMatcher.model_path)
def preprocess(self, img):
_, h, w = img.shape
orig_shape = h, w
img = resize_to_divisible(img, self.vit_patch_size)
img = self.normalize(img).unsqueeze(0)
return img, orig_shape
def _forward(self, img0, img1):
img0, img0_orig_shape = self.preprocess(img0)
img1, img1_orig_shape = self.preprocess(img1)
img_pair = [
{
"img": img0,
"idx": 0,
"instance": 0,
"true_shape": np.int32([img0.shape[-2:]]),
},
{
"img": img1,
"idx": 1,
"instance": 1,
"true_shape": np.int32([img1.shape[-2:]]),
},
]
output = inference(
[tuple(img_pair)], self.model, self.device, batch_size=1, verbose=False
)
view1, pred1 = output["view1"], output["pred1"]
view2, pred2 = output["view2"], output["pred2"]
desc1, desc2 = (
pred1["desc"].squeeze(0).detach(),
pred2["desc"].squeeze(0).detach(),
)
matches_im0, matches_im1 = fast_reciprocal_NNs(
desc1,
desc2,
subsample_or_initxy1=8,
device=self.device,
dist="dot",
block_size=2**13,
max_matches=self.max_keypoints,
)
H0, W0 = view1["true_shape"][0]
valid_matches_im0 = (
(matches_im0[:, 0] >= 3)
& (matches_im0[:, 0] < int(W0) - 3)
& (matches_im0[:, 1] >= 3)
& (matches_im0[:, 1] < int(H0) - 3)
)
H1, W1 = view2["true_shape"][0]
valid_matches_im1 = (
(matches_im1[:, 0] >= 3)
& (matches_im1[:, 0] < int(W1) - 3)
& (matches_im1[:, 1] >= 3)
& (matches_im1[:, 1] < int(H1) - 3)
)
valid_matches = valid_matches_im0 & valid_matches_im1
mkpts0, mkpts1 = matches_im0[valid_matches], matches_im1[valid_matches]
H0, W0, H1, W1 = *img0.shape[-2:], *img1.shape[-2:]
mkpts0 = self.rescale_coords(mkpts0, *img0_orig_shape, H0, W0)
mkpts1 = self.rescale_coords(mkpts1, *img1_orig_shape, H1, W1)
return mkpts0, mkpts1, None, None, None, None
|