-
Notifications
You must be signed in to change notification settings - Fork 1
/
hopenet_utils.py
117 lines (99 loc) · 4.32 KB
/
hopenet_utils.py
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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
import numpy as np
import torch
#from torch.utils.serialization import load_lua
import os
import scipy.io as sio
import cv2
import math
from math import cos, sin
def softmax_temperature(tensor, temperature):
result = torch.exp(tensor / temperature)
result = torch.div(result, torch.sum(result, 1).unsqueeze(1).expand_as(result))
return result
def get_pose_params_from_mat(mat_path):
# This functions gets the pose parameters from the .mat
# Annotations that come with the Pose_300W_LP dataset.
mat = sio.loadmat(mat_path)
# [pitch yaw roll tdx tdy tdz scale_factor]
pre_pose_params = mat['Pose_Para'][0]
# Get [pitch, yaw, roll, tdx, tdy]
pose_params = pre_pose_params[:5]
return pose_params
def get_ypr_from_mat(mat_path):
# Get yaw, pitch, roll from .mat annotation.
# They are in radians
mat = sio.loadmat(mat_path)
# [pitch yaw roll tdx tdy tdz scale_factor]
pre_pose_params = mat['Pose_Para'][0]
# Get [pitch, yaw, roll]
pose_params = pre_pose_params[:3]
return pose_params
def get_pt2d_from_mat(mat_path):
# Get 2D landmarks
mat = sio.loadmat(mat_path)
pt2d = mat['pt2d']
return pt2d
def mse_loss(input, target):
return torch.sum(torch.abs(input.data - target.data) ** 2)
def plot_pose_cube(img, yaw, pitch, roll, tdx=None, tdy=None, size=150.):
# Input is a cv2 image
# pose_params: (pitch, yaw, roll, tdx, tdy)
# Where (tdx, tdy) is the translation of the face.
# For pose we have [pitch yaw roll tdx tdy tdz scale_factor]
p = pitch * np.pi / 180
y = -(yaw * np.pi / 180)
r = roll * np.pi / 180
if tdx != None and tdy != None:
face_x = tdx - 0.50 * size
face_y = tdy - 0.50 * size
else:
height, width = img.shape[:2]
face_x = width / 2 - 0.5 * size
face_y = height / 2 - 0.5 * size
x1 = size * (cos(y) * cos(r)) + face_x
y1 = size * (cos(p) * sin(r) + cos(r) * sin(p) * sin(y)) + face_y
x2 = size * (-cos(y) * sin(r)) + face_x
y2 = size * (cos(p) * cos(r) - sin(p) * sin(y) * sin(r)) + face_y
x3 = size * (sin(y)) + face_x
y3 = size * (-cos(y) * sin(p)) + face_y
# Draw base in red
cv2.line(img, (int(face_x), int(face_y)), (int(x1),int(y1)),(0,0,255),3)
cv2.line(img, (int(face_x), int(face_y)), (int(x2),int(y2)),(0,0,255),3)
cv2.line(img, (int(x2), int(y2)), (int(x2+x1-face_x),int(y2+y1-face_y)),(0,0,255),3)
cv2.line(img, (int(x1), int(y1)), (int(x1+x2-face_x),int(y1+y2-face_y)),(0,0,255),3)
# Draw pillars in blue
cv2.line(img, (int(face_x), int(face_y)), (int(x3),int(y3)),(255,0,0),2)
cv2.line(img, (int(x1), int(y1)), (int(x1+x3-face_x),int(y1+y3-face_y)),(255,0,0),2)
cv2.line(img, (int(x2), int(y2)), (int(x2+x3-face_x),int(y2+y3-face_y)),(255,0,0),2)
cv2.line(img, (int(x2+x1-face_x),int(y2+y1-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(255,0,0),2)
# Draw top in green
cv2.line(img, (int(x3+x1-face_x),int(y3+y1-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(0,255,0),2)
cv2.line(img, (int(x2+x3-face_x),int(y2+y3-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(0,255,0),2)
cv2.line(img, (int(x3), int(y3)), (int(x3+x1-face_x),int(y3+y1-face_y)),(0,255,0),2)
cv2.line(img, (int(x3), int(y3)), (int(x3+x2-face_x),int(y3+y2-face_y)),(0,255,0),2)
return img
def draw_axis(img, yaw, pitch, roll, tdx=None, tdy=None, size = 100):
pitch = pitch * np.pi / 180
yaw = -(yaw * np.pi / 180)
roll = roll * np.pi / 180
if tdx != None and tdy != None:
tdx = tdx
tdy = tdy
else:
height, width = img.shape[:2]
tdx = width / 2
tdy = height / 2
# X-Axis pointing to right. drawn in red
x1 = size * (cos(yaw) * cos(roll)) + tdx
y1 = size * (cos(pitch) * sin(roll) + cos(roll) * sin(pitch) * sin(yaw)) + tdy
# Y-Axis | drawn in green
# v
x2 = size * (-cos(yaw) * sin(roll)) + tdx
y2 = size * (cos(pitch) * cos(roll) - sin(pitch) * sin(yaw) * sin(roll)) + tdy
# Z-Axis (out of the screen) drawn in blue
x3 = size * (sin(yaw)) + tdx
y3 = size * (-cos(yaw) * sin(pitch)) + tdy
cv2.line(img, (int(tdx), int(tdy)), (int(x1),int(y1)),(0,0,255),3)
cv2.line(img, (int(tdx), int(tdy)), (int(x2),int(y2)),(0,255,0),3)
cv2.line(img, (int(tdx), int(tdy)), (int(x3),int(y3)),(255,0,0),2)
return img