update elfotr

Author: LuoXubo

utils.py

@@ -105,4 +105,40 @@ def sp_lg(sp, lg, ref, tgt):
         ref = np.transpose(ref,(1,2,0))
         ref = cv2.cvtColor(np.uint8(ref*255), cv2.COLOR_RGB2BGR)
 
-    return mkpts_0, mkpts_1, time_det, time_mat
+    return mkpts_0, mkpts_1, time_det, time_mat
+
+def eloftr(model, ref, tgt):
+    """
+    Function to perform LoFTR pipeline.
+    
+    Args:
+        model: LoFTR model
+        ref: reference image
+        tgt: target image
+    
+    Returns:
+        kpts0: keypoints of reference image
+        kpts1: keypoints of target image
+        time_det: time taken for detection
+        time_mat: time taken for matching
+    """
+    
+    ref = cv2.resize(ref, (ref.shape[1]//32*32, ref.shape[0]//32*32))
+    tgt = cv2.resize(tgt, (tgt.shape[1]//32*32, tgt.shape[0]//32*32))
+    
+    ref = torch.from_numpy(ref)[None][None].cuda()/255.
+    tgt = torch.from_numpy(tgt)[None][None].cuda()/255.
+    
+    batch = {'image0': ref, 'image1': tgt}
+    
+    tik = time.time()
+    with torch.no_grad():
+        model(batch)
+    tok = time.time()
+    time_total = tok - tik
+    
+    mkpts0 = batch['mkpts0_f'].cpu().numpy()
+    mkpts1 = batch['mkpts1_f'].cpu().numpy()
+    mconf = batch['mconf'].cpu().numpy()
+    
+    return mkpts0, mkpts1, 0, time_total

video_matching.py

@@ -10,17 +10,19 @@
 import torch
 import argparse
 
+from copy import deepcopy
 from xfeat.xfeat import XFeat
 from lightglue import LightGlue, SuperPoint
+from efficientloftr.loftr import LoFTR, full_default_cfg, opt_default_cfg, reparameter
 
-from utils import warp_corners_and_draw_matches, sp_lg
+from utils import warp_corners_and_draw_matches, sp_lg, eloftr
 
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--ref', type=str, help='Path to the reference image', default='assets/groot/groot.jpg')
     parser.add_argument('--tgt', type=str, help='Path to the target video', default='assets/groot/groot.mp4')
-    parser.add_argument('--method', type=str, help='Method to use for image matching (xfeat+mnn, sp+lg)', default='xfeat+mnn')
+    parser.add_argument('--method', type=str, help='Method to use for image matching (xfeat+mnn, sp+lg, loftr)', default='xfeat+mnn')
     parser.add_argument('--save_path', type=str, help='Path to save the output video', default='output.mp4')
 
     args = parser.parse_args()
@@ -36,6 +38,13 @@
     elif method == 'sp+lg':
         extractor = SuperPoint(max_num_keypoints=2048).eval().to(device)
         matcher = LightGlue(features="superpoint").eval().to(device)
+        print(f'Load superpoint and lightglue to {device}')
+    elif method == 'loftr':
+        _default_cfg = deepcopy(full_default_cfg)
+        loftr = LoFTR(config=_default_cfg)
+        loftr.load_state_dict(torch.load("weights/eloftr_outdoor.ckpt")['state_dict'])
+        loftr = reparameter(loftr)
+        loftr = loftr.eval().to(device)
     else:
         raise ValueError(f'Unknown method: {method}')
 
@@ -72,8 +81,8 @@
                 mkpts_0, mkpts_1, time_det, time_mat = xfeat.match_xfeat(ref, frame, top_k = 4096)
             elif method == 'sp+lg':
                 mkpts_0, mkpts_1, time_det, time_mat = sp_lg(extractor, matcher, ref, frame)
-                
-    
+            elif method == 'loftr':
+                mkpts_0, mkpts_1, time_det, time_mat = eloftr(loftr, ref, frame)
             time_total = time_det + time_mat
 
             canvas = warp_corners_and_draw_matches(mkpts_1, mkpts_0, frame, ref, time_total)