fix: cpu running

hloc/extractors/dedode.py

@@ -64,8 +64,8 @@ class DeDoDe(BaseModel):
         # load the model
         weights_detector = torch.load(model_detector_path, map_location="cpu")
         weights_descriptor = torch.load(model_descriptor_path, map_location="cpu")
-        self.detector = dedode_detector_L(weights=weights_detector)
-        self.descriptor = dedode_descriptor_B(weights=weights_descriptor)
+        self.detector = dedode_detector_L(weights=weights_detector, device = device)
+        self.descriptor = dedode_descriptor_B(weights=weights_descriptor, device = device)
         logger.info(f"Load DeDoDe model done.")
 
     def _forward(self, data):

hloc/matchers/sgmnet.py

@@ -118,7 +118,7 @@ class SGMNet(BaseModel):
             index[:, 0],
             index2.squeeze(0),
         )
-        mask_mc = index2[index] == torch.arange(len(p)).cuda()
+        mask_mc = index2[index] == torch.arange(len(p)).to(device)
         mask = mask_th & mask_mc
         indices0 = torch.where(mask, index, index.new_tensor(-1))
         return indices0

third_party/ASpanFormer/src/ASpanFormer/aspan_module/transformer.py

@@ -4,6 +4,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from .attention import FullAttention, HierachicalAttention, layernorm2d
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 class messageLayer_ini(nn.Module):
     def __init__(self, d_model, d_flow, d_value, nhead):
@@ -171,7 +172,7 @@ class messageLayer_gla(nn.Module):
 
     def decode_flow(self, flow_feature, kshape):
         bs, h, w = flow_feature.shape[0], flow_feature.shape[2], flow_feature.shape[3]
-        scale_factor = torch.tensor([kshape[1], kshape[0]]).cuda()[None, None, None]
+        scale_factor = torch.tensor([kshape[1], kshape[0]]).to(device)[None, None, None]
         flow = (
             self.flow_decoder(flow_feature.view(bs, -1, h * w))
             .permute(0, 2, 1)

third_party/ASpanFormer/src/ASpanFormer/aspanformer.py

@@ -13,6 +13,7 @@ from .aspan_module import (
 from .utils.coarse_matching import CoarseMatching
 from .utils.fine_matching import FineMatching
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 class ASpanFormer(nn.Module):
     def __init__(self, config):
@@ -159,13 +160,14 @@ class ASpanFormer(nn.Module):
             train_res_h / data["image1"].shape[2],
             train_res_w / data["image1"].shape[3],
         ]
+        
         data["online_resize_scale0"], data["online_resize_scale1"] = (
             torch.tensor([w0 / data["image0"].shape[3], h0 / data["image0"].shape[2]])[
                 None
-            ].cuda(),
+            ].to(device),
             torch.tensor([w1 / data["image1"].shape[3], h1 / data["image1"].shape[2]])[
                 None
-            ].cuda(),
+            ].to(device),
         )
 
     def resize_df(self, image, df=32):

third_party/DarkFeat/darkfeat.py

@@ -260,7 +260,7 @@ class DarkFeat(nn.Module):
         )
         self.clf = nn.Conv2d(128, 2, kernel_size=1)
 
-        state_dict = torch.load(self.config["model_path"])
+        state_dict = torch.load(self.config["model_path"], map_location="cpu")
         new_state_dict = {}
 
         for key in state_dict:

third_party/SGMNet/sgmnet/match_model.py

@@ -3,6 +3,7 @@ import torch.nn as nn
 
 eps = 1e-8
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 def sinkhorn(M, r, c, iteration):
     p = torch.softmax(M, dim=-1)
@@ -18,10 +19,10 @@ def sinkhorn(M, r, c, iteration):
 def sink_algorithm(M, dustbin, iteration):
     M = torch.cat([M, dustbin.expand([M.shape[0], M.shape[1], 1])], dim=-1)
     M = torch.cat([M, dustbin.expand([M.shape[0], 1, M.shape[2]])], dim=-2)
-    r = torch.ones([M.shape[0], M.shape[1] - 1], device="cuda")
-    r = torch.cat([r, torch.ones([M.shape[0], 1], device="cuda") * M.shape[1]], dim=-1)
-    c = torch.ones([M.shape[0], M.shape[2] - 1], device="cuda")
-    c = torch.cat([c, torch.ones([M.shape[0], 1], device="cuda") * M.shape[2]], dim=-1)
+    r = torch.ones([M.shape[0], M.shape[1] - 1], device=device)
+    r = torch.cat([r, torch.ones([M.shape[0], 1], device=device) * M.shape[1]], dim=-1)
+    c = torch.ones([M.shape[0], M.shape[2] - 1], device=device)
+    c = torch.cat([c, torch.ones([M.shape[0], 1], device=device) * M.shape[2]], dim=-1)
     p = sinkhorn(M, r, c, iteration)
     return p
 
@@ -42,7 +43,7 @@ def seeding(
     # apply mutual check before nms
     if use_mc:
         mask_not_mutual = nn_index2.gather(dim=-1, index=nn_index1) != torch.arange(
-            nn_index1.shape[1], device="cuda"
+            nn_index1.shape[1], device=device
         )
         match_score[mask_not_mutual] = -1
     # NMS

third_party/lanet/augmentations.py

@@ -12,7 +12,7 @@ import torchvision
 import torchvision.transforms as transforms
 from PIL import Image
 
-from utils import image_grid
+from lanet_utils import image_grid
 
 
 def filter_dict(dict, keywords):

third_party/lanet/data_loader.py

@@ -2,7 +2,7 @@ from PIL import Image
 from torch.utils.data import Dataset, DataLoader
 
 from augmentations import ha_augment_sample, resize_sample, spatial_augment_sample
-from utils import to_tensor_sample
+from lanet_utils import to_tensor_sample
 
 
 def image_transforms(shape, jittering):

third_party/lanet/evaluation/descriptor_evaluation.py

@@ -8,7 +8,7 @@ from os import path as osp
 import cv2
 import numpy as np
 
-from utils import warp_keypoints
+from lanet_utils import warp_keypoints
 
 
 def select_k_best(points, descriptors, k):

third_party/lanet/evaluation/detector_evaluation.py

@@ -8,7 +8,7 @@ from os import path as osp
 import cv2
 import numpy as np
 
-from utils import warp_keypoints
+from lanet_utils import warp_keypoints
 
 
 def compute_repeatability(data, keep_k_points=300, distance_thresh=3):

third_party/lanet/main.py

@@ -2,7 +2,7 @@ import torch
 
 from train import Trainer
 from config import get_config
-from utils import prepare_dirs
+from lanet_utils import prepare_dirs
 from data_loader import get_data_loader
 
 

third_party/lanet/network_v0/modules.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from utils import image_grid
+from lanet_utils import image_grid
 
 
 class ConvBlock(nn.Module):

third_party/lanet/network_v1/modules.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from torchvision import models
-from utils import image_grid
+from lanet_utils import image_grid
 
 
 class ConvBlock(nn.Module):