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))