Vincentqyw commited on
Commit
472119d
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):
hloc/matchers/sgmnet.py CHANGED
@@ -118,7 +118,7 @@ class SGMNet(BaseModel):
118
  index[:, 0],
119
  index2.squeeze(0),
120
  )
121
- mask_mc = index2[index] == torch.arange(len(p)).cuda()
122
  mask = mask_th & mask_mc
123
  indices0 = torch.where(mask, index, index.new_tensor(-1))
124
  return indices0
 
118
  index[:, 0],
119
  index2.squeeze(0),
120
  )
121
+ mask_mc = index2[index] == torch.arange(len(p)).to(device)
122
  mask = mask_th & mask_mc
123
  indices0 = torch.where(mask, index, index.new_tensor(-1))
124
  return indices0
third_party/ASpanFormer/src/ASpanFormer/aspan_module/transformer.py CHANGED
@@ -4,6 +4,7 @@ import torch.nn as nn
4
  import torch.nn.functional as F
5
  from .attention import FullAttention, HierachicalAttention, layernorm2d
6
 
 
7
 
8
  class messageLayer_ini(nn.Module):
9
  def __init__(self, d_model, d_flow, d_value, nhead):
@@ -171,7 +172,7 @@ class messageLayer_gla(nn.Module):
171
 
172
  def decode_flow(self, flow_feature, kshape):
173
  bs, h, w = flow_feature.shape[0], flow_feature.shape[2], flow_feature.shape[3]
174
- scale_factor = torch.tensor([kshape[1], kshape[0]]).cuda()[None, None, None]
175
  flow = (
176
  self.flow_decoder(flow_feature.view(bs, -1, h * w))
177
  .permute(0, 2, 1)
 
4
  import torch.nn.functional as F
5
  from .attention import FullAttention, HierachicalAttention, layernorm2d
6
 
7
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
8
 
9
  class messageLayer_ini(nn.Module):
10
  def __init__(self, d_model, d_flow, d_value, nhead):
 
172
 
173
  def decode_flow(self, flow_feature, kshape):
174
  bs, h, w = flow_feature.shape[0], flow_feature.shape[2], flow_feature.shape[3]
175
+ scale_factor = torch.tensor([kshape[1], kshape[0]]).to(device)[None, None, None]
176
  flow = (
177
  self.flow_decoder(flow_feature.view(bs, -1, h * w))
178
  .permute(0, 2, 1)
third_party/ASpanFormer/src/ASpanFormer/aspanformer.py CHANGED
@@ -13,6 +13,7 @@ from .aspan_module import (
13
  from .utils.coarse_matching import CoarseMatching
14
  from .utils.fine_matching import FineMatching
15
 
 
16
 
17
  class ASpanFormer(nn.Module):
18
  def __init__(self, config):
@@ -159,13 +160,14 @@ 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):
 
13
  from .utils.coarse_matching import CoarseMatching
14
  from .utils.fine_matching import FineMatching
15
 
16
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
17
 
18
  class ASpanFormer(nn.Module):
19
  def __init__(self, config):
 
160
  train_res_h / data["image1"].shape[2],
161
  train_res_w / data["image1"].shape[3],
162
  ]
163
+
164
  data["online_resize_scale0"], data["online_resize_scale1"] = (
165
  torch.tensor([w0 / data["image0"].shape[3], h0 / data["image0"].shape[2]])[
166
  None
167
+ ].to(device),
168
  torch.tensor([w1 / data["image1"].shape[3], h1 / data["image1"].shape[2]])[
169
  None
170
+ ].to(device),
171
  )
172
 
173
  def resize_df(self, image, df=32):
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):