def regress_pose(self, regressor, predictions, pr_para, pi_para, K_inv, pts2d_pred_loc, pts2d_pred_var, graph_pred, sym_cor_pred, mask_pred): if mask_pred.sum() == 0: # object is not detected R = np.eye(3, dtype=np.float32) t = np.zeros((3, 1), dtype=np.float32) return R, t, R, t self.fill_intermediate_predictions(regressor, predictions, K_inv, pts2d_pred_loc, pts2d_pred_var, graph_pred, sym_cor_pred, mask_pred) # initialize pose predictions = regressor.initialize_pose(predictions, pi_para, self.args.use_keypoint, self.args.use_edge, self.args.use_symmetry) pose_init = np.zeros((4, 3), dtype=np.float32) regressor.get_pose(predictions, get_2d_ctypes(pose_init)) R_init = pose_init[1:].transpose() t_init = pose_init[0].reshape((3, 1)) # refine pose predictions = regressor.refine_pose(predictions, pr_para, self.args.use_keypoint, self.args.use_edge, self.args.use_symmetry) pose_final = np.zeros((4, 3), dtype=np.float32) regressor.get_pose(predictions, get_2d_ctypes(pose_final)) R_final = pose_final[1:].transpose() t_final = pose_final[0].reshape((3, 1)) return R_final, t_final, R_init, t_init
def search_para(self, regressor, predictions_para, poses_para, K_inv, normal_gt, diameter, val_set): para_id = 0 for data_id in range(len(val_set['pts3d'])): if val_set['mask_pred'][data_id].sum() == 0 or \ np.sum(val_set['pts2d_pred_loc'][data_id]) == 0: # object not detected continue predictions = regressor.get_prediction_container( predictions_para, para_id) # fill intermediate predictions self.fill_intermediate_predictions( regressor, predictions, K_inv, val_set['pts2d_pred_loc'][data_id], val_set['pts2d_pred_var'][data_id], val_set['graph_pred'][data_id], val_set['sym_cor_pred'][data_id], val_set['mask_pred'][data_id]) # fill ground-truth poses pose_gt = np.zeros((4, 3), dtype=np.float32) tvec = val_set['t_gt'][data_id] r = val_set['R_gt'][data_id] pose_gt[0] = tvec.transpose()[0] pose_gt[1:] = r.transpose() regressor.set_pose_gt(poses_para, para_id, get_2d_ctypes(pose_gt)) # increment number of valid examples in the val set para_id += 1 # search parameter # para_id is datasize for parameter search pi_para = regressor.search_pose_initial(predictions_para, poses_para, para_id, diameter) pr_para = regressor.search_pose_refine(predictions_para, poses_para, para_id, diameter) return pr_para, pi_para
def fill_intermediate_predictions(self, regressor, predictions, K_inv, pts2d_pred_loc, pts2d_pred_var, graph_pred, sym_cor_pred, mask_pred): # load intermediate representations to regressor n_keypts = self.args.num_keypoints n_edges = n_keypts * (n_keypts - 1) // 2 # point3D_gt #regressor.set_point3D_gt(predictions, get_2d_ctypes(pts3d), n_keypts) # point2D_pred point2D_pred = np.matrix(np.ones((3, n_keypts), dtype=np.float32)) point2D_pred[:2] = pts2d_pred_loc.transpose() point2D_pred = np.array((K_inv * point2D_pred)[:2]).transpose() regressor.set_point2D_pred(predictions, get_2d_ctypes(point2D_pred), n_keypts) # point_inv_half_var point_inv_half_var = np.zeros((n_keypts, 2, 2), dtype=np.float32) for i in range(n_keypts): # compute cov^{-1/2} cov = np.matrix(pts2d_pred_var[i]) cov = (cov + cov.transpose() ) / 2 # ensure the covariance matrix is symmetric v, u = np.linalg.eig(cov) v = np.matrix(np.diag(1. / np.sqrt(v))) point_inv_half_var[i] = u * v * u.transpose() point_inv_half_var = point_inv_half_var.reshape((n_keypts, 4)) regressor.set_point_inv_half_var(predictions, get_2d_ctypes(point_inv_half_var), n_keypts) # normal_gt #regressor.set_normal_gt(predictions, normal_gt.ctypes) # vec_pred and edge_inv_half_var graph_pred = graph_pred.reshape( (n_edges, 2, graph_pred.shape[1], graph_pred.shape[2])) vec_pred = np.zeros((n_edges, 2), dtype=np.float32) edge_inv_half_var = np.zeros((n_edges, 2, 2), dtype=np.float32) for i in range(n_edges): xs = graph_pred[i, 0][mask_pred == 1.] ys = graph_pred[i, 1][mask_pred == 1.] vec_pred[i] = [xs.mean(), ys.mean()] if self.args.dataset == 'linemod': cov = np.cov(xs, ys) cov = (cov + cov.transpose() ) / 2 # ensure the covariance matrix is symmetric v, u = np.linalg.eig(cov) v = np.matrix(np.diag(1. / np.sqrt(v))) edge_inv_half_var[i] = u * v * u.transpose() elif self.args.dataset == 'occlusion_linemod': edge_inv_half_var[i] = np.eye(2) else: # dataset not supported pdb.set_trace() vec_pred = np.array(K_inv[:2, :2] * np.matrix(vec_pred).transpose()).transpose() edge_inv_half_var = edge_inv_half_var.reshape((n_edges, 4)) regressor.set_vec_pred(predictions, get_2d_ctypes(vec_pred), n_edges) regressor.set_edge_inv_half_var(predictions, get_2d_ctypes(edge_inv_half_var), n_edges) # qs1_cross_qs2 and symmetry weight sym_cor_pred = self.flatten_sym_cor(sym_cor_pred, mask_pred) qs1_cross_qs2_all = np.zeros((sym_cor_pred.shape[0], 3), dtype=np.float32) for i in range(sym_cor_pred.shape[0]): qs1 = np.ones((3, ), dtype=np.float32) qs2 = np.ones((3, ), dtype=np.float32) qs1[:2] = sym_cor_pred[i][0] qs2[:2] = sym_cor_pred[i][1] qs1 = np.array(K_inv * np.matrix(qs1).transpose()).transpose()[0] qs2 = np.array(K_inv * np.matrix(qs2).transpose()).transpose()[0] qs1_cross_qs2_all[i] = np.cross(qs1, qs2) qs1_cross_qs2_filtered, symmetry_weight = self.filter_symmetry( qs1_cross_qs2_all) n_symmetry = qs1_cross_qs2_filtered.shape[0] regressor.set_qs1_cross_qs2(predictions, get_2d_ctypes(qs1_cross_qs2_filtered), n_symmetry) regressor.set_symmetry_weight(predictions, symmetry_weight.ctypes, n_symmetry)