def flow_metrics(batch, batch_start_idx, theta_det, theta_aff, theta_tps,
theta_afftps, results, args):
result_path = args.flow_output_dir
do_det = theta_det is not None
do_aff = theta_aff is not None
do_tps = theta_tps is not None
do_aff_tps = theta_afftps is not None
pt = PointTnf(use_cuda=args.cuda)
batch_size = batch['source_im_info'].size(0)
for b in range(batch_size):
# Get H, W of source and target image
h_src = int(batch['source_im_info'][b, 0].cpu().numpy())
w_src = int(batch['source_im_info'][b, 1].cpu().numpy())
h_tgt = int(batch['target_im_info'][b, 0].cpu().numpy())
w_tgt = int(batch['target_im_info'][b, 1].cpu().numpy())
# Generate grid for warping
grid_X, grid_Y = np.meshgrid(np.linspace(-1, 1, w_tgt),
np.linspace(-1, 1, h_tgt))
# grid_X, grid_Y.shape: (1, h_tgt, w_tgt, 1)
grid_X = torch.Tensor(grid_X).unsqueeze(0).unsqueeze(3)
grid_Y = torch.Tensor(grid_Y).unsqueeze(0).unsqueeze(3)
grid_X.requires_grad = False
grid_Y.requires_grad = False
if args.cuda:
grid_X = grid_X.cuda()
grid_Y = grid_Y.cuda()
# Reshape to vector, grid_X_vec, grid_Y_vec.shape: (1, 1, h_tgt * w_tgt)
grid_X_vec = grid_X.view(1, 1, -1)
grid_Y_vec = grid_Y.view(1, 1, -1)
# grid_XY_vec.shape: (1, 2, h_tgt * w_tgt)
grid_XY_vec = torch.cat((grid_X_vec, grid_Y_vec), 1)
# Transform vector of points to grid
def pointsToGrid(x, h_tgt=h_tgt, w_tgt=w_tgt):
return x.contiguous().view(1, 2, h_tgt, w_tgt).permute(0, 2, 3, 1)
idx = batch_start_idx + b
if do_det:
grid_det = pointsToGrid(
pt.affPointTnf(theta=theta_det[b, :].unsqueeze(0),
points=grid_XY_vec))
flow_det = th_sampling_grid_to_np_flow(source_grid=grid_det,
h_src=h_src,
w_src=w_src)
flow_det_path = os.path.join(result_path, 'det',
batch['flow_path'][b])
# create_file_path(flow_det_path)
# write_flo_file(flow_det, flow_det_path)
if do_aff:
if do_det:
key = 'det_aff'
grid_aff = pointsToGrid(
pt.affPointTnf(theta=theta_det[b, :].unsqueeze(0),
points=pt.affPointTnf(
theta=theta_aff[b, :].unsqueeze(0),
points=grid_XY_vec)))
else:
key = 'aff'
grid_aff = pointsToGrid(
pt.affPointTnf(theta=theta_aff[b, :].unsqueeze(0),
points=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, key,
batch['flow_path'][b])
# create_file_path(flow_aff_path)
# write_flo_file(flow_aff, flow_aff_path)
if do_tps:
# vis = Visdom()
# flow_gt = batch['flow_gt'][b]
# vis.heatmap(flow_gt.cpu().numpy()[:, ::-1, 0])
# vis.heatmap(flow_gt.cpu().numpy()[:, ::-1, 1])
# flow_gt_img = visualize_flow(flow_gt.cpu().numpy())
# vis.image(flow_gt_img.transpose((2, 0, 1)))
# vis.mesh(grid_XY_vec.squeeze().transpose(0, 1), opts=dict(opacity=0.3))
# vis.mesh(pt.tpsPointTnf(theta=theta_tps[b, :].unsqueeze(0), points=grid_XY_vec).squeeze().transpose(0, 1), opts=dict(opacity=0.3))
# grid_XY = pointsToGrid(grid_XY_vec).squeeze(0)
# in_bound_mask = (grid_XY[:, :, 0] > -1) & (grid_XY[:, :, 0] < 1) & (grid_XY[:, :, 1] > -1) & (grid_XY[:, :, 1] < 1)
# vis.heatmap(in_bound_mask)
# Get sampling grid with predicted TPS parameters, grid_tps.shape: (1, h_tgt, w_tgt, 2)
grid_tps = pointsToGrid(
pt.tpsPointTnf(theta=theta_tps[b, :].unsqueeze(0),
points=grid_XY_vec))
# Transform sampling grid to flow
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:
if do_det:
key = 'det_aff_tps'
grid_aff_tps = pointsToGrid(
pt.affPointTnf(
theta=theta_det[b, :].unsqueeze(0),
points=pt.affPointTnf(
theta=theta_aff[b, :].unsqueeze(0),
points=pt.tpsPointTnf(
theta=theta_afftps[b, :].unsqueeze(0),
points=grid_XY_vec))))
else:
key = 'afftps'
grid_aff_tps = pointsToGrid(
pt.affPointTnf(theta=theta_aff[b, :].unsqueeze(0),
points=pt.tpsPointTnf(
theta=theta_afftps[b, :].unsqueeze(0),
points=grid_XY_vec)))
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, key,
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 results
def area_metrics(batch, batch_start_idx, theta_det, theta_aff, theta_tps,
theta_afftps, results, args):
do_det = theta_det is not None
do_aff = theta_aff is not None
do_tps = theta_tps is not None
do_aff_tps = theta_afftps is not None
batch_size = batch['source_im_info'].size(0)
pt = PointTnf(use_cuda=args.cuda)
for b in range(batch_size):
# Get H, W of source and target image
h_src = int(batch['source_im_info'][b, 0].cpu().numpy())
w_src = int(batch['source_im_info'][b, 1].cpu().numpy())
h_tgt = int(batch['target_im_info'][b, 0].cpu().numpy())
w_tgt = int(batch['target_im_info'][b, 1].cpu().numpy())
# Transform annotated polygon to mask using given coordinates of key points
# target_mask_np.shape: (h_tgt, w_tgt), target_mask.shape: (1, 1, h_tgt, w_tgt)
target_mask_np, target_mask = poly_str_to_mask(
poly_x_str=batch['target_polygon'][0][b],
poly_y_str=batch['target_polygon'][1][b],
out_h=h_tgt,
out_w=w_tgt,
use_cuda=args.cuda)
source_mask_np, source_mask = poly_str_to_mask(
poly_x_str=batch['source_polygon'][0][b],
poly_y_str=batch['source_polygon'][1][b],
out_h=h_src,
out_w=w_src,
use_cuda=args.cuda)
# Generate grid for warping
grid_X, grid_Y = np.meshgrid(np.linspace(-1, 1, w_tgt),
np.linspace(-1, 1, h_tgt))
# grid_X, grid_Y.shape: (1, h_tgt, w_tgt, 1)
grid_X = torch.Tensor(grid_X.astype(
np.float32)).unsqueeze(0).unsqueeze(3)
grid_Y = torch.Tensor(grid_Y.astype(
np.float32)).unsqueeze(0).unsqueeze(3)
grid_X.requires_grad = False
grid_Y.requires_grad = False
if args.cuda:
grid_X = grid_X.cuda()
grid_Y = grid_Y.cuda()
# Reshape to vector, grid_X_vec, grid_Y_vec.shape: (1, 1, h_tgt * w_tgt)
grid_X_vec = grid_X.view(1, 1, -1)
grid_Y_vec = grid_Y.view(1, 1, -1)
# grid_XY_vec.shape: (1, 2, h_tgt * w_tgt)
grid_XY_vec = torch.cat((grid_X_vec, grid_Y_vec), 1)
# Transform vector of points to grid
def pointsToGrid(x, h_tgt=h_tgt, w_tgt=w_tgt):
return x.contiguous().view(1, 2, h_tgt, w_tgt).permute(0, 2, 3, 1)
idx = batch_start_idx + b
if do_det:
grid_det = pointsToGrid(
pt.affPointTnf(theta=theta_det[b, :].unsqueeze(0),
points=grid_XY_vec))
warped_mask_det = F.grid_sample(source_mask, grid_det)
flow_det = th_sampling_grid_to_np_flow(source_grid=grid_det,
h_src=h_src,
w_src=w_src)
results['det']['intersection_over_union'][
idx] = intersection_over_union(warped_mask=warped_mask_det,
target_mask=target_mask)
results['det']['label_transfer_accuracy'][
idx] = label_transfer_accuracy(warped_mask=warped_mask_det,
target_mask=target_mask)
results['det']['localization_error'][idx] = localization_error(
source_mask_np=source_mask_np,
target_mask_np=target_mask_np,
flow_np=flow_det)
if do_aff:
if do_det:
key = 'det_aff'
grid_aff = pointsToGrid(
pt.affPointTnf(theta=theta_det[b, :].unsqueeze(0),
points=pt.affPointTnf(
theta=theta_aff[b, :].unsqueeze(0),
points=grid_XY_vec)))
else:
key = 'aff'
grid_aff = pointsToGrid(
pt.affPointTnf(theta=theta_aff[b, :].unsqueeze(0),
points=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)
results[key]['intersection_over_union'][
idx] = intersection_over_union(warped_mask=warped_mask_aff,
target_mask=target_mask)
results[key]['label_transfer_accuracy'][
idx] = label_transfer_accuracy(warped_mask=warped_mask_aff,
target_mask=target_mask)
results[key]['localization_error'][idx] = localization_error(
source_mask_np=source_mask_np,
target_mask_np=target_mask_np,
flow_np=flow_aff)
if do_tps:
# Get sampling grid with predicted TPS parameters, grid_tps.shape: (1, h_tgt, w_tgt, 2)
grid_tps = pointsToGrid(
pt.tpsPointTnf(theta=theta_tps[b, :].unsqueeze(0),
points=grid_XY_vec))
warped_mask_tps = F.grid_sample(
source_mask, grid_tps) # Sampling source_mask with warped grid
# Transform sampling grid to flow
flow_tps = th_sampling_grid_to_np_flow(source_grid=grid_tps,
h_src=h_src,
w_src=w_src)
results['tps']['intersection_over_union'][
idx] = intersection_over_union(warped_mask=warped_mask_tps,
target_mask=target_mask)
results['tps']['label_transfer_accuracy'][
idx] = label_transfer_accuracy(warped_mask=warped_mask_tps,
target_mask=target_mask)
results['tps']['localization_error'][idx] = localization_error(
source_mask_np=source_mask_np,
target_mask_np=target_mask_np,
flow_np=flow_tps)
if do_aff_tps:
if do_det:
key = 'det_aff_tps'
grid_aff_tps = pointsToGrid(
pt.affPointTnf(
theta=theta_det[b, :].unsqueeze(0),
points=pt.affPointTnf(
theta=theta_aff[b, :].unsqueeze(0),
points=pt.tpsPointTnf(
theta=theta_afftps[b, :].unsqueeze(0),
points=grid_XY_vec))))
else:
key = 'afftps'
grid_aff_tps = pointsToGrid(
pt.affPointTnf(theta=theta_aff[b, :].unsqueeze(0),
points=pt.tpsPointTnf(
theta=theta_afftps[b, :].unsqueeze(0),
points=grid_XY_vec)))
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)
results[key]['intersection_over_union'][
idx] = intersection_over_union(warped_mask=warped_mask_aff_tps,
target_mask=target_mask)
results[key]['label_transfer_accuracy'][
idx] = label_transfer_accuracy(warped_mask=warped_mask_aff_tps,
target_mask=target_mask)
results[key]['localization_error'][idx] = localization_error(
source_mask_np=source_mask_np,
target_mask_np=target_mask_np,
flow_np=flow_aff_tps)
return results
