Vincentqyw commited on
Commit
514cd02
1 Parent(s): c74a070

fix: cpu running

Browse files
hloc/extractors/dedode.py CHANGED
@@ -64,8 +64,8 @@ class DeDoDe(BaseModel):
64
  # load the model
65
  weights_detector = torch.load(model_detector_path, map_location="cpu")
66
  weights_descriptor = torch.load(model_descriptor_path, map_location="cpu")
67
- self.detector = dedode_detector_L(weights=weights_detector)
68
- self.descriptor = dedode_descriptor_B(weights=weights_descriptor)
69
  logger.info(f"Load DeDoDe model done.")
70
 
71
  def _forward(self, data):
 
64
  # load the model
65
  weights_detector = torch.load(model_detector_path, map_location="cpu")
66
  weights_descriptor = torch.load(model_descriptor_path, map_location="cpu")
67
+ self.detector = dedode_detector_L(weights=weights_detector, device = device)
68
+ self.descriptor = dedode_descriptor_B(weights=weights_descriptor, device = device)
69
  logger.info(f"Load DeDoDe model done.")
70
 
71
  def _forward(self, data):
third_party/ASpanFormer/src/ASpanFormer/aspanformer.py CHANGED
@@ -159,14 +159,24 @@ class ASpanFormer(nn.Module):
159
  train_res_h / data["image1"].shape[2],
160
  train_res_w / data["image1"].shape[3],
161
  ]
162
- data["online_resize_scale0"], data["online_resize_scale1"] = (
163
- torch.tensor([w0 / data["image0"].shape[3], h0 / data["image0"].shape[2]])[
164
- None
165
- ].cuda(),
166
- torch.tensor([w1 / data["image1"].shape[3], h1 / data["image1"].shape[2]])[
167
- None
168
- ].cuda(),
169
- )
 
 
 
 
 
 
 
 
 
 
170
 
171
  def resize_df(self, image, df=32):
172
  h, w = image.shape[2], image.shape[3]
 
159
  train_res_h / data["image1"].shape[2],
160
  train_res_w / data["image1"].shape[3],
161
  ]
162
+ if torch.cuda.is_available():
163
+ data["online_resize_scale0"], data["online_resize_scale1"] = (
164
+ torch.tensor([w0 / data["image0"].shape[3], h0 / data["image0"].shape[2]])[
165
+ None
166
+ ].cuda(),
167
+ torch.tensor([w1 / data["image1"].shape[3], h1 / data["image1"].shape[2]])[
168
+ None
169
+ ].cuda(),
170
+ )
171
+ else:
172
+ data["online_resize_scale0"], data["online_resize_scale1"] = (
173
+ torch.tensor([w0 / data["image0"].shape[3], h0 / data["image0"].shape[2]])[
174
+ None
175
+ ],
176
+ torch.tensor([w1 / data["image1"].shape[3], h1 / data["image1"].shape[2]])[
177
+ None
178
+ ],
179
+ )
180
 
181
  def resize_df(self, image, df=32):
182
  h, w = image.shape[2], image.shape[3]
third_party/DarkFeat/darkfeat.py CHANGED
@@ -260,7 +260,7 @@ class DarkFeat(nn.Module):
260
  )
261
  self.clf = nn.Conv2d(128, 2, kernel_size=1)
262
 
263
- state_dict = torch.load(self.config["model_path"])
264
  new_state_dict = {}
265
 
266
  for key in state_dict:
 
260
  )
261
  self.clf = nn.Conv2d(128, 2, kernel_size=1)
262
 
263
+ state_dict = torch.load(self.config["model_path"], map_location="cpu")
264
  new_state_dict = {}
265
 
266
  for key in state_dict:
third_party/SGMNet/sgmnet/match_model.py CHANGED
@@ -3,6 +3,7 @@ import torch.nn as nn
3
 
4
  eps = 1e-8
5
 
 
6
 
7
  def sinkhorn(M, r, c, iteration):
8
  p = torch.softmax(M, dim=-1)
@@ -18,10 +19,10 @@ def sinkhorn(M, r, c, iteration):
18
  def sink_algorithm(M, dustbin, iteration):
19
  M = torch.cat([M, dustbin.expand([M.shape[0], M.shape[1], 1])], dim=-1)
20
  M = torch.cat([M, dustbin.expand([M.shape[0], 1, M.shape[2]])], dim=-2)
21
- r = torch.ones([M.shape[0], M.shape[1] - 1], device="cuda")
22
- r = torch.cat([r, torch.ones([M.shape[0], 1], device="cuda") * M.shape[1]], dim=-1)
23
- c = torch.ones([M.shape[0], M.shape[2] - 1], device="cuda")
24
- c = torch.cat([c, torch.ones([M.shape[0], 1], device="cuda") * M.shape[2]], dim=-1)
25
  p = sinkhorn(M, r, c, iteration)
26
  return p
27
 
@@ -42,7 +43,7 @@ def seeding(
42
  # apply mutual check before nms
43
  if use_mc:
44
  mask_not_mutual = nn_index2.gather(dim=-1, index=nn_index1) != torch.arange(
45
- nn_index1.shape[1], device="cuda"
46
  )
47
  match_score[mask_not_mutual] = -1
48
  # NMS
 
3
 
4
  eps = 1e-8
5
 
6
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7
 
8
  def sinkhorn(M, r, c, iteration):
9
  p = torch.softmax(M, dim=-1)
 
19
  def sink_algorithm(M, dustbin, iteration):
20
  M = torch.cat([M, dustbin.expand([M.shape[0], M.shape[1], 1])], dim=-1)
21
  M = torch.cat([M, dustbin.expand([M.shape[0], 1, M.shape[2]])], dim=-2)
22
+ r = torch.ones([M.shape[0], M.shape[1] - 1], device=device)
23
+ r = torch.cat([r, torch.ones([M.shape[0], 1], device=device) * M.shape[1]], dim=-1)
24
+ c = torch.ones([M.shape[0], M.shape[2] - 1], device=device)
25
+ c = torch.cat([c, torch.ones([M.shape[0], 1], device=device) * M.shape[2]], dim=-1)
26
  p = sinkhorn(M, r, c, iteration)
27
  return p
28
 
 
43
  # apply mutual check before nms
44
  if use_mc:
45
  mask_not_mutual = nn_index2.gather(dim=-1, index=nn_index1) != torch.arange(
46
+ nn_index1.shape[1], device=device
47
  )
48
  match_score[mask_not_mutual] = -1
49
  # NMS
third_party/lanet/augmentations.py CHANGED
@@ -12,7 +12,7 @@ import torchvision
12
  import torchvision.transforms as transforms
13
  from PIL import Image
14
 
15
- from utils import image_grid
16
 
17
 
18
  def filter_dict(dict, keywords):
 
12
  import torchvision.transforms as transforms
13
  from PIL import Image
14
 
15
+ from lanet_utils import image_grid
16
 
17
 
18
  def filter_dict(dict, keywords):
third_party/lanet/data_loader.py CHANGED
@@ -2,7 +2,7 @@ from PIL import Image
2
  from torch.utils.data import Dataset, DataLoader
3
 
4
  from augmentations import ha_augment_sample, resize_sample, spatial_augment_sample
5
- from utils import to_tensor_sample
6
 
7
 
8
  def image_transforms(shape, jittering):
 
2
  from torch.utils.data import Dataset, DataLoader
3
 
4
  from augmentations import ha_augment_sample, resize_sample, spatial_augment_sample
5
+ from lanet_utils import to_tensor_sample
6
 
7
 
8
  def image_transforms(shape, jittering):
third_party/lanet/evaluation/descriptor_evaluation.py CHANGED
@@ -8,7 +8,7 @@ from os import path as osp
8
  import cv2
9
  import numpy as np
10
 
11
- from utils import warp_keypoints
12
 
13
 
14
  def select_k_best(points, descriptors, k):
 
8
  import cv2
9
  import numpy as np
10
 
11
+ from lanet_utils import warp_keypoints
12
 
13
 
14
  def select_k_best(points, descriptors, k):
third_party/lanet/evaluation/detector_evaluation.py CHANGED
@@ -8,7 +8,7 @@ from os import path as osp
8
  import cv2
9
  import numpy as np
10
 
11
- from utils import warp_keypoints
12
 
13
 
14
  def compute_repeatability(data, keep_k_points=300, distance_thresh=3):
 
8
  import cv2
9
  import numpy as np
10
 
11
+ from lanet_utils import warp_keypoints
12
 
13
 
14
  def compute_repeatability(data, keep_k_points=300, distance_thresh=3):
third_party/lanet/{utils.py → lanet_utils.py} RENAMED
File without changes
third_party/lanet/main.py CHANGED
@@ -2,7 +2,7 @@ import torch
2
 
3
  from train import Trainer
4
  from config import get_config
5
- from utils import prepare_dirs
6
  from data_loader import get_data_loader
7
 
8
 
 
2
 
3
  from train import Trainer
4
  from config import get_config
5
+ from lanet_utils import prepare_dirs
6
  from data_loader import get_data_loader
7
 
8
 
third_party/lanet/network_v0/modules.py CHANGED
@@ -2,7 +2,7 @@ import torch
2
  import torch.nn as nn
3
  import torch.nn.functional as F
4
 
5
- from utils import image_grid
6
 
7
 
8
  class ConvBlock(nn.Module):
 
2
  import torch.nn as nn
3
  import torch.nn.functional as F
4
 
5
+ from lanet_utils import image_grid
6
 
7
 
8
  class ConvBlock(nn.Module):
third_party/lanet/network_v1/modules.py CHANGED
@@ -4,7 +4,7 @@ import torch.nn as nn
4
  import torch.nn.functional as F
5
 
6
  from torchvision import models
7
- from utils import image_grid
8
 
9
 
10
  class ConvBlock(nn.Module):
 
4
  import torch.nn.functional as F
5
 
6
  from torchvision import models
7
+ from lanet_utils import image_grid
8
 
9
 
10
  class ConvBlock(nn.Module):