def test(model,epoch):
torch.cuda.empty_cache()
# switch to evaluate mode
model.eval()
from architectures import AffNetFast
affnet = AffNetFast()
model_weights = 'pretrained/AffNet.pth'
hncheckpoint = torch.load(model_weights)
affnet.load_state_dict(hncheckpoint['state_dict'])
affnet.eval()
detector = ScaleSpaceAffinePatchExtractor( mrSize = 5.192, num_features = 3000,
border = 5, num_Baum_iters = 1,
AffNet = affnet, OriNet = model)
descriptor = HardNet()
model_weights = 'HardNet++.pth'
hncheckpoint = torch.load(model_weights)
descriptor.load_state_dict(hncheckpoint['state_dict'])
descriptor.eval()
if args.cuda:
detector = detector.cuda()
descriptor = descriptor.cuda()
input_img_fname1 = 'test-graf/img1.png'#sys.argv[1]
input_img_fname2 = 'test-graf/img6.png'#sys.argv[1]
H_fname = 'test-graf/H1to6p'#sys.argv[1]
output_img_fname = 'graf_match.png'#sys.argv[3]
img1 = load_grayscale_var(input_img_fname1)
img2 = load_grayscale_var(input_img_fname2)
H = np.loadtxt(H_fname)
H1to2 = Variable(torch.from_numpy(H).float())
SNN_threshold = 0.8
with torch.no_grad():
LAFs1, descriptors1 = get_geometry_and_descriptors(img1, detector, descriptor)
torch.cuda.empty_cache()
LAFs2, descriptors2 = get_geometry_and_descriptors(img2, detector, descriptor)
visualize_LAFs(img1.detach().cpu().numpy().squeeze(), LAFs1.detach().cpu().numpy().squeeze(), 'b', show = False, save_to = LOG_DIR + "/detections1_" + str(epoch) + '.png')
visualize_LAFs(img2.detach().cpu().numpy().squeeze(), LAFs2.detach().cpu().numpy().squeeze(), 'g', show = False, save_to = LOG_DIR + "/detection2_" + str(epoch) + '.png')
dist_matrix = distance_matrix_vector(descriptors1, descriptors2)
min_dist, idxs_in_2 = torch.min(dist_matrix,1)
dist_matrix[:,idxs_in_2] = 100000;# mask out nearest neighbour to find second nearest
min_2nd_dist, idxs_2nd_in_2 = torch.min(dist_matrix,1)
mask = (min_dist / (min_2nd_dist + 1e-8)) <= SNN_threshold
tent_matches_in_1 = indxs_in1 = torch.autograd.Variable(torch.arange(0, idxs_in_2.size(0)), requires_grad = False).cuda()[mask]
tent_matches_in_2 = idxs_in_2[mask]
tent_matches_in_1 = tent_matches_in_1.long()
tent_matches_in_2 = tent_matches_in_2.long()
LAF1s_tent = LAFs1[tent_matches_in_1,:,:]
LAF2s_tent = LAFs2[tent_matches_in_2,:,:]
min_dist, plain_indxs_in1, idxs_in_2 = get_GT_correspondence_indexes(LAF1s_tent, LAF2s_tent,H1to2.cuda(), dist_threshold = 6)
plain_indxs_in1 = plain_indxs_in1.long()
inl_ratio = float(plain_indxs_in1.size(0)) / float(tent_matches_in_1.size(0))
print 'Test epoch', str(epoch)
print 'Test on graf1-6,', tent_matches_in_1.size(0), 'tentatives', plain_indxs_in1.size(0), 'true matches', str(inl_ratio)[:5], ' inl.ratio'
visualize_LAFs(img1.detach().cpu().numpy().squeeze(), LAF1s_tent[plain_indxs_in1.long(),:,:].detach().cpu().numpy().squeeze(), 'g', show = False, save_to = LOG_DIR + "/inliers1_" + str(epoch) + '.png')
visualize_LAFs(img2.detach().cpu().numpy().squeeze(), LAF2s_tent[idxs_in_2.long(),:,:].detach().cpu().numpy().squeeze(), 'g', show = False, save_to = LOG_DIR + "/inliers2_" + str(epoch) + '.png')
return
for epoch in range(start, end):
# iterate over test loaders and test results
train(train_loader, model, optimizer1, epoch)
test(model, epoch)
return 0
if __name__ == '__main__':
LOG_DIR = args.log_dir
LOG_DIR = os.path.join(args.log_dir, suffix)
if not os.path.isdir(LOG_DIR):
os.makedirs(LOG_DIR)
from architectures import AffNetFast, AffNetFastScale, AffNetFast4, AffNetFast4RotNosc, AffNetFast52RotUp, AffNetFast52Rot, AffNetFast5Rot, AffNetFast4Rot, AffNetFast4Rot
from architectures import AffNetFast2Par, AffNetFastBias
if args.arch == 'AffNetFast':
model = AffNetFast(PS=PS)
elif args.arch == 'AffNetFastBias':
model = AffNetFastBias(PS=PS)
elif args.arch == 'AffNetFastScale':
model = AffNetFastScale(PS=PS)
elif args.arch == 'AffNetFast2Par':
model = AffNetFast2Par(PS=PS)
elif args.arch == 'AffNetFast4':
model = AffNetFast4(PS=PS)
elif args.arch == 'AffNetFast4Rot':
model = AffNetFast4Rot(PS=PS)
elif args.arch == 'AffNetFast52':
model = AffNetFast52Rot(PS=PS)
elif args.arch == 'AffNetFast52Rot':
model = AffNetFast52Rot(PS=PS)
elif args.arch == 'AffNetFast52RotUp':
output_fname = sys.argv[2]
nfeats = int(sys.argv[3])
except:
print(
"Wrong input format. Try python hesaffnet.py imgs/cat.png cat.txt 2000"
)
sys.exit(1)
img = Image.open(input_img_fname).convert('RGB')
img = np.mean(np.array(img), axis=2)
var_image = torch.autograd.Variable(torch.from_numpy(img.astype(np.float32)),
volatile=True)
var_image_reshape = var_image.view(1, 1, var_image.size(0), var_image.size(1))
AffNetPix = AffNetFast(PS=32)
weightd_fname = '../../pretrained/AffNet.pth'
checkpoint = torch.load(weightd_fname)
AffNetPix.load_state_dict(checkpoint['state_dict'])
AffNetPix.eval()
HA = ScaleSpaceAffinePatchExtractor(mrSize=5.192,
num_features=nfeats,
border=5,
num_Baum_iters=1,
th=th,
AffNet=AffNetPix)
if USE_CUDA:
HA = HA.cuda()
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torch.backends.cudnn as cudnn
import time
import os
import cv2
import math
import numpy as np
from architectures import AffNetFast
PS = 32
USE_CUDA = False
model = AffNetFast(PS = PS)
weightd_fname = '../../pretrained/AffNet.pth'
checkpoint = torch.load(weightd_fname)
model.load_state_dict(checkpoint['state_dict'])
model.eval()
if USE_CUDA:
model.cuda()
try:
input_img_fname = sys.argv[1]
output_fname = sys.argv[2]
except:
print "Wrong input format. Try ./detect_affine_shape.py imgs/ref.png out.txt"
sys.exit(1)