30 lines
1.0 KiB
Python
30 lines
1.0 KiB
Python
import numpy as np
|
|
import cv2 # pylint: disable=import-error
|
|
|
|
def rot_matrix(roll, pitch, yaw):
|
|
cr, sr = np.cos(roll), np.sin(roll)
|
|
cp, sp = np.cos(pitch), np.sin(pitch)
|
|
cy, sy = np.cos(yaw), np.sin(yaw)
|
|
rr = np.array([[1, 0, 0], [0, cr, -sr], [0, sr, cr]])
|
|
rp = np.array([[cp, 0, sp], [0, 1, 0], [-sp, 0, cp]])
|
|
ry = np.array([[cy, -sy, 0], [sy, cy, 0], [0, 0, 1]])
|
|
return ry.dot(rp.dot(rr))
|
|
|
|
def draw_pose(img, pose, loc, W=160, H=320, xyoffset=(0, 0), faceprob=0):
|
|
rcmat = np.zeros((3, 4))
|
|
rcmat[:, :3] = rot_matrix(*pose[0:3]) * 0.5
|
|
rcmat[0, 3] = (loc[0]+0.5) * W
|
|
rcmat[1, 3] = (loc[1]+0.5) * H
|
|
rcmat[2, 3] = 1.0
|
|
# draw nose
|
|
p1 = np.dot(rcmat, [0, 0, 0, 1])[0:2]
|
|
p2 = np.dot(rcmat, [0, 0, 100, 1])[0:2]
|
|
tr = tuple([int(round(x + xyoffset[i])) for i, x in enumerate(p1)])
|
|
pr = tuple([int(round(x + xyoffset[i])) for i, x in enumerate(p2)])
|
|
if faceprob > 0.4:
|
|
color = (255, 255, 0)
|
|
cv2.line(img, tr, pr, color=(255, 255, 0), thickness=3)
|
|
else:
|
|
color = (64, 64, 64)
|
|
cv2.circle(img, tr, 7, color=color)
|