def flow_metrics(batch,batch_start_idx,theta_1,theta_2,geometric_model_1,geometric_model_2,stats,args,use_cuda=True):
two_stage=(geometric_model_2 is not None)
result_path=args.flow_output_dir
pt=PointTnf(use_cuda=use_cuda)
if geometric_model_1=='affine':
tnf_1 = pt.affPointTnf
elif geometric_model_1=='hom':
tnf_1 = pt.homPointTnf
elif geometric_model_1=='tps':
tnf_1 = pt.tpsPointTnf
if two_stage:
if geometric_model_2=='affine':
tnf_2 = pt.affPointTnf
elif geometric_model_2=='hom':
tnf_2 = pt.homPointTnf
elif geometric_model_2=='tps':
tnf_2 = pt.tpsPointTnf
batch_size=batch['source_im_size'].size(0)
for b in range(batch_size):
h_src = int(batch['source_im_size'][b,0].data.cpu().numpy())
w_src = int(batch['source_im_size'][b,1].data.cpu().numpy())
h_tgt = int(batch['target_im_size'][b,0].data.cpu().numpy())
w_tgt = int(batch['target_im_size'][b,1].data.cpu().numpy())
grid_X,grid_Y = np.meshgrid(np.linspace(-1,1,w_tgt),np.linspace(-1,1,h_tgt))
grid_X = torch.FloatTensor(grid_X).unsqueeze(0).unsqueeze(3)
grid_Y = torch.FloatTensor(grid_Y).unsqueeze(0).unsqueeze(3)
grid_X = Variable(grid_X,requires_grad=False)
grid_Y = Variable(grid_Y,requires_grad=False)
if use_cuda:
grid_X = grid_X.cuda()
grid_Y = grid_Y.cuda()
grid_X_vec = grid_X.view(1,1,-1)
grid_Y_vec = grid_Y.view(1,1,-1)
grid_XY_vec = torch.cat((grid_X_vec,grid_Y_vec),1)
def pointsToGrid (x,h_tgt=h_tgt,w_tgt=w_tgt): return x.contiguous().view(1,2,h_tgt,w_tgt).transpose(1,2).transpose(2,3)
grid_1 = pointsToGrid(tnf_1(theta_1[b,:].unsqueeze(0),grid_XY_vec))
flow_1 = th_sampling_grid_to_np_flow(source_grid=grid_1,h_src=h_src,w_src=w_src)
flow_1_path = os.path.join(result_path,geometric_model_1,batch['flow_path'][b])
create_file_path(flow_1_path)
write_flo_file(flow_1,flow_1_path)
if two_stage:
grid_1_2 = pointsToGrid(tnf_1(theta_1[b,:].unsqueeze(0),tnf_2(theta_2[b,:].unsqueeze(0),grid_XY_vec)))
flow_1_2 = th_sampling_grid_to_np_flow(source_grid=grid_1_2,h_src=h_src,w_src=w_src)
flow_1_2_path = os.path.join(result_path,geometric_model_1+'_'+geometric_model_2,batch['flow_path'][b])
create_file_path(flow_1_2_path)
write_flo_file(flow_1_2,flow_1_2_path)
return stats
def flow_metrics(batch, batch_start_idx, matches, stats, args, use_cuda=True):
result_path = args.flow_output_dir
pt = PointTnf(use_cuda=use_cuda)
batch_size = batch['source_im_size'].size(0)
for b in range(batch_size):
h_src = int(batch['source_im_size'][b, 0].data.cpu().numpy())
w_src = int(batch['source_im_size'][b, 1].data.cpu().numpy())
h_tgt = int(batch['target_im_size'][b, 0].data.cpu().numpy())
w_tgt = int(batch['target_im_size'][b, 1].data.cpu().numpy())
grid_X, grid_Y = np.meshgrid(np.linspace(-1, 1, w_tgt),
np.linspace(-1, 1, h_tgt))
grid_X = torch.FloatTensor(grid_X).unsqueeze(0).unsqueeze(3)
grid_Y = torch.FloatTensor(grid_Y).unsqueeze(0).unsqueeze(3)
grid_X = Variable(grid_X, requires_grad=False)
grid_Y = Variable(grid_Y, requires_grad=False)
if use_cuda:
grid_X = grid_X.cuda()
grid_Y = grid_Y.cuda()
grid_X_vec = grid_X.view(1, 1, -1)
grid_Y_vec = grid_Y.view(1, 1, -1)
grid_XY_vec = torch.cat((grid_X_vec, grid_Y_vec), 1)
def pointsToGrid(x, h_tgt=h_tgt, w_tgt=w_tgt):
return x.contiguous().view(1, 2, h_tgt,
w_tgt).transpose(1, 2).transpose(2, 3)
idx = batch_start_idx + b
source_im_size = batch['source_im_size']
warped_points_norm = bilinearInterpPointTnf(matches, grid_XY_vec)
# warped_points = PointsToPixelCoords(warped_points_norm,source_im_size)
warped_points = pointsToGrid(warped_points_norm)
# grid_aff = pointsToGrid(pt.affPointTnf(theta_aff[b, :].unsqueeze(0), grid_XY_vec))
flow_aff = th_sampling_grid_to_np_flow(source_grid=warped_points,
h_src=h_src,
w_src=w_src)
flow_aff_path = os.path.join(result_path, 'nc', batch['flow_path'][b])
create_file_path(flow_aff_path)
write_flo_file(flow_aff, flow_aff_path)
idx = batch_start_idx + b
return stats
def flow_metrics(batch,
batch_start_idx,
theta_aff,
theta_tps,
theta_aff_tps,
stats,
args,
use_cuda=True):
result_path = args.flow_output_dir
do_aff = theta_aff is not None
do_tps = theta_tps is not None
do_aff_tps = theta_aff_tps is not None
batch_size = batch['source_im_size'].size(0)
for b in range(batch_size):
h_src = int(batch['source_im_size'][b, 0].data.cpu().numpy())
w_src = int(batch['source_im_size'][b, 1].data.cpu().numpy())
h_tgt = int(batch['target_im_size'][b, 0].data.cpu().numpy())
w_tgt = int(batch['target_im_size'][b, 1].data.cpu().numpy())
grid_aff, grid_tps, grid_aff_tps = theta_to_sampling_grid(
h_tgt,
w_tgt,
theta_aff[b, :] if do_aff else None,
theta_tps[b, :] if do_tps else None,
theta_aff_tps[b, :] if do_aff_tps else None,
use_cuda=use_cuda,
tps_reg_factor=args.tps_reg_factor)
if do_aff_tps:
flow_aff_tps = th_sampling_grid_to_np_flow(
source_grid=grid_aff_tps, h_src=h_src, w_src=w_src)
flow_aff_tps_path = os.path.join(result_path, 'aff_tps',
batch['flow_path'][b])
create_file_path(flow_aff_tps_path)
write_flo_file(flow_aff_tps, flow_aff_tps_path)
elif do_aff:
flow_aff = th_sampling_grid_to_np_flow(source_grid=grid_aff,
h_src=h_src,
w_src=w_src)
flow_aff_path = os.path.join(result_path, 'aff',
batch['flow_path'][b])
create_file_path(flow_aff_path)
write_flo_file(flow_aff, flow_aff_path)
elif do_tps:
flow_tps = th_sampling_grid_to_np_flow(source_grid=grid_tps,
h_src=h_src,
w_src=w_src)
flow_tps_path = os.path.join(result_path, 'tps',
batch['flow_path'][b])
create_file_path(flow_tps_path)
write_flo_file(flow_tps, flow_tps_path)
idx = batch_start_idx + b
return stats
def flow_metrics(batch,batch_start_idx,theta_aff,theta_tps,theta_aff_tps,model_tps,stats,args,use_cuda=True):
result_path=args.flow_output_dir
do_aff = theta_aff is not None
do_tps = theta_tps is not None
do_aff_tps = theta_aff_tps is not None
pt=PointTnf(use_cuda=use_cuda)
batch_size=batch['source_im_size'].size(0)
for b in range(batch_size):
h_src = int(batch['source_im_size'][b,0].data.cpu().numpy())
w_src = int(batch['source_im_size'][b,1].data.cpu().numpy())
h_tgt = int(batch['target_im_size'][b,0].data.cpu().numpy())
w_tgt = int(batch['target_im_size'][b,1].data.cpu().numpy())
grid_X,grid_Y = np.meshgrid(np.linspace(-1,1,w_tgt),np.linspace(-1,1,h_tgt))
grid_X = torch.FloatTensor(grid_X).unsqueeze(0).unsqueeze(3)
grid_Y = torch.FloatTensor(grid_Y).unsqueeze(0).unsqueeze(3)
grid_X = Variable(grid_X,requires_grad=False)
grid_Y = Variable(grid_Y,requires_grad=False)
if use_cuda:
grid_X = grid_X.cuda()
grid_Y = grid_Y.cuda()
grid_X_vec = grid_X.view(1,1,-1)
grid_Y_vec = grid_Y.view(1,1,-1)
grid_XY_vec = torch.cat((grid_X_vec,grid_Y_vec),1)
def pointsToGrid (x,h_tgt=h_tgt,w_tgt=w_tgt): return x.contiguous().view(1,2,h_tgt,w_tgt).transpose(1,2).transpose(2,3)
idx = batch_start_idx+b
if do_aff:
grid_aff = pointsToGrid(pt.affPointTnf(theta_aff[b,:].unsqueeze(0),grid_XY_vec))
flow_aff = th_sampling_grid_to_np_flow(source_grid=grid_aff,h_src=h_src,w_src=w_src)
flow_aff_path = os.path.join(result_path,'aff',batch['flow_path'][b])
create_file_path(flow_aff_path)
write_flo_file(flow_aff,flow_aff_path)
if do_tps:
grid_tps = pointsToGrid(pt.defPointTnf(theta_tps[b,:].unsqueeze(0),grid_XY_vec,model_tps))
flow_tps = th_sampling_grid_to_np_flow(source_grid=grid_tps,h_src=h_src,w_src=w_src)
flow_tps_path = os.path.join(result_path,'tps',batch['flow_path'][b])
create_file_path(flow_tps_path)
write_flo_file(flow_tps,flow_tps_path)
if do_aff_tps:
grid_aff_tps = pointsToGrid(pt.affPointTnf(theta_aff[b,:].unsqueeze(0),pt.defPointTnf(theta_aff_tps[b,:].unsqueeze(0),grid_XY_vec,model_tps)))
flow_aff_tps = th_sampling_grid_to_np_flow(source_grid=grid_aff_tps,h_src=h_src,w_src=w_src)
flow_aff_tps_path = os.path.join(result_path,'aff_tps',batch['flow_path'][b])
create_file_path(flow_aff_tps_path)
write_flo_file(flow_aff_tps,flow_aff_tps_path)
idx = batch_start_idx+b
return stats
def flow_metrics(batch,batch_start_idx,theta_aff,theta_tps,theta_aff_tps,stats,args,use_cuda=True):
result_path=args.flow_output_dir
do_aff = theta_aff is not None
do_tps = theta_tps is not None
do_aff_tps = theta_aff_tps is not None
batch_size=batch['source_im_size'].size(0)
for b in range(batch_size):
h_src = int(batch['source_im_size'][b,0].data.cpu().numpy())
w_src = int(batch['source_im_size'][b,1].data.cpu().numpy())
h_tgt = int(batch['target_im_size'][b,0].data.cpu().numpy())
w_tgt = int(batch['target_im_size'][b,1].data.cpu().numpy())
grid_aff,grid_tps,grid_aff_tps=theta_to_sampling_grid(h_tgt,w_tgt,
theta_aff[b,:] if do_aff else None,
theta_tps[b,:] if do_tps else None,
theta_aff_tps[b,:] if do_aff_tps else None,
use_cuda=use_cuda,
tps_reg_factor=args.tps_reg_factor)
if do_aff_tps:
flow_aff_tps = th_sampling_grid_to_np_flow(source_grid=grid_aff_tps,h_src=h_src,w_src=w_src)
flow_aff_tps_path = os.path.join(result_path,'aff_tps',batch['flow_path'][b])
create_file_path(flow_aff_tps_path)
write_flo_file(flow_aff_tps,flow_aff_tps_path)
elif do_aff:
flow_aff = th_sampling_grid_to_np_flow(source_grid=grid_aff,h_src=h_src,w_src=w_src)
flow_aff_path = os.path.join(result_path,'aff',batch['flow_path'][b])
create_file_path(flow_aff_path)
write_flo_file(flow_aff,flow_aff_path)
elif do_tps:
flow_tps = th_sampling_grid_to_np_flow(source_grid=grid_tps,h_src=h_src,w_src=w_src)
flow_tps_path = os.path.join(result_path,'tps',batch['flow_path'][b])
create_file_path(flow_tps_path)
write_flo_file(flow_tps,flow_tps_path)
idx = batch_start_idx+b
return stats