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