diff --git a/demo.py b/demo.py new file mode 100644 index 0000000..3870552 --- /dev/null +++ b/demo.py @@ -0,0 +1,158 @@ +import numpy as np +import os +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): + if args.isShow or args.isTexture: + import cv2 + from utils.cv_plot import plot_kpt, plot_vertices, plot_pose_box + + # ---- init PRN + os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu # GPU number, -1 for CPU + prn = PRN(is_dlib = args.isDlib) + + # ------------- load data + image_folder = args.inputDir + save_folder = args.outputDir + if not os.path.exists(save_folder): + os.mkdir(save_folder) + + types = ('*.jpg', '*.png') + image_path_list= [] + for files in types: + image_path_list.extend(glob(os.path.join(image_folder, files))) + total_num = len(image_path_list) + + for i, image_path in enumerate(image_path_list): + + name = image_path.strip().split('/')[-1][:-4] + + # 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) + image = (image*255).astype(np.uint8) + pos = prn.process(image) # use dlib to detect face + else: + 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: + 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.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.copy() + save_vertices[:,1] = h - 1 - save_vertices[:,1] + + if args.isImage: + imsave(os.path.join(save_folder, name + '.jpg'), image) + + if args.is3d: + # corresponding colors + colors = prn.get_colors(image, vertices) + + 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, True) + depth = get_depth_image(vertices, prn.triangles, h, w) + imsave(os.path.join(save_folder, name + '_depth.jpg'), depth_image) + sio.savemat(os.path.join(save_folder, name + '_depth.mat'), {'depth':depth}) + + if args.isMat: + sio.savemat(os.path.join(save_folder, name + '_mesh.mat'), {'vertices': vertices, 'colors': colors, 'triangles': prn.triangles}) + + if args.isKpt or args.isShow: + # get landmarks + kpt = prn.get_landmarks(pos) + np.savetxt(os.path.join(save_folder, name + '_kpt.txt'), kpt) + + if args.isPose or args.isShow: + # estimate pose + camera_matrix, pose = estimate_pose(vertices) + np.savetxt(os.path.join(save_folder, name + '_pose.txt'), pose) + np.savetxt(os.path.join(save_folder, name + '_camera_matrix.txt'), camera_matrix) + + np.savetxt(os.path.join(save_folder, name + '_pose.txt'), pose) + + if args.isShow: + # ---------- Plot + 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) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Joint 3D Face Reconstruction and Dense Alignment with Position Map Regression Network') + + 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=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('--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())