def getFlow_Coarse(pairID, flowList, finePath, coarsePath) : find = False for flowName in flowList : if flowName.split('_')[1] == str(pairID) : nbH = flowName.split('_')[2].split('H')[0] find = True break if not find : return [], [] flow = torch.from_numpy ( np.load(os.path.join(finePath, 'flow_{:d}_{}H.npy'.format(pairID, nbH))).astype(np.float32) ) param = torch.from_numpy ( np.load(os.path.join(coarsePath, 'flow_{:d}_{}H.npy'.format(pairID, nbH))).astype(np.float32) ) h, w = flow.size()[2], flow.size()[3] #### -- grid gridY = torch.linspace(-1, 1, steps = h * 8).view(1, -1, 1, 1).expand(1, h * 8, w * 8, 1) gridX = torch.linspace(-1, 1, steps = w * 8).view(1, 1, -1, 1).expand(1, h * 8, w * 8, 1) grid = torch.cat((gridX, gridY), dim=3) warper = tgm.HomographyWarper(h * 8, w * 8) coarse = warper.warp_grid(param.narrow(0, 0, 1)) return coarse, torch.ones(1, h * 8, w * 8, 1)
def getFlow(pairID, finePath, flowList, coarsePath, maskPath, multiH, th) : find = False for flowName in flowList : if flowName.split('_')[1] == str(pairID) : nbH = flowName.split('_')[2].split('H')[0] find = True break if not find : return [], [] flow = torch.from_numpy ( np.load(os.path.join(finePath, 'flow_{:d}_{}H.npy'.format(pairID, nbH))).astype(np.float32) ) param = torch.from_numpy ( np.load(os.path.join(coarsePath, 'flow_{:d}_{}H.npy'.format(pairID, nbH))).astype(np.float32) ) match = np.load(os.path.join(finePath, 'mask_{:d}_{}H.npy'.format(pairID, nbH))) matchBG = np.load(os.path.join(maskPath, 'maskBG_{:d}_{}H.npy'.format(pairID, nbH))) h, w = flow.size()[2], flow.size()[3] #### -- grid gridY = torch.linspace(-1, 1, steps = h * 8).view(1, -1, 1, 1).expand(1, h * 8, w * 8, 1) gridX = torch.linspace(-1, 1, steps = w * 8).view(1, 1, -1, 1).expand(1, h * 8, w * 8, 1) grid = torch.cat((gridX, gridY), dim=3) warper = tgm.HomographyWarper(h * 8, w * 8) coarse = warper.warp_grid(param) flow = F.interpolate(input = flow, scale_factor = 8, mode='bilinear') flow = flow.permute(0, 2, 3, 1) flowUp = torch.clamp(flow + grid, min=-1, max=1) flow = F.grid_sample(coarse.permute(0, 3, 1, 2), flowUp).permute(0, 2, 3, 1).contiguous() match = torch.from_numpy(match) match = F.interpolate(input = match, scale_factor = 8, mode='bilinear') match = match.narrow(1, 0, 1) * F.grid_sample(match.narrow(1, 1, 1), flowUp) * (((flow.narrow(3, 0, 1) >= -1) * ( flow.narrow(3, 0, 1) <= 1)).type(torch.FloatTensor) * ((flow.narrow(3, 1, 1) >= -1) * ( flow.narrow(3, 1, 1) <= 1)).type(torch.FloatTensor)).permute(0, 3, 1, 2) #match = match.narrow(1, 0, 1) * (((flow.narrow(3, 0, 1) >= -1) * ( flow.narrow(3, 0, 1) <= 1)).type(torch.FloatTensor) * ((flow.narrow(3, 1, 1) >= -1) * ( flow.narrow(3, 1, 1) <= 1)).type(torch.FloatTensor)).permute(0, 3, 1, 2) match = match.permute(0, 2, 3, 1) flow = torch.clamp(flow, min=-1, max=1) flowGlobal = flow[:1] match_binary = match[:1] >= th matchGlobal = match[:1] if multiH : for i in range(1, len(match)) : tmp_match = (match.narrow(0, i, 1) >= th) * (~ match_binary) matchGlobal[tmp_match] = match.narrow(0, i, 1)[tmp_match] match_binary = match_binary + tmp_match tmp_match = tmp_match.expand_as(flowGlobal) flowGlobal[tmp_match] = flow.narrow(0, i, 1)[tmp_match] return flowGlobal, matchGlobal
def get_info(I): w, h = I.size gridY = torch.linspace(-1, 1, steps=h).view(1, -1, 1, 1).expand(1, h, w, 1) gridX = torch.linspace(-1, 1, steps=w).view(1, 1, -1, 1).expand(1, h, w, 1) grid = torch.cat((gridX, gridY), dim=3).cuda() tensor = transforms.ToTensor()(I).unsqueeze(0).cuda() warper = tgm.HomographyWarper(h, w) return w, h, tensor, grid, warper
def _getFlow(flow, param, match, matchBG, multiH, th): h, w = flow.size()[2], flow.size()[3] #### -- grid gridY = torch.linspace(-1, 1, steps=h * 8).view(1, -1, 1, 1).expand(1, h * 8, w * 8, 1) gridX = torch.linspace(-1, 1, steps=w * 8).view(1, 1, -1, 1).expand(1, h * 8, w * 8, 1) grid = torch.cat((gridX, gridY), dim=3) warper = tgm.HomographyWarper(h * 8, w * 8) coarse = warper.warp_grid(param) flow = F.interpolate(input=flow, scale_factor=8, mode='bilinear') flow = flow.permute(0, 2, 3, 1) flowUp = torch.clamp(flow + grid, min=-1, max=1) flow = F.grid_sample(coarse.permute(0, 3, 1, 2), flowUp).permute(0, 2, 3, 1).contiguous() match = F.interpolate(input=match, scale_factor=8, mode='bilinear') match = match.narrow(1, 0, 1) * F.grid_sample(match.narrow( 1, 1, 1), flowUp) * ( ((flow.narrow(3, 0, 1) >= -1) * (flow.narrow(3, 0, 1) <= 1)).type(torch.FloatTensor) * ((flow.narrow(3, 1, 1) >= -1) * (flow.narrow(3, 1, 1) <= 1)).type(torch.FloatTensor)).permute( 0, 3, 1, 2) #match = match.narrow(1, 0, 1) * (((flow.narrow(3, 0, 1) >= -1) * ( flow.narrow(3, 0, 1) <= 1)).type(torch.FloatTensor) * ((flow.narrow(3, 1, 1) >= -1) * ( flow.narrow(3, 1, 1) <= 1)).type(torch.FloatTensor)).permute(0, 3, 1, 2) match = match.permute(0, 2, 3, 1) flow = torch.clamp(flow, min=-1, max=1) flowGlobal = flow[:1] match_binary = match[:1] >= th if multiH: for i in range(1, len(match)): tmp_match = (match.narrow(0, i, 1) >= th) * (~match_binary) match_binary = match_binary + tmp_match tmp_match = tmp_match.expand_as(flowGlobal) flowGlobal[tmp_match] = flow.narrow(0, i, 1)[tmp_match] flowGlobal, match_binary = flowGlobal.squeeze().numpy( ), match_binary.squeeze().numpy() * matchBG return flowGlobal, match_binary
def calculate_epe_hpatches(net, val_loader, device, k, inPklCoarse, onlyCoarse, transformation, Transform, coarsePlus=None, iterativeRefine=False, img_size=240): """ Compute EPE for HPatches dataset Args: net: trained model val_loader: input dataloader device: `cpu` or `gpu` img_size: size of input images Output: aepe_array: averaged EPE for the whole sequence of HPatches """ aepe_array = [] n_registered_pxs = 0 pbar = tqdm(enumerate(val_loader), total=len(val_loader)) for i, mini_batch in pbar: source_img = mini_batch['source_image'].to(device) target_img = mini_batch['target_image'].to(device) bs, _, _, _ = source_img.shape #### # net prediction gridY = torch.linspace(-1, 1, steps=target_img.size(2)).view( 1, -1, 1, 1).expand(1, target_img.size(2), target_img.size(3), 1) gridX = torch.linspace(-1, 1, steps=target_img.size(3)).view( 1, 1, -1, 1).expand(1, target_img.size(2), target_img.size(3), 1) gridFine = torch.cat((gridX, gridY), dim=3).cuda() bestParam = inPklCoarse[i] if transformation == 'Affine': grid = F.affine_grid( torch.from_numpy( bestParam[:2].astype(np.float32)).unsqueeze(0).cuda(), target_img.size()) # theta should be of size N×2×3 else: bestParam = bestParam.astype(np.float32) warper = tgm.HomographyWarper(target_img.size()[2], target_img.size()[3]) grid = warper.warp_grid( torch.from_numpy(bestParam).unsqueeze(0).cuda()) IsSample = F.grid_sample(source_img, grid) sourceFeat = F.normalize(net['netFeatCoarse'](IsSample)) targetFeat = F.normalize(net['netFeatCoarse'](target_img)) corr12 = net['netCorr'](targetFeat, sourceFeat) match12 = model.predMatchability(corr12, net['netMatch']) corr21 = net['netCorr'](sourceFeat, targetFeat) match21 = model.predMatchability(corr21, net['netMatch']) _, flow_est = model.predFlowCoarse(corr12, net['netFlowCoarse'], gridFine) match_est = match12 * F.grid_sample(match21, flow_est) _, flow_est_inverse = model.predFlowCoarse(corr21, net['netFlowCoarse'], gridFine) flow_est = F.grid_sample(grid.permute(0, 3, 1, 2), flow_est).permute(0, 2, 3, 1).contiguous() flow_est = grid if onlyCoarse else flow_est if iterativeRefine: flow_est = iterative(net, gridFine, source_img, target_img, flow_est, i, match_est, Transform, coarsePlus, transformation) flow_target = mini_batch['correspondence_map'].to(device) mask_x_gt = \ flow_target[:, :, :, 0].ge(-1) & flow_target[:, :, :, 0].le(1) mask_y_gt = \ flow_target[:, :, :, 1].ge(-1) & flow_target[:, :, :, 1].le(1) mask_xx_gt = mask_x_gt & mask_y_gt mask_gt = torch.cat((mask_xx_gt.unsqueeze(3), mask_xx_gt.unsqueeze(3)), dim=3) for i in range(bs): # unnormalize the flow: [-1; 1] -> [0; im_size - 1] flow_target[i] = (flow_target[i] + 1) * (img_size - 1) / (1 + 1) flow_est[i] = (flow_est[i] + 1) * (img_size - 1) / (1 + 1) flow_target_x = flow_target[:, :, :, 0] flow_target_y = flow_target[:, :, :, 1] flow_est_x = flow_est[:, :, :, 0] flow_est_y = flow_est[:, :, :, 1] flow_target = \ torch.cat((flow_target_x[mask_gt[:, :, :, 0]].unsqueeze(1), flow_target_y[mask_gt[:, :, :, 1]].unsqueeze(1)), dim=1) flow_est = \ torch.cat((flow_est_x[mask_gt[:, :, :, 0]].unsqueeze(1), flow_est_y[mask_gt[:, :, :, 1]].unsqueeze(1)), dim=1) # let's calculate EPE aepe = epe(flow_est, flow_target) aepe_array.append(aepe.item()) n_registered_pxs += flow_target.shape[0] return aepe_array
def iterative(net, gridFine, source_img, target_img, flow_est, i, match_est, Transform, coarsePlus, transformation, nbIter=1000, tolerance=0.03, nbPoint=4): if i not in coarsePlus: match_est = (match_est * ( ((flow_est.narrow(3, 0, 1) >= -1) * (flow_est.narrow(3, 0, 1) <= 1)).type(torch.cuda.FloatTensor) * ((flow_est.narrow(3, 1, 1) >= -1) * (flow_est.narrow(3, 1, 1) <= 1)).type(torch.cuda.FloatTensor) ).permute(0, 3, 1, 2)).squeeze().cpu().numpy() match_est = (match_est > 0.5) ix, iy = np.where(match_est) gridArr = gridFine.squeeze().cpu().numpy() flow_estArr = flow_est.squeeze().cpu().numpy() match2 = np.concatenate((gridArr[ix, iy], np.ones((len(ix), 1))), axis=1) match1 = np.concatenate((flow_estArr[ix, iy], np.ones((len(ix), 1))), axis=1) if len(match1) > nbPoint: bestParam, bestInlier, match1Inlier, match2Inlier = outil.RANSAC( nbIter, match1, match2, tolerance, nbPoint, Transform) else: bestParam = np.eye(3) bestParam = bestParam.astype(np.float32) coarsePlus[i] = bestParam else: bestParam = coarsePlus[i] if transformation == 'Affine': grid = F.affine_grid( torch.from_numpy( bestParam[:2].astype(np.float32)).unsqueeze(0).cuda(), target_img.size()) # theta should be of size N×2×3 else: bestParam = bestParam.astype(np.float32) warper = tgm.HomographyWarper(target_img.size()[2], target_img.size()[3]) grid = warper.warp_grid( torch.from_numpy(bestParam).unsqueeze(0).cuda()) IsSample = F.grid_sample(source_img, grid) sourceFeat = F.normalize(net['netFeatCoarse'](IsSample)) targetFeat = F.normalize(net['netFeatCoarse'](target_img)) corr12 = net['netCorr'](targetFeat, sourceFeat) _, flow_est = model.predFlowCoarse(corr12, net['netFlowCoarse'], gridFine) flow_est = F.grid_sample(grid.permute(0, 3, 1, 2), flow_est).permute(0, 2, 3, 1).contiguous() return flow_est
parser.add_argument('--onlyCoarse', action='store_true', help='only coarse?') args = parser.parse_args() print(args) strideNet = 16 res = {} gridY = torch.linspace(-1, 1, steps=args.minSize).view(1, -1, 1, 1).expand( 1, args.minSize, args.minSize, 1) gridX = torch.linspace(-1, 1, steps=args.minSize).view(1, 1, -1, 1).expand( 1, args.minSize, args.minSize, 1) grid = torch.cat((gridX, gridY), dim=3) warper = tgm.HomographyWarper(args.minSize, args.minSize) test_scene = os.listdir(args.finePth) getFlow = getFlow_onlyCoarse if args.onlyCoarse else getFlow_all for scene in test_scene: finePath = os.path.join(args.finePth, scene) coarsePath = os.path.join(args.coarsePth, scene) flowList = os.listdir(finePath) print('evaluating for scene {} ....'.format(scene)) res[scene] = [] csv_file = os.path.join(args.csv_path, 'hpatches_1_{}.csv'.format(scene))
dict_pairid_nbH = dict(bg) print (dict_pairid_nbH) if args.noc : args.gtPath = '../../data/Kitti/training/flow_noc/' else : args.gtPath = '../../data/Kitti/training/flow_occ/' getFlow = getFlow_onlyCoarse if args.onlyCoarse else getFlow_all for i in tqdm(range(nbImg)) : path = os.path.join(args.gtPath, '{0:06}_10.png'.format(i)) u, v, valid = readFlow(path) Ith, Itw = u.shape[0], u.shape[1] warper_org = tgm.HomographyWarper(Ith, Itw) #### -- org grid gridY = torch.linspace(-1, 1, steps = Ith).view(1, -1, 1, 1).expand(1, Ith, Itw, 1) gridX = torch.linspace(-1, 1, steps = Itw).view(1, 1, -1, 1).expand(1, Ith, Itw, 1) grid_org = torch.cat((gridX, gridY), dim=3) find = True if str(i) not in dict_pairid_nbH : flow = grid_org find = False else : nbH = dict_pairid_nbH[str(i)] flow = getFlow(str(i), args.predDir, nbH, args.resName, warper_org, args.multiH, grid_org, args.th, args.cc_th, args.interpolate) flow = flow - grid_org
It_bg_org = coarseModel.skyFromSeg(savepath / "tmp_It.jpg") It_bg = 1 - imresize(It_bg_org, (Ith, Itw)).astype(np.float32) ## 0 is bg if False: with torch.no_grad(): featt = F.normalize(network['netFeatCoarse'](coarseModel.ItTensor)) #### -- grid gridY = torch.linspace(-1, 1, steps=Ith, device=device).view(1, -1, 1, 1).expand(1, Ith, Itw, 1) gridX = torch.linspace(-1, 1, steps=Itw, device=device).view(1, 1, -1, 1).expand(1, Ith, Itw, 1) grid = torch.cat((gridX, gridY), dim=3) warper = tgm.HomographyWarper(Ith, Itw) ## update mask in every iteration Mask = np.zeros( (Ith, Itw), dtype=np.float32 ) # 0 means new region need to be explored, 1 means masked regions Coarse_Flow_Tensor = [] Coarse_Mask_Tensor = [] CoarsePlus_Flow_Tensor = [] CoarsePlus_Mask_Tensor = [] Fine_Flow_Tensor = [] Fine_Mask_Tensor = []
def align_image(source_path, output_filename, destination_path, base_path): I1 = Image.open(source_path).convert('RGB') I2 = Image.open(destination_path).convert('RGB') ### 7 scales, setting ransac parameters nbScale = 7 coarseIter = 10000 coarsetolerance = 0.05 minSize = 400 imageNet = True # we can also use MOCO feature here scaleR = 1.2 coarseModel = CoarseAlign(nbScale, coarseIter, coarsetolerance, 'Homography', minSize, 1, True, imageNet, scaleR) coarseModel.setSource(I1) coarseModel.setTarget(I2) I2w, I2h = coarseModel.It.size featt = F.normalize(network['netFeatCoarse'](coarseModel.ItTensor)) #### -- grid gridY = torch.linspace(-1, 1, steps=I2h).view(1, -1, 1, 1).expand(1, I2h, I2w, 1) gridX = torch.linspace(-1, 1, steps=I2w).view(1, 1, -1, 1).expand(1, I2h, I2w, 1) grid = torch.cat((gridX, gridY), dim=3).cuda() warper = tgm.HomographyWarper(I2h, I2w) bestPara, InlierMask = coarseModel.getCoarse(np.zeros((I2h, I2w))) bestPara = torch.from_numpy(bestPara).unsqueeze(0).cuda() ### Coarse Alignment flowCoarse = warper.warp_grid(bestPara) I1_coarse = F.grid_sample(coarseModel.IsTensor, flowCoarse) I1_coarse_pil = transforms.ToPILImage()(I1_coarse.cpu().squeeze()) plt.figure(figsize=(20, 10)) plt.subplot(1, 3, 1) plt.axis('off') plt.title('Source Image (Coarse)') plt.imshow(I1_coarse_pil) plt.subplot(1, 3, 2) plt.axis('off') plt.title('Target Image') plt.imshow(I2) plt.subplot(1, 3, 3) plt.title('Overlapped Image') plt.imshow(I2) plt.imshow(get_Avg_Image(I1_coarse_pil, coarseModel.It)) plt.show() ### Fine Alignment featsSample = F.normalize(network['netFeatCoarse'](I1_coarse.cuda())) corr12 = network['netCorr'](featt, featsSample) flowDown8 = network['netFlowCoarse']( corr12, False) ## output is with dimension B, 2, W, H flowUp = F.interpolate(flowDown8, size=(grid.size()[1], grid.size()[2]), mode='bilinear') flowUp = flowUp.permute(0, 2, 3, 1) flowUp = flowUp + grid flow12 = F.grid_sample(flowCoarse.permute(0, 3, 1, 2), flowUp).permute(0, 2, 3, 1).contiguous() I1_fine = F.grid_sample(coarseModel.IsTensor, flow12) I1_fine_pil = transforms.ToPILImage()(I1_fine.cpu().squeeze()) if not os.path.exists(os.path.join(base_path, "output")): os.mkdir(os.path.join(base_path, "output")) I1_fine_pil.save(os.path.join(base_path, "output", output_filename)) plt.figure(figsize=(20, 10)) plt.subplot(1, 3, 1) plt.axis('off') plt.title('Source Image (Fine Alignment)') plt.imshow(I1_fine_pil) plt.subplot(1, 3, 2) plt.axis('off') plt.title('Target Image') plt.imshow(I2) plt.subplot(1, 3, 3) plt.axis('off') plt.title('Overlapped Image') plt.imshow(get_Avg_Image(I1_fine_pil, coarseModel.It)) plt.show()