def test_jit(self, device, dtype):
B, C, H, W = 1, 1, 32, 32
patches = torch.rand(B, C, H, W, device=device, dtype=dtype)
patches2x = resize(patches, (48, 48))
input = {"image0": patches, "image1": patches2x}
model = LoFTR().to(patches.device, patches.dtype).eval()
model_jit = torch.jit.script(model)
out = model(input)
out_jit = model_jit(input)
for k, v in out.items():
assert_close(v, out_jit[k])
def test_gradcheck(self, device):
patches = torch.rand(1, 1, 32, 32, device=device)
patches05 = resize(patches, (48, 48))
patches = utils.tensor_to_gradcheck_var(patches) # to var
patches05 = utils.tensor_to_gradcheck_var(patches05) # to var
loftr = LoFTR().to(patches.device, patches.dtype)
def proxy_forward(x, y):
return loftr.forward({"image0": x, "image1": y})["keypoints0"]
assert gradcheck(proxy_forward, (patches, patches05), eps=1e-4, atol=1e-4, raise_exception=True)
def test_jit(self, device, dtype):
B, C, H, W = 1, 1, 32, 32
patches = torch.rand(B, C, H, W, device=device, dtype=dtype)
patches2x = resize(patches, (48, 48))
inputs = {"image0": patches, "image1": patches2x}
model = LocalFeatureMatcher(SIFTDescriptor(32),
DescriptorMatcher('snn',
0.8)).to(device).eval()
model_jit = torch.jit.script(model)
out = model(inputs)
out_jit = model_jit(inputs)
for k, v in out.items():
assert_close(v, out_jit[k])
def test_gradcheck(self, device):
matcher = LocalFeatureMatcher(SIFTFeature(5),
DescriptorMatcher('nn', 1.0)).to(device)
patches = torch.rand(1, 1, 32, 32, device=device)
patches05 = resize(patches, (48, 48))
patches = utils.tensor_to_gradcheck_var(patches) # to var
patches05 = utils.tensor_to_gradcheck_var(patches05) # to var
def proxy_forward(x, y):
return matcher({"image0": x, "image1": y})["keypoints0"]
assert gradcheck(proxy_forward, (patches, patches05),
eps=1e-4,
atol=1e-4,
raise_exception=True)
def forward(self, x):
return geometry.resize(x[None], self.size)[0]
def forward(self, data: dict) -> Dict[str, torch.Tensor]:
"""
Args:
data: dictionary containing the input data in the following format:
Keyword Args:
image0: left image with shape :math:`(N, 1, H1, W1)`.
image1: right image with shape :math:`(N, 1, H2, W2)`.
mask0 (optional): left image mask. '0' indicates a padded position :math:`(N, H1, W1)`.
mask1 (optional): right image mask. '0' indicates a padded position :math:`(N, H2, W2)`.
Returns:
- ``keypoints0``, matching keypoints from image0 :math:`(NC, 2)`.
- ``keypoints1``, matching keypoints from image1 :math:`(NC, 2)`.
- ``confidence``, confidence score [0, 1] :math:`(NC)`.
- ``batch_indexes``, batch indexes for the keypoints and lafs :math:`(NC)`.
"""
# 1. Local Feature CNN
data.update({
'bs': data['image0'].size(0),
'hw0_i': data['image0'].shape[2:], 'hw1_i': data['image1'].shape[2:]
})
if data['hw0_i'] == data['hw1_i']: # faster & better BN convergence
feats_c, feats_f = self.backbone(torch.cat([data['image0'], data['image1']], dim=0))
(feat_c0, feat_c1), (feat_f0, feat_f1) = feats_c.split(data['bs']), feats_f.split(data['bs'])
else: # handle different input shapes
(feat_c0, feat_f0), (feat_c1, feat_f1) = self.backbone(data['image0']), self.backbone(data['image1'])
data.update({
'hw0_c': feat_c0.shape[2:], 'hw1_c': feat_c1.shape[2:],
'hw0_f': feat_f0.shape[2:], 'hw1_f': feat_f1.shape[2:]
})
# 2. coarse-level loftr module
# add featmap with positional encoding, then flatten it to sequence [N, HW, C]
# feat_c0 = rearrange(self.pos_encoding(feat_c0), 'n c h w -> n (h w) c')
# feat_c1 = rearrange(self.pos_encoding(feat_c1), 'n c h w -> n (h w) c')
feat_c0 = self.pos_encoding(feat_c0).permute(0, 2, 3, 1)
n, h, w, c = feat_c0.shape
feat_c0 = feat_c0.reshape(n, -1, c)
feat_c1 = self.pos_encoding(feat_c1).permute(0, 2, 3, 1)
n1, h1, w1, c1 = feat_c1.shape
feat_c1 = feat_c1.reshape(n1, -1, c1)
mask_c0 = mask_c1 = None # mask is useful in training
if 'mask0' in data:
mask_c0 = resize(data['mask0'], data['hw0_c'], interpolation='nearest').flatten(-2)
if 'mask1' in data:
mask_c1 = resize(data['mask1'], data['hw1_c'], interpolation='nearest').flatten(-2)
feat_c0, feat_c1 = self.loftr_coarse(feat_c0, feat_c1, mask_c0, mask_c1)
# 3. match coarse-level
self.coarse_matching(feat_c0, feat_c1, data, mask_c0=mask_c0, mask_c1=mask_c1)
# 4. fine-level refinement
feat_f0_unfold, feat_f1_unfold = self.fine_preprocess(feat_f0, feat_f1, feat_c0, feat_c1, data)
if feat_f0_unfold.size(0) != 0: # at least one coarse level predicted
feat_f0_unfold, feat_f1_unfold = self.loftr_fine(feat_f0_unfold, feat_f1_unfold)
# 5. match fine-level
self.fine_matching(feat_f0_unfold, feat_f1_unfold, data)
rename_keys: Dict[str, str] = {"mkpts0_f": 'keypoints0',
"mkpts1_f": 'keypoints1',
"mconf": 'confidence',
"b_ids": 'batch_indexes'}
out = {}
for k, v in rename_keys.items():
out[v] = data[k]
return out
