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