Spaces:
Running
Running
Vincentqyw
commited on
Commit
•
472119d
1
Parent(s):
c74a070
fix: cpu running
Browse files- hloc/extractors/dedode.py +2 -2
- hloc/matchers/sgmnet.py +1 -1
- third_party/ASpanFormer/src/ASpanFormer/aspan_module/transformer.py +2 -1
- third_party/ASpanFormer/src/ASpanFormer/aspanformer.py +4 -2
- third_party/DarkFeat/darkfeat.py +1 -1
- third_party/SGMNet/sgmnet/match_model.py +6 -5
- third_party/lanet/augmentations.py +1 -1
- third_party/lanet/data_loader.py +1 -1
- third_party/lanet/evaluation/descriptor_evaluation.py +1 -1
- third_party/lanet/evaluation/detector_evaluation.py +1 -1
- third_party/lanet/{utils.py → lanet_utils.py} +0 -0
- third_party/lanet/main.py +1 -1
- third_party/lanet/network_v0/modules.py +1 -1
- third_party/lanet/network_v1/modules.py +1 -1
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)).
|
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]]).
|
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 |
-
].
|
166 |
torch.tensor([w1 / data["image1"].shape[3], h1 / data["image1"].shape[2]])[
|
167 |
None
|
168 |
-
].
|
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=
|
22 |
-
r = torch.cat([r, torch.ones([M.shape[0], 1], device=
|
23 |
-
c = torch.ones([M.shape[0], M.shape[2] - 1], device=
|
24 |
-
c = torch.cat([c, torch.ones([M.shape[0], 1], device=
|
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=
|
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
|
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
|
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
|
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
|
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
|
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
|
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
|
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):
|