vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/LiftFeat/tools/demo_match_video.py

@@ -0,0 +1,145 @@
+import os
+import cv2
+import torch
+import numpy as np
+import yaml
+import matplotlib.cm as cm
+import argparse
+
+import sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from models.liftfeat_wrapper import LiftFeat,MODEL_PATH
+from utils.post_process import match_features
+os.environ['CUDA_VISIBLE_DEVICES']='0'
+
+use_cuda = torch.cuda.is_available()
+device = torch.device("cuda" if use_cuda else "cpu")
+
+class VideoHandler:
+    def __init__(self,video_path,size=[640,360]):
+        self.video_path=video_path
+        self.size=size
+        self.cap=cv2.VideoCapture(video_path)
+        
+    def get_frame(self):
+        ret,frame=self.cap.read()
+        if ret==True:
+            frame=cv2.resize(frame,(int(self.size[0]),int(self.size[1])))
+        return ret,frame
+    
+def draw_video_match(img0,img1,kpts0,kpts1,mkpts0,mkpts1,match_scores,mask,max_match_num=512,margin=15):
+    H0, W0, c = img0.shape
+    H1, W1, c = img1.shape
+    H, W = max(H0, H1), W0 + W1 + margin
+    
+    # 构建画布，把两个图像先拼接到一起
+    out = 255*np.ones((H, W, 3), np.uint8)
+    out[:H0, :W0, :] = img0
+    out[:H1, W0+margin:, :] = img1
+    #out = np.stack([out]*3, -1)
+
+    kpts0, kpts1 = np.round(kpts0).astype(int), np.round(kpts1).astype(int)
+
+    mkpts0, mkpts1 = np.round(mkpts0).astype(int), np.round(mkpts1).astype(int)
+    mkpts0_correct,mkpts1_correct=mkpts0[mask],mkpts1[mask]
+    mkpts0_wrong,mkpts1_wrong=mkpts0[~mask],mkpts1[~mask]
+    match_s=match_scores[mask]
+    
+    print(f"correct: {mkpts0_correct.shape[0]} wrong: {mkpts0_wrong.shape[0]}")
+    
+    if mkpts0_correct.shape[0] > max_match_num:
+        # perm=np.random.randint(low=0,high=mkpts0_correct.shape[0],size=max_match_num)
+        # mkpts0_show,mkpts1_show=mkpts0_correct[perm],mkpts1_correct[perm]
+        mkpts0_show,mkpts1_show=mkpts0_correct,mkpts1_correct
+    else:
+        mkpts0_show,mkpts1_show=mkpts0_correct,mkpts1_correct
+        
+    # 普通的点
+    vis_normal_point = True
+    if (vis_normal_point):
+        for x, y in mkpts0_show:
+            cv2.circle(out, (x, y), 2, (47,132,250), -1, lineType=cv2.LINE_AA)
+        for x, y in mkpts1_show:
+            cv2.circle(out, (x + margin + W0, y), 2, (47,132,250), -1,lineType=cv2.LINE_AA)
+    
+    vis_match_line = True
+    if (vis_match_line):
+        for pt0, pt1,score in zip(mkpts0_show, mkpts1_show,match_s):
+            color_cm = cm.jet(1.0 - score, alpha=0)  
+            color = (int(color_cm[0] * 255), int(color_cm[1] * 255), int(color_cm[2] * 255))
+            cv2.line(out, pt0, (W0 + margin + pt1[0], pt1[1]), color, 1)
+            
+    return out
+    
+def run_video_demo(std_img_path,video_path):
+   
+    
+    liftfeat=LiftFeat(weight=MODEL_PATH,detect_threshold=0.15)
+    
+    std_img=cv2.imread(std_img_path)
+    std_img=cv2.resize(std_img,(640,360))
+    
+    handler=VideoHandler(video_path)
+    
+    # 定义编解码器并创建VideoWriter对象
+    if not os.path.exists('./output'):
+        os.makedirs('./output')
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 或者使用 'XVID'
+    out = cv2.VideoWriter('./output/video_demo.mp4', fourcc, 20.0, (1300, 360))
+    K=[[1084.8,0,640.24],[0,1085,354.87],[0,0,1]]
+    K=np.array(K)
+    data_std=liftfeat.extract(std_img)
+    
+    while True:
+        ret,frame=handler.get_frame()
+        if ret==False:
+            break
+       
+        if frame is not None:
+            data=liftfeat.extract(frame)
+            idx0, idx1, match_scores=match_features(data_std["descriptors"],data["descriptors"],-1)
+            mkpts0=data_std["keypoints"][idx0]
+            mkpts1=data["keypoints"][idx1]
+            mkpts0_np=mkpts0.cpu().numpy()
+            mkpts1_np=mkpts1.cpu().numpy()
+            match_scores_np=match_scores.detach().cpu().numpy()
+            kpts0 = (mkpts0_np - K[[0, 1], [2, 2]][None]) / K[[0, 1], [0, 1]][None]
+            kpts1 = (mkpts1_np - K[[0, 1], [2, 2]][None]) / K[[0, 1], [0, 1]][None]
+            
+            # normalize ransac threshold
+            ransac_thr = 0.5 / np.mean([K[0, 0], K[1, 1], K[0, 0], K[1, 1]])
+            
+            if mkpts0_np.shape[0] < 5:
+                print(f"mkpts size less then 5")
+            else:
+                # compute pose with cv2
+                
+                E, mask = cv2.findEssentialMat(kpts0, kpts1, np.eye(3), threshold=ransac_thr, prob=0.999, method=cv2.RANSAC)
+                if E is None:
+                    print("\nE is None while trying to recover pose.\n")
+                    continue
+                match_mask=mask.squeeze(axis=1)>0
+                show_kpts0,show_kpts1=mkpts0_np[match_mask],mkpts1_np[match_mask]
+                show_match_scores=match_scores_np[match_mask]
+                show_mask=np.ones(show_kpts0.shape[0])>0
+                match_img=draw_video_match(std_img,frame,show_kpts0,show_kpts1,show_kpts0,show_kpts1,show_match_scores,show_mask,margin=20)
+                kpts0_num,kpts1_num=data_std["keypoints"].shape[0],data["keypoints"].shape[0]
+                cv2.putText(match_img,f"LiftFeat",(10,20),cv2.FONT_HERSHEY_TRIPLEX,0.5,(0,0,241))
+                cv2.putText(match_img,f"Keypoints: {kpts0_num}:{kpts1_num}",(10,40),cv2.FONT_HERSHEY_TRIPLEX,0.5,(0,0,255))
+                cv2.putText(match_img,f"Matches: {show_kpts0.shape[0]}",(10,60),cv2.FONT_HERSHEY_TRIPLEX,0.5,(0,0,255))
+                out.write(match_img)
+                
+                
+    out.release()
+    
+
+
+            
+if __name__=="__main__":
+    parser = argparse.ArgumentParser(description="Run LiftFeat video matching demo.")
+    parser.add_argument('--img', type=str, required=True, help='Path to the template image')
+    parser.add_argument('--video', type=str, required=True, help='Path to the input video')
+    
+    args = parser.parse_args()
+    
+    run_video_demo(args.img, args.video)

vismatch/third_party/LiftFeat/tools/demo_vo.py

@@ -0,0 +1,163 @@
+import numpy as np
+import cv2
+import argparse
+import yaml
+import logging
+import os
+import sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from utils.VisualOdometry import VisualOdometry, AbosluteScaleComputer, create_dataloader, \
+    plot_keypoints, create_detector, create_matcher
+from models.liftfeat_wrapper import LiftFeat,MODEL_PATH
+
+
+vo_config = {
+    'dataset': {
+        'name': 'KITTILoader',
+        'root_path': '/home/yepeng_liu/code_python/dataset/visual_odometry/kitty/gray',
+        'sequence': '10',
+        'start': 0
+    },
+    'detector': {
+        'name': 'LiftFeatDetector',
+        'descriptor_dim': 64,
+        'nms_radius': 5,
+        'keypoint_threshold': 0.005,
+        'max_keypoints': 4096,
+        'remove_borders': 4,
+        'cuda': 1
+    },
+    'matcher': {
+        'name': 'FrameByFrameMatcher',
+        'type': 'FLANN',
+        'FLANN': {
+            'kdTrees': 5,
+            'searchChecks': 50
+        },
+        'distance_ratio': 0.75
+    }
+}
+
+# 可视化当前frame的关键点
+def keypoints_plot(img, vo, img_id, path2):
+    img_ = cv2.imread(path2+str(img_id-1).zfill(6)+".png")
+  
+    if not vo.match_kps:
+        img_ = plot_keypoints(img_, vo.kptdescs["cur"]["keypoints"])
+    else:
+        for index in range(vo.match_kps["ref"].shape[0]):
+            ref_point = tuple(map(int, vo.match_kps['ref'][index,:]))  # 将关键点转换为整数元组
+            cur_point = tuple(map(int, vo.match_kps['cur'][index,:]))
+            cv2.line(img_, ref_point, cur_point, (0, 255, 0), 2)  # Draw green line
+            cv2.circle(img_, cur_point, 3, (0, 0, 255), -1)  # Draw red circle at current keypoint
+
+    return img_
+
+# 负责绘制相机的轨迹并计算估计轨迹与真实轨迹的误差。
+class TrajPlotter(object):
+    def __init__(self):
+        self.errors = []
+        self.traj = np.zeros((800, 800, 3), dtype=np.uint8)
+        pass
+
+    def update(self, est_xyz, gt_xyz):
+        x, z = est_xyz[0], est_xyz[2]
+        gt_x, gt_z = gt_xyz[0], gt_xyz[2]
+        est = np.array([x, z]).reshape(2)
+        gt = np.array([gt_x, gt_z]).reshape(2)
+        error = np.linalg.norm(est - gt)
+        self.errors.append(error)
+        avg_error = np.mean(np.array(self.errors))
+        # === drawer ==================================
+        # each point
+        draw_x, draw_y = int(x) + 80, int(z) + 230
+        true_x, true_y = int(gt_x) + 80, int(gt_z) + 230
+
+        # draw trajectory
+        cv2.circle(self.traj, (draw_x, draw_y), 1, (0, 0, 255), 1)
+        cv2.circle(self.traj, (true_x, true_y), 1, (0, 255, 0), 2)
+        cv2.rectangle(self.traj, (10, 5), (450, 120), (0, 0, 0), -1)
+
+        # draw text
+        text = "[AvgError] %2.4fm" % (avg_error)
+        print(text)
+        cv2.putText(self.traj, text, (20, 40),
+                    cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
+        note = "Green: GT, Red: Predict"
+        cv2.putText(self.traj, note, (20, 80),
+                    cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
+
+        return self.traj
+
+def run_video(args):
+    # create dataloader
+    vo_config["dataset"]['root_path'] = args.path1
+    vo_config["dataset"]['sequence'] = args.id
+    loader = create_dataloader(vo_config["dataset"])
+    # create detector
+    liftfeat=LiftFeat(weight=MODEL_PATH, detect_threshold=0.25)
+    # create matcher
+    matcher = create_matcher(vo_config["matcher"])
+
+    absscale = AbosluteScaleComputer()
+    traj_plotter = TrajPlotter()
+
+  
+    if not os.path.exists('./output'):
+        os.makedirs('./output')
+    fname = "kitti_liftfeat_flannmatch"
+    log_fopen = open("output/" + fname + ".txt", mode='a')
+
+    vo = VisualOdometry(liftfeat, matcher, loader.cam)
+
+    # Initialize video writer for keypoints and trajectory videos
+    keypoints_video_path = "output/" + fname + "_keypoints_liftfeat.avi"
+    trajectory_video_path = "output/" + fname + "_trajectory_liftfeat.avi"
+
+    # Set up video writer: choose codec and set FPS and frame size
+    fourcc = cv2.VideoWriter_fourcc(*'XVID')
+    fps = 10  # Adjust the FPS according to your input data
+    frame_size = (1200, 400)  # Get frame size from first image
+
+    # Video writers for keypoints and trajectory
+    keypoints_writer = cv2.VideoWriter(keypoints_video_path, fourcc, fps, frame_size)
+    trajectory_writer = cv2.VideoWriter(trajectory_video_path, fourcc, fps, (800, 800))
+    
+    for i, img in enumerate(loader):
+        img_id = loader.img_id
+        gt_pose = loader.get_cur_pose()
+       
+        R, t = vo.update(img, absscale.update(gt_pose))
+        
+        # === log writer ==============================
+        print(i, t[0, 0], t[1, 0], t[2, 0], gt_pose[0, 3], gt_pose[1, 3], gt_pose[2, 3], file=log_fopen)
+
+        # === drawer ==================================
+        img1 = keypoints_plot(img, vo, img_id, args.path2)
+        img1 = cv2.resize(img1, (1200, 400))
+        img2 = traj_plotter.update(t, gt_pose[:, 3])
+
+        # Write frames to videos
+        keypoints_writer.write(img1)
+        trajectory_writer.write(img2)
+
+    # Release the video writers
+    keypoints_writer.release()
+    trajectory_writer.release()
+    print(f"Videos saved as {keypoints_video_path} and {trajectory_video_path}")
+   
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='python_vo')
+    parser.add_argument('--path1', type=str, default='/home/yepeng_liu/code_python/dataset/visual_odometry/kitty/gray',
+                        help='config file')
+    parser.add_argument('--path2', type=str, default="/home/yepeng_liu/code_python/dataset/visual_odometry/kitty/color/sequences/03/image_2/",
+                        help='config file')
+    parser.add_argument('--id', type=str, default="03",
+                        help='config file')
+   
+
+    args = parser.parse_args()
+    
+    run_video(args)

vismatch/third_party/LiftFeat/train.py

@@ -0,0 +1,369 @@
+"""
+	"LiftFeat: 3D Geometry-Aware Local Feature Matching"
+    training script
+"""
+
+import argparse
+import os
+import time
+import sys
+sys.path.append(os.path.dirname(__file__))
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="LiftFeat training script.")
+    parser.add_argument('--name',type=str,default='LiftFeat',help='set process name')
+    
+    # MegaDepth dataset setting
+    parser.add_argument('--use_megadepth',action='store_true')
+    parser.add_argument('--megadepth_root_path', type=str,
+                        default='/home/yepeng_liu/code_python/dataset/MegaDepth/phoenix/S6/zl548',
+                        help='Path to the MegaDepth dataset root directory.')
+    parser.add_argument('--megadepth_batch_size', type=int, default=6)
+    
+    # COCO20k dataset setting
+    parser.add_argument('--use_coco',action='store_true')
+    parser.add_argument('--coco_root_path', type=str, default='/home/yepeng_liu/code_python/dataset/coco_20k',
+                        help='Path to the COCO20k dataset root directory.')
+    parser.add_argument('--coco_batch_size',type=int,default=4)
+    
+    parser.add_argument('--ckpt_save_path', type=str, default='/home/yepeng_liu/code_python/LiftFeat/trained_weights/test',
+                        help='Path to save the checkpoints.')
+    parser.add_argument('--n_steps', type=int, default=160_000,
+                        help='Number of training steps. Default is 160000.')
+    parser.add_argument('--lr', type=float, default=3e-4,
+                        help='Learning rate. Default is 0.0003.')
+    parser.add_argument('--gamma_steplr', type=float, default=0.5,
+                        help='Gamma value for StepLR scheduler. Default is 0.5.')
+    parser.add_argument('--training_res', type=lambda s: tuple(map(int, s.split(','))),
+                        default=(800, 608), help='Training resolution as width,height. Default is (800, 608).')
+    parser.add_argument('--device_num', type=str, default='0',
+                        help='Device number to use for training. Default is "0".')
+    parser.add_argument('--dry_run', action='store_true',
+                        help='If set, perform a dry run training with a mini-batch for sanity check.')
+    parser.add_argument('--save_ckpt_every', type=int, default=500,
+                        help='Save checkpoints every N steps. Default is 500.')
+    parser.add_argument('--use_coord_loss',action='store_true',help='Enable coordinate loss')
+
+    args = parser.parse_args()
+
+    os.environ['CUDA_VISIBLE_DEVICES'] = args.device_num
+
+    return args
+
+args = parse_arguments()
+
+import torch
+from torch import nn
+from torch import optim
+import torch.nn.functional as F
+from torch.utils.tensorboard import SummaryWriter
+from torch.utils.data import Dataset, DataLoader
+
+import numpy as np
+import tqdm
+import glob
+
+from models.model import LiftFeatSPModel
+from loss.loss import LiftFeatLoss
+from utils.config import featureboost_config
+from models.interpolator import InterpolateSparse2d
+from utils.depth_anything_wrapper import DepthAnythingExtractor
+from utils.alike_wrapper import ALikeExtractor
+
+from dataset import megadepth_wrapper
+from dataset import coco_wrapper
+from dataset.megadepth import MegaDepthDataset
+from dataset.coco_augmentor import COCOAugmentor
+
+import setproctitle
+
+
+class Trainer():
+    def __init__(self, megadepth_root_path,use_megadepth,megadepth_batch_size,
+                       coco_root_path,use_coco,coco_batch_size,
+                       ckpt_save_path, 
+                       model_name = 'LiftFeat',
+                       n_steps = 160_000, lr= 3e-4, gamma_steplr=0.5, 
+                       training_res = (800, 608), device_num="0", dry_run = False,
+                       save_ckpt_every = 500, use_coord_loss = False):
+        print(f'MegeDepth: {use_megadepth}-{megadepth_batch_size}')
+        print(f'COCO20k: {use_coco}-{coco_batch_size}')
+        print(f'Coordinate loss: {use_coord_loss}')
+        self.dev = torch.device ('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        # training model
+        self.net = LiftFeatSPModel(featureboost_config, use_kenc=False, use_normal=True, use_cross=True).to(self.dev)
+        self.loss_fn=LiftFeatLoss(self.dev,lam_descs=1,lam_kpts=2,lam_heatmap=1)
+        
+        # depth-anything model
+        self.depth_net=DepthAnythingExtractor('vits',self.dev,256)
+        
+        # alike model
+        self.alike_net=ALikeExtractor('alike-t',self.dev)
+
+        #Setup optimizer 
+        self.steps = n_steps
+        self.opt = optim.Adam(filter(lambda x: x.requires_grad, self.net.parameters()) , lr = lr)
+        self.scheduler = torch.optim.lr_scheduler.StepLR(self.opt, step_size=10_000, gamma=gamma_steplr)
+
+        ##################### COCO INIT ##########################
+        self.use_coco=use_coco
+        self.coco_batch_size=coco_batch_size
+        if self.use_coco:
+            self.augmentor=COCOAugmentor(
+                img_dir=coco_root_path,
+                device=self.dev,load_dataset=True,
+                batch_size=self.coco_batch_size,
+                out_resolution=training_res,
+                warp_resolution=training_res,
+                sides_crop=0.1,
+                max_num_imgs=3000,
+                num_test_imgs=5,
+                photometric=True,
+                geometric=True,
+                reload_step=4000
+            )
+        ##################### COCO END #######################
+
+
+        ##################### MEGADEPTH INIT ##########################
+        self.use_megadepth=use_megadepth
+        self.megadepth_batch_size=megadepth_batch_size
+        if self.use_megadepth:
+            TRAIN_BASE_PATH = f"{megadepth_root_path}/train_data/megadepth_indices"
+            TRAINVAL_DATA_SOURCE = f"{megadepth_root_path}/MegaDepth_v1"
+
+            TRAIN_NPZ_ROOT = f"{TRAIN_BASE_PATH}/scene_info_0.1_0.7"
+
+            npz_paths = glob.glob(TRAIN_NPZ_ROOT + '/*.npz')[:]
+            megadepth_dataset = torch.utils.data.ConcatDataset( [MegaDepthDataset(root_dir = TRAINVAL_DATA_SOURCE,
+                            npz_path = path) for path in tqdm.tqdm(npz_paths, desc="[MegaDepth] Loading metadata")] )
+
+            self.megadepth_dataloader = DataLoader(megadepth_dataset, batch_size=megadepth_batch_size, shuffle=True)
+            self.megadepth_data_iter = iter(self.megadepth_dataloader)
+        ##################### MEGADEPTH INIT END #######################
+
+        os.makedirs(ckpt_save_path, exist_ok=True)
+        os.makedirs(ckpt_save_path + '/logdir', exist_ok=True)
+
+        self.dry_run = dry_run
+        self.save_ckpt_every = save_ckpt_every
+        self.ckpt_save_path = ckpt_save_path
+        self.writer = SummaryWriter(ckpt_save_path + f'/logdir/{model_name}_' + time.strftime("%Y_%m_%d-%H_%M_%S"))
+        self.model_name = model_name
+        self.use_coord_loss = use_coord_loss
+        
+        
+    def generate_train_data(self):
+        imgs1_t,imgs2_t=[],[]
+        imgs1_np,imgs2_np=[],[]
+        # norms0,norms1=[],[]
+        positives_coarse=[]
+        
+        if self.use_coco:
+                coco_imgs1, coco_imgs2, H1, H2 = coco_wrapper.make_batch(self.augmentor, 0.1)
+                h_coarse, w_coarse = coco_imgs1[0].shape[-2] // 8, coco_imgs1[0].shape[-1] // 8
+                _ , positives_coco_coarse = coco_wrapper.get_corresponding_pts(coco_imgs1, coco_imgs2, H1, H2, self.augmentor, h_coarse, w_coarse)
+                coco_imgs1=coco_imgs1.mean(1,keepdim=True);coco_imgs2=coco_imgs2.mean(1,keepdim=True)
+                imgs1_t.append(coco_imgs1);imgs2_t.append(coco_imgs2)
+                positives_coarse += positives_coco_coarse
+                    
+        if self.use_megadepth:
+            try:
+                megadepth_data=next(self.megadepth_data_iter)
+            except StopIteration:
+                print('End of MD DATASET')
+                self.megadepth_data_iter=iter(self.megadepth_dataloader)
+                megadepth_data=next(self.megadepth_data_iter)
+            if megadepth_data is not None:
+                for k in megadepth_data.keys():
+                    if isinstance(megadepth_data[k],torch.Tensor):
+                        megadepth_data[k]=megadepth_data[k].to(self.dev)
+                megadepth_imgs1_t,megadepth_imgs2_t=megadepth_data['image0'],megadepth_data['image1']
+                megadepth_imgs1_t=megadepth_imgs1_t.mean(1,keepdim=True);megadepth_imgs2_t=megadepth_imgs2_t.mean(1,keepdim=True)
+                imgs1_t.append(megadepth_imgs1_t);imgs2_t.append(megadepth_imgs2_t)
+                megadepth_imgs1_np,megadepth_imgs2_np=megadepth_data['image0_np'],megadepth_data['image1_np']
+                for np_idx in range(megadepth_imgs1_np.shape[0]):
+                    img1_np,img2_np=megadepth_imgs1_np[np_idx].squeeze(0).cpu().numpy(),megadepth_imgs2_np[np_idx].squeeze(0).cpu().numpy()
+                    imgs1_np.append(img1_np);imgs2_np.append(img2_np)
+                positives_megadepth_coarse=megadepth_wrapper.spvs_coarse(megadepth_data,8)
+                positives_coarse += positives_megadepth_coarse
+                
+        with torch.no_grad():
+            imgs1_t=torch.cat(imgs1_t,dim=0)
+            imgs2_t=torch.cat(imgs2_t,dim=0)
+            
+        return imgs1_t,imgs2_t,imgs1_np,imgs2_np,positives_coarse
+
+
+    def train(self):
+        self.net.train()
+
+        with tqdm.tqdm(total=self.steps) as pbar:
+            for i in range(self.steps):
+                # import pdb;pdb.set_trace()
+                imgs1_t,imgs2_t,imgs1_np,imgs2_np,positives_coarse=self.generate_train_data()
+
+                #Check if batch is corrupted with too few correspondences
+                is_corrupted = False
+                for p in positives_coarse:
+                    if len(p) < 30:
+                        is_corrupted = True
+
+                if is_corrupted:
+                    continue
+
+                # import pdb;pdb.set_trace()
+                #Forward pass
+                # start=time.perf_counter()
+                feats1,kpts1,normals1 = self.net.forward1(imgs1_t)
+                feats2,kpts2,normals2 = self.net.forward1(imgs2_t)
+                
+                coordinates,fb_coordinates=[],[]
+                alike_kpts1,alike_kpts2=[],[]
+                DA_normals1,DA_normals2=[],[]
+                
+                # import pdb;pdb.set_trace()
+                
+                fb_feats1,fb_feats2=[],[]
+                for b in range(feats1.shape[0]):
+                    feat1=feats1[b].permute(1,2,0).reshape(-1,feats1.shape[1])
+                    feat2=feats2[b].permute(1,2,0).reshape(-1,feats2.shape[1])
+                    
+                    coordinate=self.net.fine_matcher(torch.cat([feat1,feat2],dim=-1))
+                    coordinates.append(coordinate)
+                    
+                    fb_feat1=self.net.forward2(feats1[b].unsqueeze(0),kpts1[b].unsqueeze(0),normals1[b].unsqueeze(0))
+                    fb_feat2=self.net.forward2(feats2[b].unsqueeze(0),kpts2[b].unsqueeze(0),normals2[b].unsqueeze(0))
+                    
+                    fb_coordinate=self.net.fine_matcher(torch.cat([fb_feat1,fb_feat2],dim=-1))
+                    fb_coordinates.append(fb_coordinate)
+                    
+                    fb_feats1.append(fb_feat1.unsqueeze(0));fb_feats2.append(fb_feat2.unsqueeze(0))
+                    
+                    img1,img2=imgs1_t[b],imgs2_t[b]
+                    img1=img1.permute(1,2,0).expand(-1,-1,3).cpu().numpy() * 255
+                    img2=img2.permute(1,2,0).expand(-1,-1,3).cpu().numpy() * 255
+                    alike_kpt1=torch.tensor(self.alike_net.extract_alike_kpts(img1),device=self.dev)
+                    alike_kpt2=torch.tensor(self.alike_net.extract_alike_kpts(img2),device=self.dev)
+                    alike_kpts1.append(alike_kpt1);alike_kpts2.append(alike_kpt2)
+                
+                # import pdb;pdb.set_trace()
+                for b in range(len(imgs1_np)):
+                    megadepth_depth1,megadepth_norm1=self.depth_net.extract(imgs1_np[b])
+                    megadepth_depth2,megadepth_norm2=self.depth_net.extract(imgs2_np[b])
+                    DA_normals1.append(megadepth_norm1);DA_normals2.append(megadepth_norm2)
+                    
+                # import pdb;pdb.set_trace()
+                fb_feats1=torch.cat(fb_feats1,dim=0)
+                fb_feats2=torch.cat(fb_feats2,dim=0)
+                fb_feats1=fb_feats1.reshape(feats1.shape[0],feats1.shape[2],feats1.shape[3],-1).permute(0,3,1,2)
+                fb_feats2=fb_feats2.reshape(feats2.shape[0],feats2.shape[2],feats2.shape[3],-1).permute(0,3,1,2)
+                
+                coordinates=torch.cat(coordinates,dim=0)
+                coordinates=coordinates.reshape(feats1.shape[0],feats1.shape[2],feats1.shape[3],-1).permute(0,3,1,2)
+                
+                fb_coordinates=torch.cat(fb_coordinates,dim=0)
+                fb_coordinates=fb_coordinates.reshape(feats1.shape[0],feats1.shape[2],feats1.shape[3],-1).permute(0,3,1,2)
+                
+                # end=time.perf_counter()
+                # print(f"forward1 cost {end-start} seconds")
+
+                loss_items = []
+
+                # import pdb;pdb.set_trace()
+                loss_info=self.loss_fn(
+                    feats1,fb_feats1,kpts1,normals1,
+                    feats2,fb_feats2,kpts2,normals2,
+                    positives_coarse,
+                    coordinates,fb_coordinates,
+                    alike_kpts1,alike_kpts2,
+                    DA_normals1,DA_normals2,
+                    self.megadepth_batch_size,self.coco_batch_size)
+                
+                loss_descs,acc_coarse=loss_info['loss_descs'],loss_info['acc_coarse']
+                loss_coordinates,acc_coordinates=loss_info['loss_coordinates'],loss_info['acc_coordinates']
+                loss_fb_descs,acc_fb_coarse=loss_info['loss_fb_descs'],loss_info['acc_fb_coarse']
+                loss_fb_coordinates,acc_fb_coordinates=loss_info['loss_fb_coordinates'],loss_info['acc_fb_coordinates']
+                loss_kpts,acc_kpt=loss_info['loss_kpts'],loss_info['acc_kpt']
+                loss_normals=loss_info['loss_normals']
+                
+                loss_items.append(loss_fb_descs.unsqueeze(0))
+                loss_items.append(loss_kpts.unsqueeze(0))
+                loss_items.append(loss_normals.unsqueeze(0))
+                
+                if self.use_coord_loss:
+                    loss_items.append(loss_fb_coordinates.unsqueeze(0))
+
+                # nb_coarse = len(m1)
+                # nb_coarse = len(fb_m1)
+                loss = torch.cat(loss_items, -1).mean()
+
+                # Compute Backward Pass
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(self.net.parameters(), 1.)
+                self.opt.step()
+                self.opt.zero_grad()
+                self.scheduler.step()
+
+                # import pdb;pdb.set_trace()
+                if (i+1) % self.save_ckpt_every == 0:
+                    print('saving iter ', i+1)
+                    torch.save(self.net.state_dict(), self.ckpt_save_path + f'/{self.model_name}_{i+1}.pth')
+
+                pbar.set_description(
+'Loss: {:.4f} \
+loss_descs: {:.3f} acc_coarse: {:.3f} \
+loss_coordinates: {:.3f} acc_coordinates: {:.3f} \
+loss_fb_descs: {:.3f} acc_fb_coarse: {:.3f} \
+loss_fb_coordinates: {:.3f} acc_fb_coordinates: {:.3f} \
+loss_kpts: {:.3f} acc_kpts: {:.3f} \
+loss_normals: {:.3f}'.format( \
+loss.item(), \
+loss_descs.item(), acc_coarse, \
+loss_coordinates.item(), acc_coordinates, \
+loss_fb_descs.item(), acc_fb_coarse, \
+loss_fb_coordinates.item(), acc_fb_coordinates, \
+loss_kpts.item(), acc_kpt, \
+loss_normals.item()) )
+                pbar.update(1)
+
+                # Log metrics
+                self.writer.add_scalar('Loss/total', loss.item(), i)
+                self.writer.add_scalar('Accuracy/acc_coarse', acc_coarse, i)
+                self.writer.add_scalar('Accuracy/acc_coordinates', acc_coordinates, i)
+                self.writer.add_scalar('Accuracy/acc_fb_coarse', acc_fb_coarse, i)
+                self.writer.add_scalar('Accuracy/acc_fb_coordinates', acc_fb_coordinates, i)
+                self.writer.add_scalar('Loss/descs', loss_descs.item(), i)
+                self.writer.add_scalar('Loss/coordinates', loss_coordinates.item(), i)
+                self.writer.add_scalar('Loss/fb_descs', loss_fb_descs.item(), i)
+                self.writer.add_scalar('Loss/fb_coordinates', loss_fb_coordinates.item(), i)
+                self.writer.add_scalar('Loss/kpts', loss_kpts.item(), i)
+                self.writer.add_scalar('Loss/normals', loss_normals.item(), i)
+
+
+
+if __name__ == '__main__':
+    
+    setproctitle.setproctitle(args.name)
+
+    trainer = Trainer(
+        megadepth_root_path=args.megadepth_root_path, 
+        use_megadepth=args.use_megadepth,
+        megadepth_batch_size=args.megadepth_batch_size,
+        coco_root_path=args.coco_root_path, 
+        use_coco=args.use_coco,
+        coco_batch_size=args.coco_batch_size,
+        ckpt_save_path=args.ckpt_save_path,
+        n_steps=args.n_steps,
+        lr=args.lr,
+        gamma_steplr=args.gamma_steplr,
+        training_res=args.training_res,
+        device_num=args.device_num,
+        dry_run=args.dry_run,
+        save_ckpt_every=args.save_ckpt_every,
+        use_coord_loss=args.use_coord_loss
+    )
+
+    #The most fun part
+    trainer.train()