diff --git a/run.py b/run.py index 188c6a0..253af4a 100644 --- a/run.py +++ b/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)