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