def icp_refine(self, pose_pred, anno, output, K): depth = read_depth(anno['depth_path']) mask = torch.argmax(output['seg'], dim=1)[0].detach().cpu().numpy() if pose_pred[2, 3] <= 0 or np.sum(mask) < 20: return pose_pred depth[mask != 1] = 0 pose_pred_tmp = pose_pred.copy() pose_pred_tmp[:3, 3] = pose_pred_tmp[:3, 3] * 1000 R_refined, t_refined = icp_utils.icp_refinement( depth, self.icp_render, pose_pred_tmp[:3, :3], pose_pred_tmp[:3, 3], K.copy(), (depth.shape[1], depth.shape[0]), depth_only=True, max_mean_dist_factor=5.0) R_refined, _ = icp_utils.icp_refinement( depth, self.icp_render, R_refined, t_refined, K.copy(), (depth.shape[1], depth.shape[0]), no_depth=True) pose_pred = np.hstack((R_refined, t_refined.reshape((3, 1)) / 1000)) return pose_pred
def icp_refine_(self, pose, anno, output): depth = read_depth(anno['depth_path']).astype(np.uint16) mask = torch.argmax(output['seg'], dim=1)[0].detach().cpu().numpy() mask = mask.astype(np.int32) pose = pose.astype(np.float32) poses = np.zeros([1, 7], dtype=np.float32) poses[0, :4] = mat2quat(pose[:, :3]) poses[0, 4:] = pose[:, 3] poses_new = np.zeros([1, 7], dtype=np.float32) poses_icp = np.zeros([1, 7], dtype=np.float32) fx = 572.41140 fy = 573.57043 px = 325.26110 py = 242.04899 zfar = 6.0 znear = 0.25 factor = 1000.0 error_threshold = 0.01 rois = np.zeros([1, 6], dtype=np.float32) rois[:, :] = 1 self.icp_refiner.solveICP(mask, depth, self.height, self.width, fx, fy, px, py, znear, zfar, factor, rois.shape[0], rois, poses, poses_new, poses_icp, error_threshold) pose_icp = np.zeros([3, 4], dtype=np.float32) pose_icp[:, :3] = quat2mat(poses_icp[0, :4]) pose_icp[:, 3] = poses_icp[0, 4:] return pose_icp