def pck_metric(batch,batch_start_idx,theta_aff,theta_tps,theta_aff_tps,model_tps,stats,args,use_cuda=True):
alpha = args.pck_alpha
do_aff = theta_aff is not None
do_tps = theta_tps is not None
do_aff_tps = theta_aff_tps is not None
source_im_size = batch['source_im_size']
target_im_size = batch['target_im_size']
source_points = batch['source_points']
target_points = batch['target_points']
# Instantiate point transformer
pt = PointTnf(use_cuda=use_cuda,
tps_reg_factor=args.tps_reg_factor)
# warp points with estimated transformations
target_points_norm = PointsToUnitCoords(target_points,target_im_size)
if do_aff:
# do affine only
warped_points_aff_norm = pt.affPointTnf(theta_aff,target_points_norm)
warped_points_aff = PointsToPixelCoords(warped_points_aff_norm,source_im_size)
if do_tps:
# do tps only
warped_points_tps_norm = pt.defPointTnf(theta_tps,target_points_norm,model_tps)
warped_points_tps = PointsToPixelCoords(warped_points_tps_norm,source_im_size)
if do_aff_tps:
# do tps+affine
warped_points_aff_tps_norm = pt.defPointTnf(theta_aff_tps,target_points_norm,model_tps)
warped_points_aff_tps_norm = pt.affPointTnf(theta_aff,warped_points_aff_tps_norm)
warped_points_aff_tps = PointsToPixelCoords(warped_points_aff_tps_norm,source_im_size)
L_pck = batch['L_pck'].data
current_batch_size=batch['source_im_size'].size(0)
indices = range(batch_start_idx,batch_start_idx+current_batch_size)
# import pdb; pdb.set_trace()
if do_aff:
pck_aff = pck(source_points.data, warped_points_aff.data, L_pck, alpha)
if do_tps:
pck_tps = pck(source_points.data, warped_points_tps.data, L_pck, alpha)
if do_aff_tps:
pck_aff_tps = pck(source_points.data, warped_points_aff_tps.data, L_pck, alpha)
if do_aff:
stats['aff']['pck'][indices] = pck_aff.unsqueeze(1).cpu().numpy()
if do_tps:
stats['tps']['pck'][indices] = pck_tps.unsqueeze(1).cpu().numpy()
if do_aff_tps:
stats['aff_tps']['pck'][indices] = pck_aff_tps.unsqueeze(1).cpu().numpy()
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
class TransformedGridLoss(nn.Module):
def __init__(self, geometric_model='affine', use_cuda=True, grid_size=20):
super(TransformedGridLoss, self).__init__()
self.geometric_model = geometric_model
# define virtual grid of points to be transformed
axis_coords = np.linspace(-1,1,grid_size)
self.N = grid_size*grid_size
X,Y = np.meshgrid(axis_coords,axis_coords)
X = np.reshape(X,(1,1,self.N))
Y = np.reshape(Y,(1,1,self.N))
P = np.concatenate((X,Y),1)
self.P = Variable(torch.FloatTensor(P),requires_grad=False)
self.pointTnf = PointTnf(use_cuda=use_cuda)
if use_cuda:
self.P = self.P.cuda();
def forward(self, theta, theta_GT, model=None):
# expand grid according to batch size
batch_size = theta.size()[0]
P = self.P.expand(batch_size,2,self.N)
# compute transformed grid points using estimated and GT tnfs
if self.geometric_model=='affine':
P_prime = self.pointTnf.affPointTnf(theta,P)
P_prime_GT = self.pointTnf.affPointTnf(theta_GT,P)
elif self.geometric_model=='tps':
P_prime = self.pointTnf.defPointTnf(theta,P,model)
P_prime_GT = self.pointTnf.tpsPointTnf(theta_GT,P)
# compute MSE loss on transformed grid points
loss = torch.sum(torch.pow(P_prime - P_prime_GT,2),1)
loss = torch.mean(loss)
return loss
def area_metrics(batch,
batch_start_idx,
theta_aff,
theta_tps,
theta_aff_tps,
model_tps,
stats,
args,
use_cuda=True):
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)
pt = PointTnf(use_cuda=use_cuda)
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())
target_mask_np, target_mask = poly_str_to_mask(
batch['target_polygon'][0][b],
batch['target_polygon'][1][b],
h_tgt,
w_tgt,
use_cuda=use_cuda)
source_mask_np, source_mask = poly_str_to_mask(
batch['source_polygon'][0][b],
batch['source_polygon'][1][b],
h_src,
w_src,
use_cuda=use_cuda)
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))
warped_mask_aff = F.grid_sample(source_mask, grid_aff)
flow_aff = th_sampling_grid_to_np_flow(source_grid=grid_aff,
h_src=h_src,
w_src=w_src)
stats['aff']['intersection_over_union'][
idx] = intersection_over_union(warped_mask_aff, target_mask)
stats['aff']['label_transfer_accuracy'][
idx] = label_transfer_accuracy(warped_mask_aff, target_mask)
stats['aff']['localization_error'][idx] = localization_error(
source_mask_np, target_mask_np, flow_aff)
if do_tps:
grid_tps = pointsToGrid(
pt.defPointTnf(theta_tps[b, :].unsqueeze(0), grid_XY_vec,
model_tps))
warped_mask_tps = F.grid_sample(source_mask, grid_tps)
flow_tps = th_sampling_grid_to_np_flow(source_grid=grid_tps,
h_src=h_src,
w_src=w_src)
stats['tps']['intersection_over_union'][
idx] = intersection_over_union(warped_mask_tps, target_mask)
stats['tps']['label_transfer_accuracy'][
idx] = label_transfer_accuracy(warped_mask_tps, target_mask)
stats['tps']['localization_error'][idx] = localization_error(
source_mask_np, target_mask_np, flow_tps)
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)))
warped_mask_aff_tps = F.grid_sample(source_mask, grid_aff_tps)
flow_aff_tps = th_sampling_grid_to_np_flow(
source_grid=grid_aff_tps, h_src=h_src, w_src=w_src)
stats['aff_tps']['intersection_over_union'][
idx] = intersection_over_union(warped_mask_aff_tps,
target_mask)
stats['aff_tps']['label_transfer_accuracy'][
idx] = label_transfer_accuracy(warped_mask_aff_tps,
target_mask)
stats['aff_tps']['localization_error'][idx] = localization_error(
source_mask_np, target_mask_np, flow_aff_tps)
return stats