Spaces:
Build error
Build error
temp
Browse files- app.py +1 -3
- fire_network.py +3 -15
- how/networks/how_net.py +2 -75
app.py
CHANGED
@@ -2,8 +2,6 @@ import gradio as gr
|
|
2 |
|
3 |
import torch
|
4 |
|
5 |
-
from how.networks import how_net
|
6 |
-
|
7 |
import fire_network
|
8 |
|
9 |
import cv2
|
@@ -30,7 +28,7 @@ sf_idx_ = [55, 14, 5, 4, 52, 57, 40, 9]
|
|
30 |
|
31 |
col = plt.get_cmap('tab10')
|
32 |
|
33 |
-
def generate_matching_superfeatures(im1, im2,
|
34 |
|
35 |
im1_tensor = transform(im1)
|
36 |
im2_tensor = transform(im2)
|
|
|
2 |
|
3 |
import torch
|
4 |
|
|
|
|
|
5 |
import fire_network
|
6 |
|
7 |
import cv2
|
|
|
28 |
|
29 |
col = plt.get_cmap('tab10')
|
30 |
|
31 |
+
def generate_matching_superfeatures(im1, im2, scale_id=6, threshold=50):
|
32 |
|
33 |
im1_tensor = transform(im1)
|
34 |
im2_tensor = transform(im2)
|
fire_network.py
CHANGED
@@ -7,11 +7,10 @@ from torch import nn
|
|
7 |
import torchvision
|
8 |
|
9 |
from how import layers
|
10 |
-
from how.layers import functional as HF
|
11 |
|
12 |
from lit import LocalfeatureIntegrationTransformer
|
13 |
|
14 |
-
from how.networks.how_net import HOWNet
|
15 |
|
16 |
class FIReNet(HOWNet):
|
17 |
|
@@ -62,20 +61,9 @@ class FIReNet(HOWNet):
|
|
62 |
return feats, attns, strengths
|
63 |
|
64 |
def forward(self, x):
|
65 |
-
if self.return_global:
|
66 |
-
return self.forward_global(x, scales=self.runtime['training_scales'])
|
67 |
return self.get_superfeatures(x, scales=self.runtime['training_scales'])
|
68 |
|
69 |
-
|
70 |
-
"""Return global descriptor"""
|
71 |
-
feats, _, strengths = self.get_superfeatures(x, scales=scales)
|
72 |
-
return HF.weighted_spoc(feats, strengths)
|
73 |
-
|
74 |
-
def forward_local(self, x, *, features_num, scales):
|
75 |
-
"""Return selected super features"""
|
76 |
-
feats, _, strengths = self.get_superfeatures(x, scales=scales)
|
77 |
-
return HF.how_select_local(feats, strengths, scales=scales, features_num=features_num)
|
78 |
-
|
79 |
def init_network(architecture, pretrained, skip_layer, dim_reduction, lit, runtime):
|
80 |
"""Initialize FIRe network
|
81 |
:param str architecture: Network backbone architecture (e.g. resnet18)
|
@@ -115,7 +103,7 @@ def init_network(architecture, pretrained, skip_layer, dim_reduction, lit, runti
|
|
115 |
"architecture": architecture,
|
116 |
"backbone_dim": lit['dim'],
|
117 |
"outputdim": reduction_layer.out_channels if dim_reduction else lit['dim'],
|
118 |
-
"corercf_size":
|
119 |
}
|
120 |
net = FIReNet(nn.Sequential(*features), att_layer, lit_layer, reduction_layer, meta, runtime)
|
121 |
|
|
|
7 |
import torchvision
|
8 |
|
9 |
from how import layers
|
|
|
10 |
|
11 |
from lit import LocalfeatureIntegrationTransformer
|
12 |
|
13 |
+
from how.networks.how_net import HOWNet
|
14 |
|
15 |
class FIReNet(HOWNet):
|
16 |
|
|
|
61 |
return feats, attns, strengths
|
62 |
|
63 |
def forward(self, x):
|
|
|
|
|
64 |
return self.get_superfeatures(x, scales=self.runtime['training_scales'])
|
65 |
|
66 |
+
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
67 |
def init_network(architecture, pretrained, skip_layer, dim_reduction, lit, runtime):
|
68 |
"""Initialize FIRe network
|
69 |
:param str architecture: Network backbone architecture (e.g. resnet18)
|
|
|
103 |
"architecture": architecture,
|
104 |
"backbone_dim": lit['dim'],
|
105 |
"outputdim": reduction_layer.out_channels if dim_reduction else lit['dim'],
|
106 |
+
"corercf_size": 32 // (2 ** skip_layer),
|
107 |
}
|
108 |
net = FIReNet(nn.Sequential(*features), att_layer, lit_layer, reduction_layer, meta, runtime)
|
109 |
|
how/networks/how_net.py
CHANGED
@@ -5,20 +5,12 @@ import torch
|
|
5 |
import torch.nn as nn
|
6 |
import torchvision
|
7 |
|
8 |
-
from cirtorch.networks import imageretrievalnet
|
9 |
-
|
10 |
from .. import layers
|
11 |
from ..layers import functional as HF
|
12 |
from ..utils import io_helpers
|
13 |
|
14 |
NUM_WORKERS = 6
|
15 |
|
16 |
-
CORERCF_SIZE = {
|
17 |
-
'resnet18': 32,
|
18 |
-
'resnet50': 32,
|
19 |
-
'resnet101': 32,
|
20 |
-
}
|
21 |
-
|
22 |
|
23 |
class HOWNet(nn.Module):
|
24 |
"""Network for the HOW method
|
@@ -100,22 +92,6 @@ class HOWNet(nn.Module):
|
|
100 |
|
101 |
return feats, masks
|
102 |
|
103 |
-
def forward(self, x):
|
104 |
-
return self.forward_global(x, scales=self.runtime['training_scales'])
|
105 |
-
|
106 |
-
def forward_global(self, x, *, scales):
|
107 |
-
"""Return global descriptor"""
|
108 |
-
feats, masks = self.features_attentions(x, scales=scales)
|
109 |
-
return HF.weighted_spoc(feats, masks)
|
110 |
-
|
111 |
-
def forward_local(self, x, *, features_num, scales):
|
112 |
-
"""Return local descriptors"""
|
113 |
-
feats, masks = self.features_attentions(x, scales=scales)
|
114 |
-
return HF.how_select_local(feats, masks, scales=scales, features_num=features_num)
|
115 |
-
|
116 |
-
|
117 |
-
# String conversion
|
118 |
-
|
119 |
def __repr__(self):
|
120 |
meta_str = "\n".join(" %s: %s" % x for x in self.meta.items())
|
121 |
return "%s(meta={\n%s\n})" % (self.__class__.__name__, meta_str)
|
@@ -151,7 +127,7 @@ def init_network(architecture, pretrained, skip_layer, dim_reduction, smoothing,
|
|
151 |
|
152 |
if skip_layer > 0:
|
153 |
features = features[:-skip_layer]
|
154 |
-
backbone_dim =
|
155 |
|
156 |
att_layer = layers.attention.L2Attention()
|
157 |
smooth_layer = None
|
@@ -165,57 +141,8 @@ def init_network(architecture, pretrained, skip_layer, dim_reduction, smoothing,
|
|
165 |
"architecture": architecture,
|
166 |
"backbone_dim": backbone_dim,
|
167 |
"outputdim": reduction_layer.out_channels if dim_reduction else backbone_dim,
|
168 |
-
"corercf_size":
|
169 |
}
|
170 |
return HOWNet(nn.Sequential(*features), att_layer, smooth_layer, reduction_layer, meta, runtime)
|
171 |
|
172 |
|
173 |
-
def extract_vectors(net, dataset, device, *, scales):
|
174 |
-
"""Return global descriptors in torch.Tensor"""
|
175 |
-
net.eval()
|
176 |
-
loader = torch.utils.data.DataLoader(dataset, shuffle=False, pin_memory=True, num_workers=NUM_WORKERS)
|
177 |
-
|
178 |
-
with torch.no_grad():
|
179 |
-
vecs = torch.zeros(len(loader), net.meta['outputdim'])
|
180 |
-
for i, inp in io_helpers.progress(enumerate(loader), size=len(loader), print_freq=100):
|
181 |
-
vecs[i] = net.forward_global(inp.to(device), scales=scales).cpu().squeeze()
|
182 |
-
|
183 |
-
return vecs
|
184 |
-
|
185 |
-
|
186 |
-
def extract_vectors_local(net, dataset, device, *, features_num, scales):
|
187 |
-
"""Return tuple (local descriptors, image ids, strenghts, locations and scales) where locations
|
188 |
-
consists of (coor_x, coor_y, scale) and elements of each list correspond to each other"""
|
189 |
-
net.eval()
|
190 |
-
loader = torch.utils.data.DataLoader(dataset, shuffle=False, pin_memory=True, num_workers=NUM_WORKERS)
|
191 |
-
|
192 |
-
with torch.no_grad():
|
193 |
-
vecs, strengths, locs, scls, imids = [], [], [], [], []
|
194 |
-
for imid, inp in io_helpers.progress(enumerate(loader), size=len(loader), print_freq=100):
|
195 |
-
output = net.forward_local(inp.to(device), features_num=features_num, scales=scales)
|
196 |
-
|
197 |
-
vecs.append(output[0].cpu().numpy())
|
198 |
-
strengths.append(output[1].cpu().numpy())
|
199 |
-
locs.append(output[2].cpu().numpy())
|
200 |
-
scls.append(output[3].cpu().numpy())
|
201 |
-
imids.append(np.full((output[0].shape[0],), imid))
|
202 |
-
|
203 |
-
return np.vstack(vecs), np.hstack(imids), np.hstack(strengths), np.vstack(locs), np.hstack(scls)
|
204 |
-
|
205 |
-
|
206 |
-
|
207 |
-
def extract_vectors_all(net, dataset, device, *, features_num, scales):
|
208 |
-
"""Return tuple (local descriptors, image ids, strenghts, locations and scales) where locations
|
209 |
-
consists of (coor_x, coor_y, scale) and elements of each list correspond to each other"""
|
210 |
-
net.eval()
|
211 |
-
loader = torch.utils.data.DataLoader(dataset, shuffle=False, pin_memory=True, num_workers=NUM_WORKERS)
|
212 |
-
|
213 |
-
with torch.no_grad():
|
214 |
-
feats, attns, strenghts = [], [], []
|
215 |
-
for imid, inp in io_helpers.progress(enumerate(loader), size=len(loader), print_freq=100):
|
216 |
-
output = net.get_superfeatures(inp.to(device), scales=scales)
|
217 |
-
feats.append(output[0])
|
218 |
-
attns.append(output[1])
|
219 |
-
strenghts.append(output[2])
|
220 |
-
|
221 |
-
return feats, attns, strenghts
|
|
|
5 |
import torch.nn as nn
|
6 |
import torchvision
|
7 |
|
|
|
|
|
8 |
from .. import layers
|
9 |
from ..layers import functional as HF
|
10 |
from ..utils import io_helpers
|
11 |
|
12 |
NUM_WORKERS = 6
|
13 |
|
|
|
|
|
|
|
|
|
|
|
|
|
14 |
|
15 |
class HOWNet(nn.Module):
|
16 |
"""Network for the HOW method
|
|
|
92 |
|
93 |
return feats, masks
|
94 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
95 |
def __repr__(self):
|
96 |
meta_str = "\n".join(" %s: %s" % x for x in self.meta.items())
|
97 |
return "%s(meta={\n%s\n})" % (self.__class__.__name__, meta_str)
|
|
|
127 |
|
128 |
if skip_layer > 0:
|
129 |
features = features[:-skip_layer]
|
130 |
+
backbone_dim = 2048 // (2 ** skip_layer)
|
131 |
|
132 |
att_layer = layers.attention.L2Attention()
|
133 |
smooth_layer = None
|
|
|
141 |
"architecture": architecture,
|
142 |
"backbone_dim": backbone_dim,
|
143 |
"outputdim": reduction_layer.out_channels if dim_reduction else backbone_dim,
|
144 |
+
"corercf_size": 32 // (2 ** skip_layer),
|
145 |
}
|
146 |
return HOWNet(nn.Sequential(*features), att_layer, smooth_layer, reduction_layer, meta, runtime)
|
147 |
|
148 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|