-
Notifications
You must be signed in to change notification settings - Fork 37
/
eval_inloc.py
226 lines (185 loc) · 9.28 KB
/
eval_inloc.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
from __future__ import print_function, division
import os
from os.path import exists, join, basename
from collections import OrderedDict
import numpy as np
import numpy.random
import scipy as sc
import scipy.misc
from scipy.io import loadmat
from scipy.io import savemat
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torch.autograd import Variable
from lib.model import ImMatchNet,MutualMatching
from lib.normalization import NormalizeImageDict
from lib.torch_util import str_to_bool
from lib.point_tnf import normalize_axis,unnormalize_axis,corr_to_matches
from lib.plot import plot_image
import argparse
print('NCNet evaluation script - InLoc dataset')
use_cuda = torch.cuda.is_available()
# Argument parsing
parser = argparse.ArgumentParser(description='Compute InLoc matches')
parser.add_argument('--checkpoint', type=str, default='')
parser.add_argument('--inloc_shortlist', type=str, default='datasets/inloc/densePE_top100_shortlist_cvpr18.mat')
parser.add_argument('--k_size', type=int, default=2)
parser.add_argument('--image_size', type=int, default=3200)
parser.add_argument('--n_queries', type=int, default=356)
parser.add_argument('--n_panos', type=int, default=10)
parser.add_argument('--softmax', type=str_to_bool, default=True)
parser.add_argument('--matching_both_directions', type=str_to_bool, default=True)
parser.add_argument('--flip_matching_direction', type=str_to_bool, default=False)
parser.add_argument('--pano_path', type=str, default='datasets/inloc/pano/', help='path to InLoc panos - should contain CSE3,CSE4,CSE5,DUC1 and DUC2 folders')
parser.add_argument('--query_path', type=str, default='datasets/inloc/query/iphone7/', help='path to InLoc queries')
args = parser.parse_args()
image_size = args.image_size
k_size = args.k_size
matching_both_directions = args.matching_both_directions
flip_matching_direction = args.flip_matching_direction
# Load pretrained model
half_precision=True # use for memory saving
print(args)
model = ImMatchNet(use_cuda=use_cuda,
checkpoint=args.checkpoint,
half_precision=half_precision,
relocalization_k_size=args.k_size)
# Generate output folder path
output_folder = args.inloc_shortlist.split('/')[-1].split('.')[0]+'_SZ_NEW_'+str(image_size)+'_K_'+str(k_size)
if matching_both_directions:
output_folder += '_BOTHDIRS'
elif flip_matching_direction:
output_folder += '_AtoB'
else:
output_folder += '_BtoA'
if args.softmax==True:
output_folder += '_SOFTMAX'
if args.checkpoint!='':
checkpoint_name=args.checkpoint.split('/')[-1].split('.')[0]
output_folder += '_CHECKPOINT_'+checkpoint_name
print('Output matches folder: '+output_folder)
# Data preprocessing
# Manually change image resolution for this test. On training, image_size=400 was used, with squared images
scale_factor = 0.0625
imreadth = lambda x: torch.Tensor(sc.misc.imread(x).astype(np.float32)).transpose(1,2).transpose(0,1)
normalize = lambda x: NormalizeImageDict(['im'])({'im':x})['im']
# allow rectangular images. Does not modify aspect ratio.
if k_size==1:
resize = lambda x: nn.functional.upsample(Variable(x.unsqueeze(0).cuda(),volatile=True),
size=(int(x.shape[1]/(np.max(x.shape[1:])/image_size)),int(x.shape[2]/(np.max(x.shape[1:])/image_size))),mode='bilinear')
else:
resize = lambda x: nn.functional.upsample(Variable(x.unsqueeze(0).cuda(),volatile=True),
size=(int(np.floor(x.shape[1]/(np.max(x.shape[1:])/image_size)*scale_factor/k_size)/scale_factor*k_size),
int(np.floor(x.shape[2]/(np.max(x.shape[1:])/image_size)*scale_factor/k_size)/scale_factor*k_size)),mode='bilinear')
padim = lambda x,h_max: torch.cat((x,x.view(-1)[0].clone().expand(1,3,h_max-x.shape[2],x.shape[3])/1e20),dim=2) if x.shape[2]<h_max else x
# Get shortlists for each query image
shortlist_fn = args.inloc_shortlist
dbmat = loadmat(shortlist_fn)
db = dbmat['ImgList'][0,:]
query_fn_all=np.squeeze(np.vstack(tuple([db[q][0] for q in range(len(db))])))
pano_fn_all=np.vstack(tuple([db[q][1] for q in range(len(db))]))
Nqueries=args.n_queries
Npanos=args.n_panos
try:
os.mkdir('matches/')
except FileExistsError:
pass
try:
os.mkdir('matches/'+output_folder)
except FileExistsError:
pass
N=int((image_size*scale_factor/k_size)*np.floor((image_size*scale_factor/k_size)*(3/4)))
if matching_both_directions:
N=2*N
do_softmax = args.softmax
plot=False
for q in range(Nqueries):
print(q)
matches=numpy.zeros((1,Npanos,N,5))
# load query image
query_fn = os.path.join(args.query_path,db[q][0].item())
src=resize(normalize(imreadth(query_fn)))
# load database image
for idx in range(Npanos):
pano_fn = os.path.join(args.pano_path,db[q][1].ravel()[idx].item())
tgt=resize(normalize(imreadth(pano_fn)))
if k_size>1:
corr4d,delta4d=model({'source_image':src,'target_image':tgt})
else:
corr4d=model({'source_image':src,'target_image':tgt})
delta4d=None
# reshape corr tensor and get matches for each point in image B
batch_size,ch,fs1,fs2,fs3,fs4 = corr4d.size()
# pad image and plot
if plot:
h_max=int(np.max([src.shape[2],tgt.shape[2]]))
im=plot_image(torch.cat((padim(src,h_max),padim(tgt,h_max)),dim=3),return_im=True)
plt.imshow(im)
if matching_both_directions:
(xA_,yA_,xB_,yB_,score_)=corr_to_matches(corr4d,scale='positive',do_softmax=do_softmax,delta4d=delta4d,k_size=k_size)
(xA2_,yA2_,xB2_,yB2_,score2_)=corr_to_matches(corr4d,scale='positive',do_softmax=do_softmax,delta4d=delta4d,k_size=k_size,invert_matching_direction=True)
xA_=torch.cat((xA_,xA2_),1)
yA_=torch.cat((yA_,yA2_),1)
xB_=torch.cat((xB_,xB2_),1)
yB_=torch.cat((yB_,yB2_),1)
score_=torch.cat((score_,score2_),1)
# sort in descending score (this will keep the max-score instance in the duplicate removal step)
sorted_index=torch.sort(-score_)[1].squeeze()
xA_=xA_.squeeze()[sorted_index].unsqueeze(0)
yA_=yA_.squeeze()[sorted_index].unsqueeze(0)
xB_=xB_.squeeze()[sorted_index].unsqueeze(0)
yB_=yB_.squeeze()[sorted_index].unsqueeze(0)
score_=score_.squeeze()[sorted_index].unsqueeze(0)
# remove duplicates
concat_coords=np.concatenate((xA_.cpu().data.numpy(),yA_.cpu().data.numpy(),xB_.cpu().data.numpy(),yB_.cpu().data.numpy()),0)
_,unique_index=np.unique(concat_coords,axis=1,return_index=True)
xA_=xA_.squeeze()[torch.cuda.LongTensor(unique_index)].unsqueeze(0)
yA_=yA_.squeeze()[torch.cuda.LongTensor(unique_index)].unsqueeze(0)
xB_=xB_.squeeze()[torch.cuda.LongTensor(unique_index)].unsqueeze(0)
yB_=yB_.squeeze()[torch.cuda.LongTensor(unique_index)].unsqueeze(0)
score_=score_.squeeze()[torch.cuda.LongTensor(unique_index)].unsqueeze(0)
elif flip_matching_direction:
(xA_,yA_,xB_,yB_,score_)=corr_to_matches(corr4d,scale='positive',do_softmax=do_softmax,delta4d=delta4d,k_size=k_size,invert_matching_direction=True)
else:
(xA_,yA_,xB_,yB_,score_)=corr_to_matches(corr4d,scale='positive',do_softmax=do_softmax,delta4d=delta4d,k_size=k_size)
# recenter
if k_size>1:
yA_=yA_*(fs1*k_size-1)/(fs1*k_size)+0.5/(fs1*k_size)
xA_=xA_*(fs2*k_size-1)/(fs2*k_size)+0.5/(fs2*k_size)
yB_=yB_*(fs3*k_size-1)/(fs3*k_size)+0.5/(fs3*k_size)
xB_=xB_*(fs4*k_size-1)/(fs4*k_size)+0.5/(fs4*k_size)
else:
yA_=yA_*(fs1-1)/fs1+0.5/fs1
xA_=xA_*(fs2-1)/fs2+0.5/fs2
yB_=yB_*(fs3-1)/fs3+0.5/fs3
xB_=xB_*(fs4-1)/fs4+0.5/fs4
xA = xA_.view(-1).data.cpu().float().numpy()
yA = yA_.view(-1).data.cpu().float().numpy()
xB = xB_.view(-1).data.cpu().float().numpy()
yB = yB_.view(-1).data.cpu().float().numpy()
score = score_.view(-1).data.cpu().float().numpy()
Npts=len(xA)
if Npts>0:
matches[0,idx,:Npts,0]=xA
matches[0,idx,:Npts,1]=yA
matches[0,idx,:Npts,2]=xB
matches[0,idx,:Npts,3]=yB
matches[0,idx,:Npts,4]=score
# plot top N matches
if plot:
c=numpy.random.rand(Npts,3)
for i in range(Npts):
if score[i]>0.75:
ax = plt.gca()
ax.add_artist(plt.Circle((float(xA[i])*src.shape[3],float(yA[i])*src.shape[2]), radius=3, color=c[i,:]))
ax.add_artist(plt.Circle((float(xB[i])*tgt.shape[3]+src.shape[3] ,float(yB[i])*tgt.shape[2]), radius=3, color=c[i,:]))
#plt.plot([float(xA[i])*src.shape[3], float(xB[i])*tgt.shape[3]+src.shape[3]], [float(yA[i])*src.shape[2], float(yB[i])*tgt.shape[2]], c='g', linestyle='-', linewidth=0.1)
corr4d=None
delta4d=None
if idx%10==0:
print(">>>"+str(idx))
savemat(os.path.join('matches/',output_folder,str(q+1)+'.mat'),{'matches':matches,'query_fn':db[q][0].item(),'pano_fn':pano_fn_all},do_compression=True)
if plot:
plt.gcf().set_dpi(200)
plt.show()