add pose, plot

run.py

@@ -7,15 +7,17 @@
 
 from api import PRN
 from utils.write import write_obj
+from utils.estimate_pose import estimate_pose
+from utils.cv_plot import *
 
 # ---- init PRN
 os.environ['CUDA_VISIBLE_DEVICES'] = '0' # GPU number, -1 for CPU
-prn = PRN(is_dlib = True, is_opencv = False) 
+prn = PRN(is_dlib = True, is_opencv = True) 
 
 
 # ------------- load data
 image_folder = 'TestImages/'
-save_folder = 'TestImages/'
+save_folder = 'TestImages/results/'
 if not os.path.exists(save_folder):
     os.mkdir(save_folder)
 
@@ -40,7 +42,22 @@
     # corresponding colors
     colors = prn.get_colors(image, vertices)
 
-    # -- save
+    # -- More
+    # estimate pose
+    camera_matrix, pose = estimate_pose(vertices)
+
+
+    # ---------- Plot
+    print vertices.shape
+    image_pose = plot_pose_box(image, camera_matrix, kpt)
+    cv2.imshow('sparse alignment', plot_kpt(image, kpt))
+    cv2.imshow('dense alignment', plot_vertices(image, vertices))
+    cv2.imshow('pose', plot_pose_box(image, camera_matrix, kpt))
+    cv2.waitKey(0)
+
+
+    # ---- Save
     name = image_path.strip().split('/')[-1][:-4]
-    np.savetxt(os.path.join(save_folder, name + '.txt'), kpt) 
+    np.savetxt(os.path.join(save_folder, name + '_pose.txt'), pose) 
+    np.savetxt(os.path.join(save_folder, name + '_kpt.txt'), kpt) 
     write_obj(os.path.join(save_folder, name + '.obj'), vertices, colors, prn.triangles) #save 3d face(can open with meshlab)