add front, depth, uv texture, uv mask(visibility)

demo.py

@@ -3,30 +3,26 @@
 from glob import glob
 import scipy.io as sio
 from skimage.io import imread, imsave
+from skimage.transform import rescale, resize
 from time import time
 import argparse
 import ast
 
 from api import PRN
 
+from utils.estimate_pose import estimate_pose
+from utils.rotate_vertices import frontalize
+from utils.render_app import get_visibility, get_uv_mask, get_depth_image
+from utils.write import write_obj, write_obj_with_texture
 
 def main(args):
-    print args.isDlib
-
-    if args.isShow:
-        args.isOpencv = True
+    if args.isShow or args.isTexture:
         import cv2
         from utils.cv_plot import plot_kpt, plot_vertices, plot_pose_box
-        from utils.write import write_obj
-        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) 
+    prn = PRN(is_dlib = args.isDlib) 
 
     # ------------- load data
     image_folder = args.inputDir
@@ -46,25 +42,57 @@ def main(args):
 
         # read image
         image = imread(image_path)
+        [h, w, _] = image.shape
 
         # the core: regress position map    
         if args.isDlib:
+            max_size = max(image.shape[0], image.shape[1]) 
+            if max_size> 1000:
+                image = rescale(image, 1000./max_size)
             pos = prn.process(image) # use dlib to detect face
         else:
-            if image.shape[1] == 256:
+            if image.shape[1] == image.shape[2]:
+                image = resize(image, (256,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()
+                box = np.array([0, image.shape[1]-1, 0, image.shape[0]-1]) # cropped with bounding box
+                pos = prn.process(image, box)
+
+        image = image/255.
         if pos is None:
             continue
 
-        if args.is3d or args.isShow:        
+        if args.is3d or args.isMat or args.isPose or args.isShow:        
             # 3D vertices
             vertices = prn.get_vertices(pos)
+            if args.isFront:
+                save_vertices = frontalize(vertices)
+            else:
+                save_vertices = vertices
+
+        if args.isImage:
+            imsave(os.path.join(save_folder, name + '.jpg'), image) 
+
+        if args.is3d:
             # corresponding colors
             colors = prn.get_colors(image, vertices)
-            write_obj(os.path.join(save_folder, name + '.obj'), vertices, colors, prn.triangles) #save 3d face(can open with meshlab)
+
+            if args.isTexture:
+                texture = cv2.remap(image, pos[:,:,:2].astype(np.float32), None, interpolation=cv2.INTER_NEAREST, borderMode=cv2.BORDER_CONSTANT,borderValue=(0))
+                if args.isMask:
+                    vertices_vis = get_visibility(vertices, prn.triangles, h, w)
+                    uv_mask = get_uv_mask(vertices_vis, prn.triangles, prn.uv_coords, h, w, prn.resolution_op)
+                    texture = texture*uv_mask[:,:,np.newaxis]
+                write_obj_with_texture(os.path.join(save_folder, name + '.obj'), save_vertices, colors, prn.triangles, texture, prn.uv_coords/prn.resolution_op)#save 3d face with texture(can open with meshlab)
+            else:
+                write_obj(os.path.join(save_folder, name + '.obj'), save_vertices, colors, prn.triangles) #save 3d face(can open with meshlab)
+
+        if args.isDepth:
+            depth_image = get_depth_image(vertices, prn.triangles, h, w) 
+            imsave(os.path.join(save_folder, name + '_depth.jpg'), depth_image) 
+
+        if args.isMat:
+            sio.savemat(os.path.join(save_folder, name + '_mesh.mat'), {'vertices': save_vertices, 'colors': colors, 'triangles': prn.triangles})
 
         if args.isKpt or args.isShow:
             # get landmarks
@@ -96,15 +124,28 @@ def main(args):
                         help='set gpu id, -1 for CPU')
     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('--is3d', default=True, type=ast.literal_eval, 
                         help='whether to output 3D face(.obj)')
+    parser.add_argument('--isMat', default=False, type=ast.literal_eval, 
+                        help='whether to save vertices,color,triangles as mat for matlab showing')
     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)')
+    parser.add_argument('--isImage', default=False, type=ast.literal_eval, 
+                        help='whether to save input image')
+    # update in 2017/4/10
+    parser.add_argument('--isFront', default=False, type=ast.literal_eval, 
+                        help='whether to frontalize vertices(mesh)')
+    # update in 2017/4/25
+    parser.add_argument('--isDepth', default=False, type=ast.literal_eval, 
+                        help='whether to output depth image')
+    # update in 2017/4/27
+    parser.add_argument('--isTexture', default=False, type=ast.literal_eval, 
+                        help='whether to save texture in obj file')
+    parser.add_argument('--isMask', default=False, type=ast.literal_eval, 
+                        help='whether to set invisible pixels(due to self-occlusion) in texture as 0')
 
     main(parser.parse_args())