From 1a3aa424e79d2e73eebdf0d503658588ced08305 Mon Sep 17 00:00:00 2001
From: Nataniel Ruiz <nruiz9@gatech.edu>
Date: 星期日, 25 二月 2018 07:55:31 +0800
Subject: [PATCH] Update train_hopenet.py
---
code/utils.py | 131 ++++++++++++++++++++++++++++---------------
1 files changed, 86 insertions(+), 45 deletions(-)
diff --git a/code/utils.py b/code/utils.py
index abba4f3..b75d51f 100644
--- a/code/utils.py
+++ b/code/utils.py
@@ -7,54 +7,14 @@
import math
from math import cos, sin
-def plot_pose_cube(img, pose_params, size=150.):
- # 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]
-
- pitch, yaw, roll, tdx, tdy = pose_params
-
- p = pitch
- y = -yaw
- r = roll
-
- # TODO: probably useless
- # Translation of box depends on head pointing to right or left.
- if y < 0:
- face_x = tdx - 0.50 * size
- face_y = tdy - 0.50 * size
- else:
- face_x = tdx - 0.50 * size
- face_y = tdy - 0.50 * 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)),(255,0,0),3)
- cv2.line(img, (int(face_x), int(face_y)), (int(x2),int(y2)),(255,0,0),3)
- cv2.line(img, (int(x2), int(y2)), (int(x2+x1-face_x),int(y2+y1-face_y)),(255,0,0),3)
- cv2.line(img, (int(x1), int(y1)), (int(x1+x2-face_x),int(y1+y2-face_y)),(255,0,0),3)
- # Draw pillars in blue
- cv2.line(img, (int(face_x), int(face_y)), (int(x3),int(y3)),(0,0,255),2)
- cv2.line(img, (int(x1), int(y1)), (int(x1+x3-face_x),int(y1+y3-face_y)),(0,0,255),2)
- cv2.line(img, (int(x2), int(y2)), (int(x2+x3-face_x),int(y2+y3-face_y)),(0,0,255),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)),(0,0,255),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 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 300W_LP dataset.
+ # 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]
@@ -72,5 +32,86 @@
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
--
Gitblit v1.8.0