SolvePnP for Head Pose Estimation

pnp.py

"""
Light weight head pose estimation with SolvePnP

Author: Yuanjun Xiong
"""

# parameters
fx = 1

# model points
model_points_xy = [
    [-0.0949517,  -0.0844475],      # left eye
    [0.0949518, -0.0844475],        # right eye
    [0, 0.0421711],                 # nose
    [6.20882e-10, 0.133588],        # mouth
]

model_points_z = [0, 0, 15, 5]
model_points = np.empty((4, 3), np.float64)
model_points[:, :2] = model_points_xy
model_points[:, 2] = model_points_z
model_points[:, 0] *= 150
model_points[:, 1] *= 150


def head_pose_estimate(image_size, landmarks):
    h, w = image_size
    K = np.float64(
                [[fx * w,   0,      0.5*(w-1)],
                [0,         fx * h, 0.5*(h-1)],
                [0.0,       0.0,    1.0]])
    dist_coef = np.zeros((4, 1))
    ret, rvec, tvec = cv2.solvePnP(model_points, landmarks, K, dist_coef, flags=cv2.SOLVEPNP_ITERATIVE)

    rot_mat = cv2.Rodrigues(rvec)[0]
    P = np.hstack((rot_mat, np.zeros((3, 1), dtype=np.float64)))
    eulerAngles =  cv2.decomposeProjectionMatrix(P)[6]
    yaw   = eulerAngles[1, 0]
    pitch = eulerAngles[0, 0]
    roll  = eulerAngles[2, 0]
    return roll, yaw, pitch