add demo

demo.py

@@ -4,29 +4,33 @@
 import scipy.io as sio
 from skimage.io import imread, imsave
 from time import time
+import argparse
+import ast
 
 from api import PRN
-import argparse
 
 
 def main(args):
+    print args.isDlib
+
     if args.isShow:
         args.isOpencv = True
+        import cv2
         from utils.cv_plot import plot_kpt, plot_vertices, plot_pose_box
-    if args.isObj:
         from utils.write import write_obj
-    if args.isPose:
+        from utils.estimate_pose import estimate_pose
+    elif args.is3d:
+        from utils.write import write_obj
+    elif args.isPose:
         from utils.estimate_pose import estimate_pose
 
     # ---- init PRN
     os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu # GPU number, -1 for CPU
-
     prn = PRN(is_dlib = args.isDlib, is_opencv = args.isOpencv) 
 
-
     # ------------- load data
-    image_folder = args.inputFolder
-    save_folder = args.outputFolder
+    image_folder = args.inputDir
+    save_folder = args.outputDir
     if not os.path.exists(save_folder):
         os.mkdir(save_folder)
 
@@ -44,10 +48,18 @@ def main(args):
         image = imread(image_path)
 
         # the core: regress position map    
-        pos = prn.process(image) # use dlib to detect face
-
-
-        if args.isObj or args.isShow:        
+        if args.isDlib:
+            pos = prn.process(image) # use dlib to detect face
+        else:
+            if image.shape[1] == 256:
+                pos = prn.net_forward(image/255.) # input image has been cropped to 256x256
+            else:
+                print('please make sure the image has been cropped')
+                exit()
+        if pos is None:
+            continue
+
+        if args.is3d or args.isShow:        
             # 3D vertices
             vertices = prn.get_vertices(pos)
             # corresponding colors
@@ -76,22 +88,23 @@ def main(args):
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='Joint 3D Face Reconstruction and Dense Alignment with Position Map Regression Network')
 
-    parser.add_argument('-i', '--inputFolder', default='TestImages/', type=str,
-                        help='input images folder path')
-    parser.add_argument('-o', '--outputFolder', default='TestImages/results', type=str, 
-                        help='output obj file')
+    parser.add_argument('-i', '--inputDir', default='TestImages/', type=str,
+                        help='path to the input directory, where input images are stored.')
+    parser.add_argument('-o', '--outputDir', default='TestImages/results', type=str, 
+                        help='path to the output directory, where results(obj,txt files) will be stored.')
     parser.add_argument('--gpu', default='0', type=str, 
                         help='set gpu id, -1 for CPU')
-    parser.add_argument('--isDlib', default=True, type=bool, 
-                        help='whether to use dlib to detect face')
-    parser.add_argument('--isOpencv', default=True, type=bool, 
+    parser.add_argument('--isDlib', default=True, type=ast.literal_eval, 
+                        help='whether to use dlib for detecting face, default is True, if False, the input image should be cropped in advance')
+    parser.add_argument('--isOpencv', default=False, type=ast.literal_eval, 
                         help='whether to use opencv')
-    parser.add_argument('--isObj', default=True, type=bool, 
-                        help='whether to output .obj file')
-    parser.add_argument('--isKpt', default=False, type=bool, 
-                        help='whether to output key points')
-    parser.add_argument('--isPose', default=False, type=bool, 
-                        help='whether to output estimated pose')
-    parser.add_argument('--isShow', default=False, type=bool, 
-                        help='whether to show the results with opencv(need set isOpencv=True')
+    parser.add_argument('--is3d', default=True, type=ast.literal_eval, 
+                        help='whether to output 3D face(.obj)')
+    parser.add_argument('--isKpt', default=False, type=ast.literal_eval,  
+                        help='whether to output key points(.txt)')
+    parser.add_argument('--isPose', default=False, type=ast.literal_eval,  
+                        help='whether to output estimated pose(.txt)')
+    parser.add_argument('--isShow', default=False, type=ast.literal_eval,  
+                        help='whether to show the results with opencv(need opencv)')
+
     main(parser.parse_args())