From af51d0ecb51ad4d6c8ed086855bd3c411ebc4aa0 Mon Sep 17 00:00:00 2001
From: natanielruiz <nruiz9@gatech.edu>
Date: 星期一, 30 十月 2017 06:29:51 +0800
Subject: [PATCH] Fixed stuff

---
 code/utils.py |   89 +++++++++++++++++++++++++++++++++-----------
 1 files changed, 66 insertions(+), 23 deletions(-)

diff --git a/code/utils.py b/code/utils.py
index 645ae19..b75d51f 100644
--- a/code/utils.py
+++ b/code/utils.py
@@ -7,6 +7,40 @@
 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)
@@ -16,14 +50,13 @@
     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.15 - size
-        face_y = height / 2 - 0.15 - size
+        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
@@ -50,25 +83,35 @@
 
     return img
 
-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.
-    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 draw_axis(img, yaw, pitch, roll, tdx=None, tdy=None, size = 100):
 
-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
+    pitch = pitch * np.pi / 180
+    yaw = -(yaw * np.pi / 180)
+    roll = roll * np.pi / 180
 
-def mse_loss(input, target):
-    return torch.sum(torch.abs(input.data - target.data) ** 2)
+    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