Spaces:
No application file
No application file
File size: 1,194 Bytes
ddec2b7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | import cv2
import numpy as np
MODEL_POINTS = np.array([
(0.0, 0.0, 0.0),
(0.0, -330.0, -65.0),
(-225.0, 170.0, -135.0),
(225.0, 170.0, -135.0),
(-150.0, -150.0, -125.0),
(150.0, -150.0, -125.0)
], dtype=np.float64)
LANDMARK_INDICES = [1, 152, 33, 263, 61, 291]
def calculate_head_pose(frame, landmarks):
h, w = frame.shape[:2]
focal = w
center = (w/2, h/2)
cam_matrix = np.array([
[focal, 0, center[0]],
[0, focal, center[1]],
[0, 0, 1]
], dtype=np.float64)
dist_coeffs = np.zeros((4, 1))
points_2d = []
for idx in LANDMARK_INDICES:
lm = landmarks[idx]
points_2d.append([lm.x * w, lm.y * h])
points_2d = np.array(points_2d, dtype=np.float64)
success, rot_vec, trans_vec = cv2.solvePnP(
MODEL_POINTS, points_2d, cam_matrix, dist_coeffs
)
if not success:
return 0, 0, 0
rot_mat, _ = cv2.Rodrigues(rot_vec)
pose_mat = cv2.hconcat((rot_mat, trans_vec))
_, _, _, _, _, _, euler = cv2.decomposeProjectionMatrix(pose_mat)
pitch = euler[0][0]
yaw = euler[1][0]
roll = euler[2][0]
return pitch, yaw, roll
|