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