add: scores to liftfeat

imcui/hloc/extractors/liftfeat.py

@@ -12,12 +12,6 @@ sys.path.append(str(fire_path))
 from models.liftfeat_wrapper import LiftFeat, MODEL_PATH
 
 
-def select_idx(N, M):
-    numbers = list(range(0, N))
-    selected = random.sample(numbers, M)
-    return selected
-
-
 class Liftfeat(BaseModel):
     default_conf = {
         "keypoint_threshold": 0.05,
@@ -42,9 +36,9 @@ class Liftfeat(BaseModel):
 
         keypoints = pred["keypoints"]
         descriptors = pred["descriptors"]
-        scores = torch.ones_like(pred["keypoints"][:, 0])
+        scores = pred["scores"]
         if self.conf["max_keypoints"] < len(keypoints):
-            idxs = select_idx(len(keypoints), self.conf["max_keypoints"])
+            idxs = scores.argsort()[-self.conf["max_keypoints"] or None :]
             keypoints = keypoints[idxs, :2]
             descriptors = descriptors[idxs]
             scores = scores[idxs]

imcui/third_party/LiftFeat/models/liftfeat_wrapper.py

@@ -9,22 +9,21 @@ from models.model import LiftFeatSPModel
 from models.interpolator import InterpolateSparse2d
 from utils.config import featureboost_config
 
-device=torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 
-MODEL_PATH=os.path.join(os.path.dirname(__file__),'../weights/LiftFeat.pth')
+MODEL_PATH = os.path.join(os.path.dirname(__file__), "../weights/LiftFeat.pth")
 
 
 class NonMaxSuppression(torch.nn.Module):
     def __init__(self, rep_thr=0.1, top_k=4096):
-        super(NonMaxSuppression,self).__init__()
+        super(NonMaxSuppression, self).__init__()
         self.max_filter = torch.nn.MaxPool2d(kernel_size=5, stride=1, padding=2)
         self.rep_thr = rep_thr
-        self.top_k=top_k
-        
-        
-    def NMS(self, x, threshold = 0.05, kernel_size = 5):
+        self.top_k = top_k
+
+    def NMS(self, x, threshold=0.05, kernel_size=5):
         B, _, H, W = x.shape
-        pad=kernel_size//2
+        pad = kernel_size // 2
         local_max = nn.MaxPool2d(kernel_size=kernel_size, stride=1, padding=pad)(x)
         pos = (x == local_max) & (x > threshold)
         pos_batched = [k.nonzero()[..., 1:].flip(-1) for k in pos]
@@ -32,17 +31,18 @@ class NonMaxSuppression(torch.nn.Module):
         pad_val = max([len(x) for x in pos_batched])
         pos = torch.zeros((B, pad_val, 2), dtype=torch.long, device=x.device)
 
-        #Pad kpts and build (B, N, 2) tensor
+        # Pad kpts and build (B, N, 2) tensor
         for b in range(len(pos_batched)):
-            pos[b, :len(pos_batched[b]), :] = pos_batched[b]
+            pos[b, : len(pos_batched[b]), :] = pos_batched[b]
 
         return pos
-        
+
     def forward(self, score):
-        pos = self.NMS(score,self.rep_thr)
-        
+        pos = self.NMS(score, self.rep_thr)
+
         return pos
 
+
 def load_model(model, weight_path):
     pretrained_weights = torch.load(weight_path, map_location="cpu")
 
@@ -72,82 +72,82 @@ def load_model(model, weight_path):
 
 
 import torch.nn as nn
+
+
 class LiftFeat(nn.Module):
-    def __init__(self,weight=MODEL_PATH,top_k=4096,detect_threshold=0.1):
+    def __init__(self, weight=MODEL_PATH, top_k=4096, detect_threshold=0.1):
         super().__init__()
-        self.device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        self.net=LiftFeatSPModel(featureboost_config).to(self.device).eval()
-        self.top_k=top_k
-        self.sampler=InterpolateSparse2d('bicubic')
-        self.net=load_model(self.net,weight)
-        self.detector=NonMaxSuppression(rep_thr=detect_threshold)
-        self.net=self.net.to(self.device)
-        self.detector=self.detector.to(self.device)
-        self.sampler=self.sampler.to(self.device)
-        
-    def image_preprocess(self,image: np.ndarray):
-        H,W,C=image.shape[0],image.shape[1],image.shape[2]
-
-        _H=math.ceil(H/32)*32
-        _W=math.ceil(W/32)*32
-
-        pad_h=_H-H
-        pad_w=_W-W
-
-        image=cv2.copyMakeBorder(image,0,pad_h,0,pad_w,cv2.BORDER_CONSTANT,None,(0, 0, 0))
-        
-        pad_info=[0,pad_h,0,pad_w]
-
-        if len(image.shape)==3:
-            image=image[None,...]
-
-        image=torch.tensor(image).permute(0,3,1,2)/255
-        image=image.to(device)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.net = LiftFeatSPModel(featureboost_config).to(self.device).eval()
+        self.top_k = top_k
+        self.sampler = InterpolateSparse2d("bicubic")
+        self.net = load_model(self.net, weight)
+        self.detector = NonMaxSuppression(rep_thr=detect_threshold)
+        self.net = self.net.to(self.device)
+        self.detector = self.detector.to(self.device)
+        self.sampler = self.sampler.to(self.device)
+
+    def image_preprocess(self, image: np.ndarray):
+        H, W, C = image.shape[0], image.shape[1], image.shape[2]
+
+        _H = math.ceil(H / 32) * 32
+        _W = math.ceil(W / 32) * 32
+
+        pad_h = _H - H
+        pad_w = _W - W
+
+        image = cv2.copyMakeBorder(image, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, None, (0, 0, 0))
+
+        pad_info = [0, pad_h, 0, pad_w]
+
+        if len(image.shape) == 3:
+            image = image[None, ...]
+
+        image = torch.tensor(image).permute(0, 3, 1, 2) / 255
+        image = image.to(device)
 
         return image, pad_info
-    
+
     @torch.inference_mode()
-    def extract(self,image: np.ndarray):
-        image,pad_info=self.image_preprocess(image)
-        B,_,_H1,_W1=image.shape
-
-        M1,K1,D1=self.net.forward1(image)
-        refine_M=self.net.forward2(M1,K1,D1)
-        
-        refine_M=refine_M.reshape(M1.shape[0],M1.shape[2],M1.shape[3],-1).permute(0,3,1,2)
-        refine_M=torch.nn.functional.normalize(refine_M,2,dim=1)
-        
-        descs_map=refine_M
-        # descs_map=M1
-        
-        scores=torch.softmax(K1,dim=1)[:,:64]
-        heatmap=scores.permute(0,2,3,1).reshape(scores.shape[0],scores.shape[2],scores.shape[3],8,8)
-        heatmap=heatmap.permute(0,1,3,2,4).reshape(scores.shape[0],1,scores.shape[2]*8,scores.shape[3]*8)
-        
-        pos=self.detector(heatmap)
-        kpts=pos.squeeze(0)
-        mask_w=kpts[...,0]<(_W1-pad_info[-1])
-        kpts=kpts[mask_w]
-        mask_h=kpts[..., 1]<(_H1-pad_info[1])
-        kpts=kpts[mask_h]
-        
-        descs=self.sampler(descs_map,kpts.unsqueeze(0),_H1,_W1)
-        descs=torch.nn.functional.normalize(descs,p=2,dim=1)
-        descs=descs.squeeze(0)
-        
-        return {
-            'descriptors':descs,
-            'keypoints':kpts
-        }
-        
+    def extract(self, image: np.ndarray):
+        image, pad_info = self.image_preprocess(image)
+        B, _, _H1, _W1 = image.shape
+
+        M1, K1, D1 = self.net.forward1(image)
+        refine_M = self.net.forward2(M1, K1, D1)
+
+        refine_M = refine_M.reshape(M1.shape[0], M1.shape[2], M1.shape[3], -1).permute(0, 3, 1, 2)
+        refine_M = torch.nn.functional.normalize(refine_M, 2, dim=1)
+
+        descs_map = refine_M
+
+        scores = torch.softmax(K1, dim=1)[:, :64]
+        heatmap = scores.permute(0, 2, 3, 1).reshape(scores.shape[0], scores.shape[2], scores.shape[3], 8, 8)
+        heatmap = heatmap.permute(0, 1, 3, 2, 4).reshape(scores.shape[0], 1, scores.shape[2] * 8, scores.shape[3] * 8)
+
+        pos = self.detector(heatmap)
+        kpts = pos.squeeze(0)
+        mask_w = kpts[..., 0] < (_W1 - pad_info[-1])
+        kpts = kpts[mask_w]
+        mask_h = kpts[..., 1] < (_H1 - pad_info[1])
+        kpts = kpts[mask_h]
+
+        scores = self.sampler(heatmap, kpts.unsqueeze(0), _H1, _W1)
+        scores = scores.squeeze(0).reshape(-1)
+        descs = self.sampler(descs_map, kpts.unsqueeze(0), _H1, _W1)
+        descs = torch.nn.functional.normalize(descs, p=2, dim=1)
+        descs = descs.squeeze(0)
+
+        return {"descriptors": descs, "keypoints": kpts, "scores": scores}
+
     def match_liftfeat(self, img1, img2, min_cossim=-1):
         # import pdb;pdb.set_trace()
-        data1=self.extract(img1)
-        data2=self.extract(img2)
-        
-        kpts1,feats1=data1['keypoints'],data1['descriptors']
-        kpts2,feats2=data2['keypoints'],data2['descriptors']
-        
+        data1 = self.extract(img1)
+        data2 = self.extract(img2)
+
+        kpts1, feats1 = data1["keypoints"], data1["descriptors"]
+        kpts2, feats2 = data2["keypoints"], data2["descriptors"]
+
         cossim = feats1 @ feats2.t()
         cossim_t = feats2 @ feats1.t()
 
@@ -165,9 +165,8 @@ class LiftFeat(nn.Module):
         else:
             idx0 = idx0[mutual]
             idx1 = match12[mutual]
-            
-        mkpts1,mkpts2=kpts1[idx0],kpts2[idx1]
-        mkpts1,mkpts2=mkpts1.cpu().numpy(),mkpts2.cpu().numpy()
 
-        return mkpts1, mkpts2
+        mkpts1, mkpts2 = kpts1[idx0], kpts2[idx1]
+        mkpts1, mkpts2 = mkpts1.cpu().numpy(), mkpts2.cpu().numpy()
 
+        return mkpts1, mkpts2