diff --git a/demo.py b/demo.py index 0e07ca4..60ac52e 100644 --- a/demo.py +++ b/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()) \ No newline at end of file