YannisK commited on
Commit
4e8ced7
1 Parent(s): c9dadbf
Files changed (3) hide show
  1. app.py +1 -3
  2. fire_network.py +3 -15
  3. 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, scale=6):
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, CORERCF_SIZE
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
- def forward_global(self, x, *, scales):
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": CORERCF_SIZE[architecture] // (2 ** skip_layer),
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 = imageretrievalnet.OUTPUT_DIM[architecture] // (2 ** skip_layer)
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": CORERCF_SIZE[architecture] // (2 ** skip_layer),
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