diff --git a/configs/loftr.yml b/configs/loftr.yml index f5c2f8e..51672c9 100644 --- a/configs/loftr.yml +++ b/configs/loftr.yml @@ -3,12 +3,17 @@ default: &default ckpt: 'pretrained/loftr/outdoor_ds.ckpt' match_threshold: 0.2 imsize: -1 - no_matches_upscale: False + aspect_ratio: -1 + no_match_upscale: False + nbr_matches: -1 example: <<: *default match_threshold: 0.5 imsize: -1 hpatch: <<: *default - imsize: 480 + match_threshold: 0.5 + imsize: 640 + aspect_ratio: 1.333333 # images will be 640x480 or 480x640 no_match_upscale: True + nbr_matches: 1000 diff --git a/immatch/modules/loftr.py b/immatch/modules/loftr.py index e2f84f8..f6a3e47 100644 --- a/immatch/modules/loftr.py +++ b/immatch/modules/loftr.py @@ -14,8 +14,10 @@ def __init__(self, args): args = Namespace(**args) self.imsize = args.imsize + self.aspect_ratio = args.aspect_ratio self.match_threshold = args.match_threshold self.no_match_upscale = args.no_match_upscale + self.nbr_matches = args.nbr_matches # Load model conf = dict(default_cfg) @@ -34,7 +36,7 @@ def __init__(self, args): def load_im(self, im_path): return load_gray_scale_tensor_cv( - im_path, self.device, imsize=self.imsize, dfactor=8 + im_path, self.device, imsize=self.imsize, aspect_ratio=self.aspect_ratio, dfactor=8 ) def match_inputs_(self, gray1, gray2): @@ -43,6 +45,12 @@ def match_inputs_(self, gray1, gray2): kpts1 = batch['mkpts0_f'].cpu().numpy() kpts2 = batch['mkpts1_f'].cpu().numpy() scores = batch['mconf'].cpu().numpy() + if self.nbr_matches and self.nbr_matches > 0: + ids = np.argsort(scores)[-self.nbr_matches:] + kpts1 = kpts1[ids] + kpts2 = kpts2[ids] + scores = scores[ids] + matches = np.concatenate([kpts1, kpts2], axis=1) return matches, kpts1, kpts2, scores diff --git a/immatch/utils/data_io.py b/immatch/utils/data_io.py index 5c9da38..ffbb31b 100644 --- a/immatch/utils/data_io.py +++ b/immatch/utils/data_io.py @@ -10,10 +10,19 @@ def lprint(ms, log=None): log.write(ms+'\n') log.flush() -def resize_im(wo, ho, imsize=None, dfactor=1, value_to_scale=max): +def resize_im(wo, ho, imsize=None, aspect_ratio=None, dfactor=1): wt, ht = wo, ho - if imsize and value_to_scale(wo, ho) > imsize and imsize > 0: - scale = imsize / value_to_scale(wo, ho) + if aspect_ratio and aspect_ratio > 0: + if not imsize or imsize < 0: + raise ValueError('imsize must be given when using fixed aspect_ratio') + max_wh = imsize + min_wh = round(imsize / aspect_ratio) + if wo < ho: + wt, ht = min_wh, max_wh + else: + wt, ht = max_wh, min_wh + elif imsize and max(wo, ho) > imsize and imsize > 0: + scale = imsize / max(wo, ho) ht, wt = int(round(ho * scale)), int(round(wo * scale)) # Make sure new sizes are divisible by the given factor @@ -41,11 +50,11 @@ def load_gray_scale_tensor(im_path, device, imsize=None, dfactor=1): gray = transforms.functional.to_tensor(gray).unsqueeze(0).to(device) return gray, scale -def load_gray_scale_tensor_cv(im_path, device, imsize=None, dfactor=1): +def load_gray_scale_tensor_cv(im_path, device, imsize=None, aspect_ratio=None, dfactor=1): # Used for LoFTR im = cv2.imread(im_path, cv2.IMREAD_GRAYSCALE) ho, wo = im.shape - wt, ht, scale = resize_im(wo, ho, imsize=imsize, dfactor=dfactor, value_to_scale=min) + wt, ht, scale = resize_im(wo, ho, imsize=imsize, aspect_ratio=aspect_ratio, dfactor=dfactor) im = cv2.resize(im, (wt, ht)) im = transforms.functional.to_tensor(im).unsqueeze(0).to(device) return im, scale diff --git a/immatch/utils/hpatches_helper.py b/immatch/utils/hpatches_helper.py index 2ee51c7..03d0803 100644 --- a/immatch/utils/hpatches_helper.py +++ b/immatch/utils/hpatches_helper.py @@ -197,7 +197,7 @@ def eval_hpatches( if 'homography' in task: try: if 'cv' in h_solver: - H_pred, inliers = cv2.findHomography(matches[:, :2], matches[:, 2:4], cv2.RANSAC, ransac_thres) + H_pred, inliers = cv2.findHomography(matches[:, :2], matches[:, 2:4], cv2.RANSAC, ransac_thres, confidence=0.99999) else: H_pred, inliers = pydegensac.findHomography(matches[:, :2], matches[:, 2:4], ransac_thres) except: