30 lines
990 B
Python
30 lines
990 B
Python
import numpy as np
|
|
import cv2
|
|
|
|
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)
|
|
|