def par_generate_gt(config, pair_rec, flow_depth_rendered=None): from lib.pair_matching.flow import calc_flow target_size, max_size = config.SCALES[0][0], config.SCALES[0][1] if flow_depth_rendered is None: flow_depth_rendered = cv2.imread(pair_rec['depth_rendered'], cv2.IMREAD_UNCHANGED).astype(np.float32) flow_depth_rendered /= config.dataset.DEPTH_FACTOR flow_depth_rendered, _ = resize(flow_depth_rendered, target_size, max_size) if 'depth_render_real' in pair_rec: flow_depth_real = cv2.imread(pair_rec['depth_render_real'], cv2.IMREAD_UNCHANGED).astype(np.float32) flow_depth_real /= config.dataset.DEPTH_FACTOR else: print('not using render_real depth in par_generate_gt') flow_depth_real = cv2.imread(pair_rec['depth_real'], cv2.IMREAD_UNCHANGED).astype(np.float32) flow_depth_real /= config.dataset.DEPTH_FACTOR flow_depth_real, _ = resize(flow_depth_real, target_size, max_size) if 'mask_real_gt' or 'mask_real_est' in pair_rec: mask_real_path = pair_rec['mask_real_gt'] assert os.path.exists(mask_real_path), '%s does not exist'.format(pair_rec['mask_real_gt']) mask_real = cv2.imread(mask_real_path, cv2.IMREAD_UNCHANGED) mask_real, _ = resize(mask_real, target_size, max_size) flow_depth_real[mask_real != pair_rec['mask_idx']] = 0 if config.network.FLOW_I2R: if config.FLOW_CLASS_AGNOSTIC: flow_i2r, visible, _ = calc_flow(flow_depth_rendered, pair_rec['pose_est'], pair_rec['pose_real'], config.dataset.INTRINSIC_MATRIX, flow_depth_real, standard_rep=config.network.STANDARD_FLOW_REP) flow_i2r_list = [flow_i2r, visible, np.logical_and(visible == 0, flow_depth_rendered == 0)] else: raise Exception('NOT_IMPLEMENTED') else: flow_i2r_list = None if config.network.FLOW_R2I: if config.FLOW_CLASS_AGNOSTIC: flow_r2i, visible, _ = calc_flow(flow_depth_real, pair_rec['pose_real'], pair_rec['pose_est'], config.dataset.INTRINSIC_MATRIX, flow_depth_rendered, standard_rep=config.network.STANDARD_FLOW_REP) flow_r2i_list = [flow_r2i, visible, np.logical_and(visible == 0, flow_depth_real == 0)] else: raise Exception('NOT_IMPLEMENTED') else: flow_r2i_list = None return {'flow_i2r': flow_i2r_list, 'flow_r2i': flow_r2i_list}
def get_pair_flow(pairdb, config, scale_ind_list, phase="train", random_k=[]): from lib.pair_matching.flow import calc_flow num_pairs = len(pairdb) flow_tensor = [] flow_weights_tensor = [] X_rendered_valid_list = [] X_observed_valid_list = [] for i in range(num_pairs): pair_rec = pairdb[i] flow_depth_rendered = cv2.imread(pair_rec["depth_rendered"], cv2.IMREAD_UNCHANGED).astype( np.float32) flow_depth_rendered /= config.dataset.DEPTH_FACTOR if "depth_gt_observed" in pair_rec: flow_depth_observed_raw = cv2.imread(pair_rec["depth_gt_observed"], cv2.IMREAD_UNCHANGED).astype( np.float32) else: flow_depth_observed_raw = cv2.imread(pair_rec["depth_observed"], cv2.IMREAD_UNCHANGED).astype( np.float32) flow_depth_observed_raw /= config.dataset.DEPTH_FACTOR flow_depth_observed = flow_depth_observed_raw if config.network.PRED_FLOW or (config.train_iter.SE3_PM_LOSS): flow, visible, X_rendered_valid = calc_flow( flow_depth_rendered, pair_rec["pose_rendered"], pair_rec["pose_observed"], config.dataset.INTRINSIC_MATRIX, flow_depth_observed, standard_rep=config.network.STANDARD_FLOW_REP, ) # print('flow *'*20, flow.shape, np.unique(flow)) # print('flow weights *' * 20, visible.shape, np.unique(visible)) flow_tensor.append(flow.transpose((2, 0, 1))[np.newaxis, :, :, :]) if config.TRAIN.FLOW_WEIGHT_TYPE == "all": flow_weights = np.ones(visible.shape, dtype=np.float32) elif config.TRAIN.FLOW_WEIGHT_TYPE == "viz": flow_weights = visible elif config.TRAIN.FLOW_WEIGHT_TYPE == "valid": flow_weights = np.logical_or( np.squeeze(flow_depth_rendered == 0), visible) flow_weights_tensor.append( np.tile(flow_weights[np.newaxis, np.newaxis, :, :], (1, 2, 1, 1))) # flow_weights_tensor.append(flow_weights[np.newaxis, np.newaxis, :, :]) X_rendered_valid_list.append(X_rendered_valid) return ( flow_tensor, flow_weights_tensor, X_rendered_valid_list, X_observed_valid_list, )
def par_generate_gt(config, pair_rec, flow_depth_rendered=None): from lib.pair_matching.flow import calc_flow target_size, max_size = config.SCALES[0][0], config.SCALES[0][1] if flow_depth_rendered is None: flow_depth_rendered = cv2.imread(pair_rec["depth_rendered"], cv2.IMREAD_UNCHANGED).astype( np.float32) flow_depth_rendered /= config.dataset.DEPTH_FACTOR flow_depth_rendered, _ = resize(flow_depth_rendered, target_size, max_size) if "depth_gt_observed" in pair_rec: flow_depth_observed = cv2.imread(pair_rec["depth_gt_observed"], cv2.IMREAD_UNCHANGED).astype( np.float32) flow_depth_observed /= config.dataset.DEPTH_FACTOR else: logger.info("not using gt_observed depth in par_generate_gt") flow_depth_observed = cv2.imread(pair_rec["depth_observed"], cv2.IMREAD_UNCHANGED).astype( np.float32) flow_depth_observed /= config.dataset.DEPTH_FACTOR flow_depth_observed, _ = resize(flow_depth_observed, target_size, max_size) if "mask_gt_observed" or "mask_observed" in pair_rec: mask_observed_path = pair_rec["mask_gt_observed"] assert os.path.exists(mask_observed_path), "%s does not exist".format( pair_rec["mask_gt_observed"]) mask_observed = cv2.imread(mask_observed_path, cv2.IMREAD_UNCHANGED) mask_observed, _ = resize(mask_observed, target_size, max_size) flow_depth_observed[mask_observed != pair_rec["mask_idx"]] = 0 if config.network.PRED_FLOW: flow_i2r, visible, _ = calc_flow( flow_depth_rendered, pair_rec["pose_rendered"], pair_rec["pose_observed"], config.dataset.INTRINSIC_MATRIX, flow_depth_observed, standard_rep=config.network.STANDARD_FLOW_REP, ) flow_i2r_list = [ flow_i2r, visible, np.logical_and(visible == 0, flow_depth_rendered == 0) ] return {"flow": flow_i2r_list}
def visualize_flow(cfg, pair_rec, flow_i2r_est, image_src, depth_src, pose_src, image_tgt, depth_tgt, pose_tgt, visible_est=None): import matplotlib.pyplot as plt K = cfg.dataset.INTRINSIC_MATRIX flow, visible, _ = calc_flow(depth_src, pose_src, pose_tgt, K, depth_tgt, standard_rep=config.network.STANDARD_FLOW_REP) visible = visible def kernel_calc_epe(flow_pred, flow_gt, visible): h_diff = (flow_gt[:, :, 0] - flow_pred[:, :, 0])[visible != 0] w_diff = (flow_gt[:, :, 1] - flow_pred[:, :, 1])[visible != 0] diff = np.sqrt(np.square(h_diff) + np.square(w_diff)) print(flow_gt.shape, diff.shape, flow_pred.shape, visible.shape, np.sum(visible), np.max(depth_tgt), np.max(depth_src)) return diff.sum()/np.logical_and(flow_gt[:,:,0]!=0, flow_gt[:,:,1]!=0).sum() r_dist, t_dist = calc_rt_dist_m(pose_src, pose_tgt) print('epe: {}, rotation dist: {}, translation dist: {}'.format(kernel_calc_epe(flow_i2r_est, flow, visible), r_dist, t_dist)) fig = plt.figure() ax_im_src = fig.add_subplot(3,4,1) ax_im_src.set_title('image_src') plt.axis('off') plt.imshow(image_src) ax_im_tgt = fig.add_subplot(3,4,2) ax_im_tgt.set_title('image_tgt') plt.axis('off') plt.imshow(image_tgt) ax_d_src = fig.add_subplot(3,4,3) ax_d_src.set_title('depth_src') plt.axis('off') plt.imshow(depth_src) ax_d_tgt = fig.add_subplot(3,4,4) ax_d_tgt.set_title('depth_tgt') plt.axis('off') plt.imshow(depth_tgt) height = depth_tgt.shape[0] width = depth_tgt.shape[1] mesh_tgt = np.zeros((height, width, 3), np.uint8) mesh_src = np.zeros((height, width, 3), np.uint8) mesh_tgt_est = np.zeros((height, width, 3), np.uint8) def sigmoid(x): return 1 / (1 + np.exp(-x*20)) for h in range(height): for w in range(width): if visible[h, w]: cur_flow = flow[h,w,:].flatten() if not config.network.STANDARD_FLOW_REP: cur_flow = cur_flow[[1,0]] cur_flow_est = flow_i2r_est[h,w,:] if not config.network.STANDARD_FLOW_REP: cur_flow_est = cur_flow_est[[1,0]] point_color = [sigmoid(float(h)/height-0.5)*200+50, sigmoid(0.5-float(w)/width)*200+50, sigmoid(float(h)/height+float(w)/width-1)*200+50] mesh_src = cv2.circle(mesh_src, (np.round(w).astype(int), np.round(h).astype(int)), 1, point_color) mesh_tgt = cv2.circle(mesh_tgt, (np.round(w + cur_flow[0]).astype(int), np.round(h + cur_flow[1]).astype(int)), 1, point_color) mesh_tgt_est = cv2.circle(mesh_tgt_est, (np.round(w + cur_flow_est[0]).astype(int), np.round(h + cur_flow_est[1]).astype(int)), 1, point_color) weight = np.dstack((visible, visible)) ax_m_src = fig.add_subplot(3,4,5) ax_m_src.set_title('mesh_src') plt.axis('off') plt.imshow(mesh_src) ax_m_tgt = fig.add_subplot(3,4,6) ax_m_tgt.set_title('mesh_tgt_gt') plt.axis('off') plt.imshow(mesh_tgt) ax_m_tgt_est = fig.add_subplot(3,4,7) ax_m_tgt_est.set_title('mesh_tgt_est') plt.axis('off') sintel_vis = sintel_compute_color(np.vstack((flow[:,:,[1,0]], flow_i2r_est[:,:,[1,0]]*weight))) flow_i2r_vis = sintel_vis[:sintel_vis.shape[0]/2, :, :] flow_i2r_est_vis = sintel_vis[sintel_vis.shape[0]/2:, :, :] plt.imshow(mesh_tgt_est) ax_m_diff = fig.add_subplot(3,4,8) ax_m_diff.set_title('mesh_diff') plt.axis('off') mask_real = (mesh_tgt[:, :, 0] != 0).astype(np.float32) mask_real_est = (mesh_tgt_est[:, :, 0] != 0).astype(np.float32) mesh_diff = np.dstack((mask_real_est, mask_real, mask_real_est)) plt.imshow(mesh_diff) ax_flow_i2r = fig.add_subplot(3,4,9) ax_flow_i2r.set_title('flow_i2r') plt.imshow(flow_i2r_vis) plt.axis('off') ax_flow_i2r_est = fig.add_subplot(3,4,10) ax_flow_i2r_est.set_title('flow_i2r_est') plt.imshow(flow_i2r_est_vis) plt.axis('off') if not (visible_est is None): ax_visible_i2r_diff_est = fig.add_subplot(3, 4, 11) ax_visible_i2r_diff_est.set_title('visible_est_diff') plt.imshow(np.dstack((visible_est, visible, visible_est)).astype(np.float32)) plt.axis('off') ax_visible_i2r_est = fig.add_subplot(3, 4, 12) ax_visible_i2r_est.set_title('visible_est') plt.imshow(np.dstack((visible_est, visible_est, visible_est)).astype(np.float32)) plt.axis('off') ax_rendered_list = {'ax_image_rendered': ax_im_src, 'ax_depth_rendered': ax_d_src, 'ax_mesh_rendered': ax_m_src, 'ax_flow_i2r': ax_flow_i2r, 'ax_flow_i2r_est': ax_flow_i2r_est} ax_real_list = {'ax_image_real': ax_im_tgt, 'ax_depth_real': ax_d_tgt, 'ax_mesh_real': ax_m_tgt, 'ax_mesh_real_est': ax_m_tgt_est} data = {'flow_i2r': flow, 'visible': visible, 'flow_i2r_est': flow_i2r_est} cid = fig.canvas.mpl_connect('button_press_event', get_onclick(ax_real_list, ax_rendered_list, data, fig)) manager = plt.get_current_fig_manager() manager.window.showMaximized() # plt.show() return mesh_src, mesh_tgt, mesh_tgt_est
'../../data/render_v5/data/rendered/0006/{}_{}-depth.png'.format( img_idx[idx], sub_idx1[idx]), cv2.IMREAD_UNCHANGED) / 10000. v_depth_rendered[idx, 0] = cv2.imread( '../../data/render_v5/data/rendered/0006/{}_{}-depth.png'.format( img_idx[idx], sub_idx2[idx]), cv2.IMREAD_UNCHANGED) / 10000. pose_real_path = '../../data/render_v5/data/rendered/0006/{}_{}-pose.txt'.format( img_idx[idx], sub_idx1[idx]) pose_est_path = '../../data/render_v5/data/rendered/0006/{}_{}-pose.txt'.format( img_idx[idx], sub_idx2[idx]) pose_real = np.loadtxt(pose_real_path, skiprows=1) pose_est = np.loadtxt(pose_est_path, skiprows=1) flow, visible, _ = calc_flow(v_depth_rendered[idx, 0], pose_est, pose_real, K, v_depth_real[idx, 0], thresh=3E-3, standard_rep=False) v_zoom_factor[:, :2] = 1.5 v_zoom_factor[:, 2:] = 0.2 v_flow[idx] = flow.transpose((2, 0, 1)) v_flow_weights[idx] = np.tile(visible, (2, 1, 1)) zoom_flow_mx = mx.sym.Custom(name='updater_in', op_type='ZoomFlow', height=height, width=width, flow=flow_sym, flow_weights=flow_weights_sym, zoom_factor=zoom_factor,
def par_generate_gt(config, pair_rec, flow_depth_rendered=None): from lib.pair_matching.flow import calc_flow target_size, max_size = config.SCALES[0][0], config.SCALES[0][1] if flow_depth_rendered is None: flow_depth_rendered = cv2.imread( pair_rec["depth_rendered"], cv2.IMREAD_UNCHANGED ).astype(np.float32) flow_depth_rendered /= config.dataset.DEPTH_FACTOR flow_depth_rendered, _ = resize(flow_depth_rendered, target_size, max_size) if "depth_render_real" in pair_rec: flow_depth_real = cv2.imread( pair_rec["depth_render_real"], cv2.IMREAD_UNCHANGED ).astype(np.float32) flow_depth_real /= config.dataset.DEPTH_FACTOR else: print("not using render_real depth in par_generate_gt") flow_depth_real = cv2.imread( pair_rec["depth_real"], cv2.IMREAD_UNCHANGED ).astype(np.float32) flow_depth_real /= config.dataset.DEPTH_FACTOR flow_depth_real, _ = resize(flow_depth_real, target_size, max_size) if "mask_gt_observed" or "mask_observed" in pair_rec: mask_observed_path = pair_rec["mask_gt_observed"] assert os.path.exists(mask_observed_path), "%s does not exist".format( pair_rec["mask_gt_observed"] ) mask_observed = cv2.imread(mask_observed_path, cv2.IMREAD_UNCHANGED) mask_observed, _ = resize(mask_observed, target_size, max_size) flow_depth_real[mask_observed != pair_rec["mask_idx"]] = 0 if config.network.FLOW_I2R: if config.FLOW_CLASS_AGNOSTIC: flow_i2r, visible, _ = calc_flow( flow_depth_rendered, pair_rec["pose_rendered"], pair_rec["pose_observed"], config.dataset.INTRINSIC_MATRIX, flow_depth_real, standard_rep=config.network.STANDARD_FLOW_REP, ) flow_i2r_list = [ flow_i2r, visible, np.logical_and(visible == 0, flow_depth_rendered == 0), ] else: raise Exception("NOT_IMPLEMENTED") else: flow_i2r_list = None if config.network.FLOW_R2I: if config.FLOW_CLASS_AGNOSTIC: flow_r2i, visible, _ = calc_flow( flow_depth_real, pair_rec["pose_observed"], pair_rec["pose_rendered"], config.dataset.INTRINSIC_MATRIX, flow_depth_rendered, standard_rep=config.network.STANDARD_FLOW_REP, ) flow_r2i_list = [ flow_r2i, visible, np.logical_and(visible == 0, flow_depth_real == 0), ] else: raise Exception("NOT_IMPLEMENTED") else: flow_r2i_list = None return {"flow_i2r": flow_i2r_list, "flow_r2i": flow_r2i_list}