vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/LiftFeat/evaluation/MegaDepth1500_evaluation.py

@@ -0,0 +1,105 @@
+import os
+import sys
+import cv2
+from pathlib import Path
+import numpy as np
+import torch
+import torch.utils.data as data
+import tqdm
+from copy import deepcopy
+from torchvision.transforms import ToTensor
+import torch.nn.functional as F
+import json
+
+import scipy.io as scio
+import poselib
+
+import argparse
+import datetime
+
+parser=argparse.ArgumentParser(description='MegaDepth dataset evaluation script')
+parser.add_argument('--name',type=str,default='LiftFeat',help='experiment name')
+parser.add_argument('--gpu',type=str,default='0',help='GPU ID')
+args=parser.parse_args()
+
+os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
+
+sys.path.append(os.path.join(os.path.dirname(__file__),'../'))
+from models.liftfeat_wrapper import LiftFeat
+from evaluation.eval_utils import *
+
+from torch.utils.data import Dataset,DataLoader
+
+use_cuda = torch.cuda.is_available()
+device = "cuda" if use_cuda else "cpu"
+
+DATASET_ROOT = os.path.join(os.path.dirname(__file__),'../data/megadepth_test_1500')
+DATASET_JSON = os.path.join(os.path.dirname(__file__),'../data/megadepth_1500.json')
+
+class MegaDepth1500(Dataset):
+    """
+        Streamlined MegaDepth-1500 dataloader. The camera poses & metadata are stored in a formatted json for facilitating 
+        the download of the dataset and to keep the setup as simple as possible.
+    """
+    def __init__(self, json_file, root_dir):
+        # Load the info & calibration from the JSON
+        with open(json_file, 'r') as f:
+            self.data = json.load(f)
+
+        self.root_dir = root_dir
+
+        if not os.path.exists(self.root_dir):
+            raise RuntimeError(
+            f"Dataset {self.root_dir} does not exist! \n \
+              > If you didn't download the dataset, use the downloader tool: python3 -m modules.dataset.download -h")
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        data = deepcopy(self.data[idx])
+
+        h1, w1 = data['size0_hw']
+        h2, w2 = data['size1_hw']
+
+        # Here we resize the images to max_dim = 1200, as described in the paper, and adjust the image such that it is divisible by 32
+        # following the protocol of the LoFTR's Dataloader (intrinsics are corrected accordingly). 
+        # For adapting this with different resolution, you would need to re-scale intrinsics below.
+        image0 = cv2.resize(cv2.imread(f"{self.root_dir}/{data['pair_names'][0]}"),(w1, h1))
+
+        image1 = cv2.resize(cv2.imread(f"{self.root_dir}/{data['pair_names'][1]}"),(w2, h2))
+
+        data['image0'] = torch.tensor(image0.astype(np.float32)/255).permute(2,0,1)
+        data['image1'] = torch.tensor(image1.astype(np.float32)/255).permute(2,0,1)
+
+        for k,v in data.items():
+            if k not in ('dataset_name', 'scene_id', 'pair_id', 'pair_names', 'size0_hw', 'size1_hw', 'image0', 'image1'):
+                data[k] = torch.tensor(np.array(v, dtype=np.float32))
+
+        return data
+
+if __name__ == "__main__":
+    weights=os.path.join(os.path.dirname(__file__),'../weights/LiftFeat.pth')
+    liftfeat=LiftFeat(weight=weights)
+    
+    dataset = MegaDepth1500(json_file = DATASET_JSON, root_dir = DATASET_ROOT)
+
+    loader = DataLoader(dataset, batch_size=1, shuffle=False)
+
+    metrics = {}
+    R_errs = []
+    t_errs = []
+    inliers = []
+    
+    results=[]
+
+    cur_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+
+    for d in tqdm.tqdm(loader, desc="processing"):
+        error_infos = compute_pose_error(liftfeat.match_liftfeat,d)
+        results.append(error_infos)
+            
+    print(f'\n==={cur_time}==={args.name}===')
+    d_err_auc,errors=compute_maa(results)
+    for s_k,s_v in d_err_auc.items():
+        print(f'{s_k}: {s_v*100}')

vismatch/third_party/LiftFeat/evaluation/eval_utils.py

@@ -0,0 +1,127 @@
+import numpy as np
+import torch
+import poselib
+
+
+def relative_pose_error(T_0to1, R, t, ignore_gt_t_thr=0.0):
+    # angle error between 2 vectors
+    t_gt = T_0to1[:3, 3]
+    n = np.linalg.norm(t) * np.linalg.norm(t_gt)
+    t_err = np.rad2deg(np.arccos(np.clip(np.dot(t, t_gt) / n, -1.0, 1.0)))
+    t_err = np.minimum(t_err, 180 - t_err)  # handle E ambiguity
+    if np.linalg.norm(t_gt) < ignore_gt_t_thr:  # pure rotation is challenging
+        t_err = 0
+
+    # angle error between 2 rotation matrices
+    R_gt = T_0to1[:3, :3]
+    cos = (np.trace(np.dot(R.T, R_gt)) - 1) / 2
+    cos = np.clip(cos, -1.0, 1.0)  # handle numercial errors
+    R_err = np.rad2deg(np.abs(np.arccos(cos)))
+
+    return t_err, R_err
+
+def intrinsics_to_camera(K):
+    px, py = K[0, 2], K[1, 2]
+    fx, fy = K[0, 0], K[1, 1]
+    return {
+        "model": "PINHOLE",
+        "width": int(2 * px),
+        "height": int(2 * py),
+        "params": [fx, fy, px, py],
+    }
+
+
+def estimate_pose(kpts0, kpts1, K0, K1, thresh, conf=0.99999):
+    M, info = poselib.estimate_relative_pose(
+        kpts0, kpts1,
+        intrinsics_to_camera(K0),
+        intrinsics_to_camera(K1),
+        {"max_epipolar_error": thresh,
+         "success_prob": conf,
+         "min_iterations": 20,
+         "max_iterations": 1_000},
+    )
+
+    R, t, inl = M.R, M.t, info["inliers"]
+    inl = np.array(inl)
+    ret = (R, t, inl)
+
+    return ret
+
+def tensor2bgr(t):
+    return (t.cpu()[0].permute(1,2,0).numpy()*255).astype(np.uint8)
+
+def compute_pose_error(match_fn,data):
+    result = {}
+    
+    with torch.no_grad():
+        mkpts0,mkpts1=match_fn(tensor2bgr(data["image0"]),tensor2bgr(data["image1"]))
+
+    mkpts0=mkpts0 * data["scale0"].numpy()
+    mkpts1=mkpts1 * data["scale1"].numpy()
+
+    K0, K1 = data["K0"][0].numpy(), data["K1"][0].numpy()
+    T_0to1 = data["T_0to1"][0].numpy()
+    T_1to0 = data["T_1to0"][0].numpy()
+
+    result={}
+    conf = 0.99999
+    
+    ret = estimate_pose(mkpts0,mkpts1,K0,K1,4.0,conf)
+    if ret is not None:
+        R, t, inliers = ret
+        t_err, R_err = relative_pose_error(T_0to1, R, t, ignore_gt_t_thr=0.0)
+        result['R_err'] = R_err
+        result['t_err'] = t_err
+
+    return result
+
+
+def error_auc(errors, thresholds=[5, 10, 20]):
+    """
+    Args:
+        errors (list): [N,]
+        thresholds (list)
+    """
+    errors = [0] + sorted(list(errors))
+    recall = list(np.linspace(0, 1, len(errors)))
+
+    aucs = []
+
+    for thr in thresholds:
+        last_index = np.searchsorted(errors, thr)
+        y = recall[:last_index] + [recall[last_index-1]]
+        x = errors[:last_index] + [thr]
+        aucs.append(np.trapz(y, x) / thr)
+
+    return {f'auc@{t}': auc for t, auc in zip(thresholds, aucs)}
+
+def compute_maa(pairs, thresholds=[5, 10, 20]):
+    # print("auc / mAcc on %d pairs" % (len(pairs)))
+    errors = []
+
+    for p in pairs:
+        et = p['t_err']
+        er = p['R_err']
+        errors.append(max(et, er))
+
+    d_err_auc = error_auc(errors)
+
+    # for k,v in d_err_auc.items():
+    #     print(k, ': ', '%.1f'%(v*100))
+
+    errors = np.array(errors)
+
+    for t in thresholds:
+        acc = (errors <= t).sum() / len(errors)
+        # print("mAcc@%d: %.1f "%(t, acc*100))
+        
+    return d_err_auc,errors
+        
+def homo_trans(coord, H):
+    kpt_num = coord.shape[0]
+    homo_coord = np.concatenate((coord, np.ones((kpt_num, 1))), axis=-1)
+    proj_coord = np.matmul(H, homo_coord.T).T
+    proj_coord = proj_coord / proj_coord[:, 2][..., None]
+    proj_coord = proj_coord[:, 0:2]
+    return proj_coord    

vismatch/third_party/LiftFeat/loss/loss.py

@@ -0,0 +1,291 @@
+import os
+import sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import time
+
+
+def dual_softmax_loss(X, Y, temp = 0.2):
+    if X.size() != Y.size() or X.dim() != 2 or Y.dim() != 2:
+        raise RuntimeError('Error: X and Y shapes must match and be 2D matrices')
+
+    dist_mat = (X @ Y.t()) * temp
+    conf_matrix12 = F.log_softmax(dist_mat, dim=1)
+    conf_matrix21 = F.log_softmax(dist_mat.t(), dim=1)
+
+    with torch.no_grad():
+        conf12 = torch.exp( conf_matrix12 ).max(dim=-1)[0]
+        conf21 = torch.exp( conf_matrix21 ).max(dim=-1)[0]
+        conf = conf12 * conf21
+
+    target = torch.arange(len(X), device = X.device)
+
+    loss = F.nll_loss(conf_matrix12, target) + \
+           F.nll_loss(conf_matrix21, target)
+
+    return loss, conf
+
+
+class LiftFeatLoss(nn.Module):
+    def __init__(self,device,lam_descs=1,lam_fb_descs=1,lam_kpts=1,lam_heatmap=1,lam_normals=1,lam_coordinates=1,lam_fb_coordinates=1,depth_spvs=False):
+        super().__init__()
+        
+        # loss parameters
+        self.lam_descs=lam_descs
+        self.lam_fb_descs=lam_fb_descs
+        self.lam_kpts=lam_kpts
+        self.lam_heatmap=lam_heatmap
+        self.lam_normals=lam_normals
+        self.lam_coordinates=lam_coordinates
+        self.lam_fb_coordinates=lam_fb_coordinates
+        self.depth_spvs=depth_spvs
+        self.running_descs_loss=0
+        self.running_kpts_loss=0
+        self.running_heatmaps_loss=0
+        self.loss_descs=0
+        self.loss_fb_descs=0
+        self.loss_kpts=0
+        self.loss_heatmaps=0
+        self.loss_normals=0
+        self.loss_coordinates=0
+        self.loss_fb_coordinates=0
+        self.acc_coarse=0
+        self.acc_fb_coarse=0
+        self.acc_kpt=0
+        self.acc_coordinates=0
+        self.acc_fb_coordinates=0
+        
+        # device
+        self.dev=device
+        
+    
+    def check_accuracy(self,m1,m2,pts1=None,pts2=None,plot=False):
+        with torch.no_grad():
+            #dist_mat = torch.cdist(X,Y)
+            dist_mat = m1 @ m2.t()
+            nn = torch.argmax(dist_mat, dim=1)
+            #nn = torch.argmin(dist_mat, dim=1)
+            correct = nn == torch.arange(len(m1), device = m1.device)
+
+            if pts1 is not None and plot:
+                import matplotlib.pyplot as plt
+                canvas = torch.zeros((60, 80),device=m1.device)
+                pts1 = pts1[~correct]
+                canvas[pts1[:,1].long(), pts1[:,0].long()] = 1
+                canvas = canvas.cpu().numpy()
+                plt.imshow(canvas), plt.show()
+
+            acc = correct.sum().item() / len(m1)
+            return acc    
+    
+    def compute_descriptors_loss(self,descs1,descs2,pts):
+        loss=[]
+        acc=0
+        B,_,H,W=descs1.shape
+        conf_list=[]
+        
+        for b in range(B):
+            pts1,pts2=pts[b][:,:2],pts[b][:,2:]
+            m1=descs1[b,:,pts1[:,1].long(),pts1[:,0].long()].permute(1,0)
+            m2=descs2[b,:,pts2[:,1].long(),pts2[:,0].long()].permute(1,0)
+            
+            loss_per,conf_per=dual_softmax_loss(m1,m2)
+            loss.append(loss_per.unsqueeze(0))
+            conf_list.append(conf_per)
+            
+            acc_coarse_per=self.check_accuracy(m1,m2)
+            acc += acc_coarse_per
+            
+        loss=torch.cat(loss,dim=-1).mean()
+        acc /= B
+        return loss,acc,conf_list
+    
+    
+    def alike_distill_loss(self,kpts,alike_kpts):
+        C, H, W = kpts.shape
+        kpts = kpts.permute(1,2,0)
+        # get ALike keypoints
+        with torch.no_grad():
+            labels = torch.ones((H, W), dtype = torch.long, device = kpts.device) * 64 # -> Default is non-keypoint (bin 64)
+            offsets = (((alike_kpts/8) - (alike_kpts/8).long())*8).long()
+            offsets =  offsets[:, 0] + 8*offsets[:, 1]  # Linear IDX
+            labels[(alike_kpts[:,1]/8).long(), (alike_kpts[:,0]/8).long()] = offsets
+
+        kpts = kpts.view(-1,C)
+        labels = labels.view(-1)
+
+        mask = labels < 64
+        idxs_pos = mask.nonzero().flatten()
+        idxs_neg = (~mask).nonzero().flatten()
+        perm = torch.randperm(idxs_neg.size(0))[:len(idxs_pos)//32]
+        idxs_neg = idxs_neg[perm]
+        idxs = torch.cat([idxs_pos, idxs_neg])
+
+        kpts = kpts[idxs]
+        labels = labels[idxs]
+
+        with torch.no_grad():
+            predicted = kpts.max(dim=-1)[1]
+            acc =  (labels == predicted)
+            acc = acc.sum() / len(acc)
+
+        kpts = F.log_softmax(kpts,dim=-1)
+        loss = F.nll_loss(kpts, labels, reduction = 'mean')
+
+        return loss, acc
+    
+    
+    def compute_keypoints_loss(self,kpts1,kpts2,alike_kpts1,alike_kpts2):
+        loss=[]
+        acc=0
+        B,_,H,W=kpts1.shape
+        
+        for b in range(B):
+            loss_per1,acc_per1=self.alike_distill_loss(kpts1[b],alike_kpts1[b])
+            loss_per2,acc_per2=self.alike_distill_loss(kpts2[b],alike_kpts2[b])
+            loss_per=(loss_per1+loss_per2)
+            acc_per=(acc_per1+acc_per2)/2
+            loss.append(loss_per.unsqueeze(0))
+            acc += acc_per
+            
+        loss=torch.cat(loss,dim=-1).mean()
+        acc /= B
+        return loss,acc
+    
+    
+    def compute_heatmaps_loss(self,heatmaps1,heatmaps2,pts,conf_list):
+        loss=[]
+        B,_,H,W=heatmaps1.shape
+        
+        for b in range(B):
+            pts1,pts2=pts[b][:,:2],pts[b][:,2:]
+            h1=heatmaps1[b,0,pts1[:,1].long(),pts1[:,0].long()]
+            h2=heatmaps2[b,0,pts2[:,1].long(),pts2[:,0].long()]
+            
+            conf=conf_list[b]
+            loss_per1=F.l1_loss(h1,conf)
+            loss_per2=F.l1_loss(h2,conf)
+            loss_per=(loss_per1+loss_per2)
+            loss.append(loss_per.unsqueeze(0))
+            
+        loss=torch.cat(loss,dim=-1).mean()
+        return loss
+    
+    
+    def normal_loss(self,normal,target_normal):
+        # import pdb;pdb.set_trace()
+        normal = normal.permute(1, 2, 0)
+        target_normal = target_normal.permute(1,2,0)
+        # loss = F.l1_loss(d_feat, depth_anything_normal_feat)
+        dot = torch.cosine_similarity(normal, target_normal, dim=2)
+        valid_mask = target_normal[:, :, 0].float() \
+                    * (dot.detach() < 0.999).float() \
+                    * (dot.detach() > -0.999).float()
+        valid_mask = valid_mask > 0.0
+        al = torch.acos(dot[valid_mask])
+        loss = torch.mean(al)
+        return loss
+    
+    
+    def compute_normals_loss(self,normals1,normals2,DA_normals1,DA_normals2,megadepth_batch_size,coco_batch_size):
+        loss=[]
+        
+        # import pdb;pdb.set_trace()
+        
+        # only MegaDepth image need depth-normal
+        normals1=normals1[coco_batch_size:,...]
+        normals2=normals2[coco_batch_size:,...]
+        for b in range(len(DA_normals1)):
+            normal1,normal2=normals1[b],normals2[b]
+            loss_per1=self.normal_loss(normal1,DA_normals1[b].permute(2,0,1))
+            loss_per2=self.normal_loss(normal2,DA_normals2[b].permute(2,0,1))
+            loss_per=(loss_per1+loss_per2)
+            loss.append(loss_per.unsqueeze(0))
+        
+        loss=torch.cat(loss,dim=-1).mean()
+        return loss
+    
+    
+    def coordinate_loss(self,coordinate,conf,pts1):
+        with torch.no_grad():
+            coordinate_detached = pts1 * 8
+            offset_detached = (coordinate_detached/8) - (coordinate_detached/8).long()
+            offset_detached = (offset_detached * 8).long()
+            label = offset_detached[:, 0] + 8*offset_detached[:, 1]
+
+        #pdb.set_trace()
+        coordinate_log = F.log_softmax(coordinate, dim=-1)
+
+        predicted = coordinate.max(dim=-1)[1]
+        acc =  (label == predicted)
+        acc = acc[conf > 0.1]
+        acc = acc.sum() / len(acc)
+
+        loss = F.nll_loss(coordinate_log, label, reduction = 'none')
+        
+        #Weight loss by confidence, giving more emphasis on reliable matches
+        conf = conf / conf.sum()
+        loss = (loss * conf).sum()
+
+        return loss*2., acc
+    
+    def compute_coordinates_loss(self,coordinates,pts,conf_list):
+        loss=[]
+        acc=0
+        B,_,H,W=coordinates.shape
+        
+        for b in range(B):
+            pts1,pts2=pts[b][:,:2],pts[b][:,2:]
+            coordinate=coordinates[b,:,pts1[:,1].long(),pts1[:,0].long()].permute(1,0)
+            conf=conf_list[b]
+            
+            loss_per,acc_per=self.coordinate_loss(coordinate,conf,pts1)
+            loss.append(loss_per.unsqueeze(0))
+            acc += acc_per
+            
+        loss=torch.cat(loss,dim=-1).mean()
+        acc /= B
+        
+        return loss,acc
+        
+        
+    def forward(self,
+                descs1,fb_descs1,kpts1,normals1,
+                descs2,fb_descs2,kpts2,normals2,
+                pts,coordinates,fb_coordinates,
+                alike_kpts1,alike_kpts2,
+                DA_normals1,DA_normals2,
+                megadepth_batch_size,coco_batch_size
+                ):
+        # import pdb;pdb.set_trace()
+        self.loss_descs,self.acc_coarse,conf_list=self.compute_descriptors_loss(descs1,descs2,pts)
+        self.loss_fb_descs,self.acc_fb_coarse,fb_conf_list=self.compute_descriptors_loss(fb_descs1,fb_descs2,pts)
+        
+        # start=time.perf_counter()
+        self.loss_kpts,self.acc_kpt=self.compute_keypoints_loss(kpts1,kpts2,alike_kpts1,alike_kpts2)
+        # end=time.perf_counter()
+        # print(f"kpts loss cost {end-start} seconds")
+        
+        # start=time.perf_counter()
+        self.loss_normals=self.compute_normals_loss(normals1,normals2,DA_normals1,DA_normals2,megadepth_batch_size,coco_batch_size)
+        # end=time.perf_counter()
+        # print(f"normal loss cost {end-start} seconds")
+        
+        self.loss_coordinates,self.acc_coordinates=self.compute_coordinates_loss(coordinates,pts,conf_list)
+        self.loss_fb_coordinates,self.acc_fb_coordinates=self.compute_coordinates_loss(fb_coordinates,pts,fb_conf_list)
+        
+        return {
+            'loss_descs':self.lam_descs*self.loss_descs,
+            'acc_coarse':self.acc_coarse,
+            'loss_coordinates':self.lam_coordinates*self.loss_coordinates,
+            'acc_coordinates':self.acc_coordinates,
+            'loss_fb_descs':self.lam_fb_descs*self.loss_fb_descs,
+            'acc_fb_coarse':self.acc_fb_coarse,
+            'loss_fb_coordinates':self.lam_fb_coordinates*self.loss_fb_coordinates,
+            'acc_fb_coordinates':self.acc_fb_coordinates,
+            'loss_kpts':self.lam_kpts*self.loss_kpts,
+            'acc_kpt':self.acc_kpt,
+            'loss_normals':self.lam_normals*self.loss_normals,
+        }
+

vismatch/third_party/LiftFeat/models/interpolator.py

@@ -0,0 +1,34 @@
+"""
+	"LiftFeat: 3D Geometry-Aware Local Feature Matching"
+ 
+    This script is used to interpolate rough descriptors from LiftFeat
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class InterpolateSparse2d(nn.Module):
+    """ Efficiently interpolate tensor at given sparse 2D positions. """ 
+    def __init__(self, mode = 'bicubic', align_corners = False): 
+        super().__init__()
+        self.mode = mode
+        self.align_corners = align_corners
+
+    def normgrid(self, x, H, W):
+        """ Normalize coords to [-1,1]. """
+        return 2. * (x/(torch.tensor([W-1, H-1], device = x.device, dtype = x.dtype))) - 1.
+
+    def forward(self, x, pos, H, W):
+        """
+        Input
+            x: [B, C, H, W] feature tensor
+            pos: [B, N, 2] tensor of positions
+            H, W: int, original resolution of input 2d positions -- used in normalization [-1,1]
+
+        Returns
+            [B, N, C] sampled channels at 2d positions
+        """
+        grid = self.normgrid(pos, H, W).unsqueeze(-2).to(x.dtype)
+        x = F.grid_sample(x, grid, mode = self.mode , align_corners = False)
+        return x.permute(0,2,3,1).squeeze(-2)