text: 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()
  
  