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):

third_party/ASpanFormer/src/ASpanFormer/aspanformer.py

@@ -159,14 +159,24 @@ 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(),
-            torch.tensor([w1 / data["image1"].shape[3], h1 / data["image1"].shape[2]])[
-                None
-            ].cuda(),
-        )
+        if torch.cuda.is_available():
+            data["online_resize_scale0"], data["online_resize_scale1"] = (
+                torch.tensor([w0 / data["image0"].shape[3], h0 / data["image0"].shape[2]])[
+                    None
+                ].cuda(),
+                torch.tensor([w1 / data["image1"].shape[3], h1 / data["image1"].shape[2]])[
+                    None
+                ].cuda(),
+            )
+        else:
+            data["online_resize_scale0"], data["online_resize_scale1"] = (
+                torch.tensor([w0 / data["image0"].shape[3], h0 / data["image0"].shape[2]])[
+                    None
+                ],
+                torch.tensor([w1 / data["image1"].shape[3], h1 / data["image1"].shape[2]])[
+                    None
+                ],
+            )
 
     def resize_df(self, image, df=32):
         h, w = image.shape[2], image.shape[3]

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):