def getAffmaps_from_Affnet(patches_np):
sp1, sp2 = np.shape(patches_np[0])
subpatches = torch.autograd.Variable(
torch.zeros([len(patches_np), 1, 32, 32], dtype=torch.float32),
volatile=True).view(len(patches_np), 1, 32, 32)
for k in range(0, len(patches_np)):
subpatch = patches_np[k][int(sp1 / 2) - 16:int(sp2 / 2) + 16,
int(sp1 / 2) - 16:int(sp2 / 2) + 16].reshape(
1, 1, 32, 32)
subpatches[k, :, :, :] = torch.from_numpy(subpatch.astype(
np.float32)) #=subpatch
x, y = subpatches.shape[3] / 2.0 + 2, subpatches.shape[2] / 2.0 + 2
LAFs = normalizeLAFs(
torch.tensor([[AffNetPix.PS / 2, 0, x], [0, AffNetPix.PS / 2,
y]]).reshape(1, 2, 3),
subpatches.shape[3], subpatches.shape[2])
baseLAFs = torch.zeros([subpatches.shape[0], 2, 3], dtype=torch.float32)
for m in range(subpatches.shape[0]):
baseLAFs[m, :, :] = LAFs
if USE_CUDA:
# or ---> A = AffNetPix(subpatches.cuda()).cpu()
with torch.no_grad():
A = batched_forward(AffNetPix, subpatches.cuda(), 256).cpu()
else:
with torch.no_grad():
A = AffNetPix(subpatches)
LAFs = torch.cat([torch.bmm(A, baseLAFs[:, :, 0:2]), baseLAFs[:, :, 2:]],
dim=2)
dLAFs = denormalizeLAFs(LAFs, subpatches.shape[3], subpatches.shape[2])
Alist = convertLAFs_to_A23format(dLAFs.detach().cpu().numpy().astype(
np.float32))
return Alist
def extract_patches_from_pyr(self, dLAFs, PS = 41):
pyr_idxs, level_idxs = get_pyramid_and_level_index_for_LAFs(dLAFs, self.sigmas, self.pix_dists, PS)
pyr_inv_idxs = get_inverted_pyr_index(self.scale_pyr, pyr_idxs, level_idxs)
patches = extract_patches_from_pyramid_with_inv_index(self.scale_pyr,
pyr_inv_idxs,
normalizeLAFs(dLAFs, self.scale_pyr[0][0].size(3), self.scale_pyr[0][0].size(2)),
PS = PS)
return patches
def AffNetHardNet_describeFromKeys(img_np, KPlist):
img = torch.autograd.Variable(torch.from_numpy(img_np.astype(np.float32)),
volatile=True)
img = img.view(1, 1, img.size(0), img.size(1))
HessianAffine = ScaleSpaceAffinePatchExtractor(mrSize=5.192,
num_features=0,
border=0,
num_Baum_iters=0)
if USE_CUDA:
HessianAffine = HessianAffine.cuda()
img = img.cuda()
with torch.no_grad():
HessianAffine.createScaleSpace(
img) # to generate scale pyramids and stuff
descriptors = []
Alist = []
n = 0
# for patch_np in patches:
for kp in KPlist:
x, y = np.float32(kp.pt)
LAFs = normalizeLAFs(
torch.tensor([[AffNetPix.PS / 2, 0, x], [0, AffNetPix.PS / 2,
y]]).reshape(1, 2, 3),
img.size(3), img.size(2))
with torch.no_grad():
patch = HessianAffine.extract_patches_from_pyr(denormalizeLAFs(
LAFs, img.size(3), img.size(2)),
PS=AffNetPix.PS)
if WRITE_IMGS_DEBUG:
SaveImageWithKeys(patch.detach().cpu().numpy().reshape([32, 32]),
[], 'p2/' + str(n) + '.png')
if USE_CUDA:
# or ---> A = AffNetPix(subpatches.cuda()).cpu()
with torch.no_grad():
A = batched_forward(AffNetPix, patch.cuda(), 256).cpu()
else:
with torch.no_grad():
A = AffNetPix(patch)
new_LAFs = torch.cat([torch.bmm(A, LAFs[:, :, 0:2]), LAFs[:, :, 2:]],
dim=2)
dLAFs = denormalizeLAFs(new_LAFs, img.size(3), img.size(2))
with torch.no_grad():
patchaff = HessianAffine.extract_patches_from_pyr(dLAFs, PS=32)
if WRITE_IMGS_DEBUG:
SaveImageWithKeys(
patchaff.detach().cpu().numpy().reshape([32, 32]), [],
'p1/' + str(n) + '.png')
SaveImageWithKeys(img_np, [kp], 'im1/' + str(n) + '.png')
descriptors.append(
HardNetDescriptor(patchaff).cpu().numpy().astype(np.float32))
Alist.append(
convertLAFs_to_A23format(LAFs.detach().cpu().numpy().astype(
np.float32)))
n = n + 1
return descriptors, Alist
def train(train_loader, model, optimizer, epoch, cuda=True):
# switch to train mode
model.train()
log_interval = 1
total_loss = 0
total_feats = 0
spatial_only = True
pbar = enumerate(train_loader)
for batch_idx, data in pbar:
print 'Batch idx', batch_idx
#print model.detector.shift_net[0].weight.data.cpu().numpy()
img1, img2, H1to2 = data
#if np.abs(np.sum(H.numpy()) - 3.0) > 0.01:
# continue
H1to2 = H1to2.squeeze(0)
if (img1.size(3) * img1.size(4) > 1340 * 1000):
print img1.shape, ' too big, skipping'
continue
img1 = img1.float().squeeze(0)
img2 = img2.float().squeeze(0)
if cuda:
img1, img2, H1to2 = img1.cuda(), img2.cuda(), H1to2.cuda()
img1, img2, H1to2 = Variable(img1, requires_grad=False), Variable(
img2, requires_grad=False), Variable(H1to2, requires_grad=False)
LAFs1, aff_norm_patches1, resp1 = HA(img1, True, True, True)
LAFs2, aff_norm_patches2, resp2 = HA(img2, True, True)
if (len(LAFs1) == 0) or (len(LAFs2) == 0):
optimizer.zero_grad()
continue
fro_dists, idxs_in1, idxs_in2, LAFs2_in_1 = get_GT_correspondence_indexes_Fro_and_center(
LAFs1,
LAFs2,
H1to2,
dist_threshold=4.,
center_dist_th=7.0,
skip_center_in_Fro=True,
do_up_is_up=True,
return_LAF2_in_1=True)
if len(fro_dists.size()) == 0:
optimizer.zero_grad()
print 'skip'
continue
aff_patches_from_LAFs2_in_1 = extract_patches(
img1,
normalizeLAFs(LAFs2_in_1[idxs_in2.data.long(), :, :], img1.size(3),
img1.size(2)))
#loss = fro_dists.mean()
patch_dist = torch.sqrt(
(aff_norm_patches1[idxs_in1.data.long(), :, :, :] / 100. -
aff_patches_from_LAFs2_in_1 / 100.)**2 + 1e-8).view(
fro_dists.size(0), -1).mean(dim=1)
loss = (fro_dists * patch_dist).mean()
print 'Fro dist', fro_dists.mean().data.cpu().numpy(
)[0], loss.data.cpu().numpy()[0]
total_loss += loss.data.cpu().numpy()[0]
#loss += patch_dist
total_feats += fro_dists.size(0)
optimizer.zero_grad()
loss.backward()
optimizer.step()
#adjust_learning_rate(optimizer)
print epoch, batch_idx, loss.data.cpu().numpy()[0], idxs_in1.shape
print 'Train total loss:', total_loss / float(
batch_idx + 1), ' features ', float(total_feats) / float(batch_idx + 1)
torch.save({
'epoch': epoch + 1,
'state_dict': model.state_dict()
}, '{}/elu_new_checkpoint_{}.pth'.format(LOG_DIR, epoch))
def AffNetHardNet_describe(patches):
descriptors = np.zeros(shape=[patches.shape[0], 128], dtype=np.float32)
HessianAffine = []
subpatches = torch.autograd.Variable(torch.zeros([len(patches), 1, 32, 32],
dtype=torch.float32),
volatile=True).view(
len(patches), 1, 32, 32)
baseLAFs = torch.zeros([len(patches), 2, 3], dtype=torch.float32)
for m in range(patches.shape[0]):
patch_np = patches[m, :, :, 0].reshape(np.shape(patches)[1:3])
HessianAffine.append(
ScaleSpaceAffinePatchExtractor(mrSize=5.192,
num_features=0,
border=0,
num_Baum_iters=0))
with torch.no_grad():
var_image = torch.autograd.Variable(torch.from_numpy(
patch_np.astype(np.float32)),
volatile=True)
patch = var_image.view(1, 1, var_image.size(0), var_image.size(1))
with torch.no_grad():
HessianAffine[m].createScaleSpace(
patch) # to generate scale pyramids and stuff
x, y = patch.size(3) / 2.0 + 2, patch.size(2) / 2.0 + 2
LAFs = normalizeLAFs(
torch.tensor([[AffNetPix.PS / 2, 0, x], [0, AffNetPix.PS / 2,
y]]).reshape(1, 2, 3),
patch.size(3), patch.size(2))
baseLAFs[m, :, :] = LAFs
with torch.no_grad():
subpatch = HessianAffine[m].extract_patches_from_pyr(
denormalizeLAFs(LAFs, patch.size(3), patch.size(2)),
PS=AffNetPix.PS)
if WRITE_IMGS_DEBUG:
SaveImageWithKeys(
subpatch.detach().cpu().numpy().reshape([32, 32]), [],
'p1/' + str(n) + '.png')
# This subpatch has been blured by extract_patches _from_pyr...
# let't us crop it manually to obtain fair results agains other methods
subpatch = patch_np[16:48, 16:48].reshape(1, 1, 32, 32)
#var_image = torch.autograd.Variable(torch.from_numpy(subpatch.astype(np.float32)), volatile = True)
#subpatch = var_image.view(1, 1, 32,32)
subpatches[m, :, :, :] = torch.from_numpy(
subpatch.astype(np.float32)) #=subpatch
if WRITE_IMGS_DEBUG:
SaveImageWithKeys(
subpatch.detach().cpu().numpy().reshape([32, 32]), [],
'p2/' + str(n) + '.png')
if USE_CUDA:
# or ---> A = AffNetPix(subpatches.cuda()).cpu()
with torch.no_grad():
A = batched_forward(AffNetPix, subpatches.cuda(), 256).cpu()
else:
with torch.no_grad():
A = AffNetPix(subpatches)
LAFs = torch.cat([torch.bmm(A, baseLAFs[:, :, 0:2]), baseLAFs[:, :, 2:]],
dim=2)
dLAFs = denormalizeLAFs(LAFs, patch.size(3), patch.size(2))
Alist = convertLAFs_to_A23format(dLAFs.detach().cpu().numpy().astype(
np.float32))
for m in range(patches.shape[0]):
with torch.no_grad():
patchaff = HessianAffine[m].extract_patches_from_pyr(
dLAFs[m, :, :].reshape(1, 2, 3), PS=32)
if WRITE_IMGS_DEBUG:
SaveImageWithKeys(
patchaff.detach().cpu().numpy().reshape([32, 32]), [],
'im1/' + str(n) + '.png')
SaveImageWithKeys(patch_np, [], 'im2/' + str(n) + '.png')
subpatches[m, :, :, :] = patchaff
if USE_CUDA:
with torch.no_grad():
# descriptors = HardNetDescriptor(subpatches.cuda()).detach().cpu().numpy().astype(np.float32)
descriptors = batched_forward(HardNetDescriptor, subpatches.cuda(),
256).cpu().numpy().astype(np.float32)
else:
with torch.no_grad():
descriptors = HardNetDescriptor(
subpatches).detach().cpu().numpy().astype(np.float32)
return descriptors, Alist
def train(train_loader, model, optimizer, epoch, cuda = True):
# switch to train mode
model.train()
log_interval = 1
total_loss = 0
total_feats = 0
spatial_only = True
pbar = enumerate(train_loader)
for batch_idx, data in pbar:
#if batch_idx > 0:
# continue
print 'Batch idx', batch_idx
#print model.detector.shift_net[0].weight.data.cpu().numpy()
img1, img2, H1to2 = data
#if np.abs(np.sum(H.numpy()) - 3.0) > 0.01:
# continue
H1to2 = H1to2.squeeze(0)
do_aug = True
if torch.abs(H1to2 - torch.eye(3)).sum() > 0.05:
do_aug = False
if (img1.size(3) *img1.size(4) > 1340*1000):
print img1.shape, ' too big, skipping'
continue
img1 = img1.float().squeeze(0)
img2 = img2.float().squeeze(0)
if cuda:
img1, img2, H1to2 = img1.cuda(), img2.cuda(), H1to2.cuda()
img1, img2, H1to2 = Variable(img1, requires_grad = False), Variable(img2, requires_grad = False), Variable(H1to2, requires_grad = False)
if do_aug:
new_img2, H_Orig2New = affineAug(img2, max_add = 0.2 )
H1to2new = torch.mm(H_Orig2New, H1to2)
else:
new_img2 = img2
H1to2new = H1to2
#print H1to2
LAFs1, aff_norm_patches1, resp1, dets1, A1 = HA(img1, True, False, True)
LAFs2Aug, aff_norm_patches2, resp2, dets2, A2 = HA(new_img2, True, False)
if (len(LAFs1) == 0) or (len(LAFs2Aug) == 0):
optimizer.zero_grad()
continue
geom_loss, idxs_in1, idxs_in2, LAFs2_in_1 = LAFMagic(LAFs1,
LAFs2Aug,
H1to2new,
3.0, scale_log = 0.3)
if len(idxs_in1.size()) == 0:
optimizer.zero_grad()
print 'skip'
continue
aff_patches_from_LAFs2_in_1 = extract_patches(img1,
normalizeLAFs(LAFs2_in_1[idxs_in2.long(),:,:],
img1.size(3), img1.size(2)))
SIFTs1 = SIFT(aff_norm_patches1[idxs_in1.long(),:,:,:]).cuda()
SIFTs2 = SIFT(aff_patches_from_LAFs2_in_1).cuda()
#sift_snn_loss = loss_HardNet(SIFTs1, SIFTs2, column_row_swap = True,
# margin = 1.0, batch_reduce = 'min', loss_type = "triplet_margin");
patch_dist = 2.0 * torch.sqrt((aff_norm_patches1[idxs_in1.long(),:,:,:]/100. - aff_patches_from_LAFs2_in_1/100.) **2 + 1e-8).view(idxs_in1.size(0),-1).mean(dim = 1)
sift_dist = torch.sqrt(((SIFTs1 - SIFTs2)**2 + 1e-12).mean(dim=1))
loss = geom_loss.cuda() .mean()
total_loss += loss.data.cpu().numpy()[0]
#loss += patch_dist
total_feats += aff_patches_from_LAFs2_in_1.size(0)
optimizer.zero_grad()
loss.backward()
optimizer.step()
adjust_learning_rate(optimizer)
if batch_idx % 10 == 0:
print 'A', A1.data.cpu().numpy()[0:1,:,:]
print 'crafted loss', pr_l(geom_loss), 'patch', pr_l(patch_dist), 'sift', pr_l(sift_dist)#, 'hardnet', pr_l(sift_snn_loss)
print epoch,batch_idx, loss.data.cpu().numpy()[0], idxs_in1.shape
print 'Train total loss:', total_loss / float(batch_idx+1), ' features ', float(total_feats) / float(batch_idx+1)
torch.save({'epoch': epoch + 1, 'state_dict': model.state_dict()},
'{}/new_loss_sep_checkpoint_{}.pth'.format(LOG_DIR, epoch))
