def prepare_test_data_RGB(img_dir, out_dir): mtcnn_dir = os.path.join(out_dir, "mtcnn_result/") lmk3D_ori_txt_path = os.path.join(out_dir, "lmk_3D_86pts_ori.txt") lmk2D_ori_txt_path = os.path.join(out_dir, "lmk_2D_68pts_ori.txt") lmk3D_crop_txt_path = os.path.join(out_dir, "lmk_3D_86pts.txt") lmk2D_crop_txt_path = os.path.join(out_dir, "lmk_2D_68pts.txt") crop_dir = out_dir seg_dir = out_dir print("start MTCNN") pb_path = os.path.join(FLAGS.pb_path, "mtcnn_model.pb") names_list = detect_face_with_mtcnn.detect_with_MTCNN(img_dir, mtcnn_dir, pb_path) print("start detect 86pt 3D lmk") tf.reset_default_graph() pb_path = os.path.join(FLAGS.pb_path, "lmk3D_86_model.pb") detect_3D_landmark.detect_lmk86( img_dir, mtcnn_dir, lmk3D_ori_txt_path, names_list, pb_path ) print("start detect 68pt 2D lmk") # need to transfer RGB in the function tf.reset_default_graph() pb_path = os.path.join(FLAGS.pb_path, "lmk2D_68_model.pb") detect_2Dlmk_all_imgs( pb_path, img_dir, lmk3D_ori_txt_path, lmk2D_ori_txt_path ) # N x 68 x 2 print("start crop by 3D lmk") crop_img_by_lmk( lmk2D_ori_txt_path, lmk3D_ori_txt_path, lmk2D_crop_txt_path, lmk3D_crop_txt_path, img_dir, crop_dir, orig=True, ) print("start face seg") # need to transfer RGB in the function tf.reset_default_graph() pb_path = os.path.join(FLAGS.pb_path, "faceseg_model.pb") face_seg(pb_path, lmk3D_crop_txt_path, crop_dir, seg_dir) print("finish RGB data preparation")
def prepare_test_data_RGBD(img_dir, out_dir): mtcnn_dir = os.path.join(out_dir, "mtcnn_result/") lmk3D_ori_txt_path = os.path.join(out_dir, "lmk_3D_86pts_ori.txt") lmk2D_ori_txt_path = os.path.join(out_dir, "lmk_2D_68pts_ori.txt") lmk3D_crop_txt_path = os.path.join(out_dir, "lmk_3D_86pts.txt") lmk2D_crop_txt_path = os.path.join(out_dir, "lmk_2D_68pts.txt") out_img_names_path = os.path.join(out_dir, "img_names.txt") crop_dir = out_dir seg_dir = out_dir print("check RGBD data") files = glob.glob(osp.join(img_dir, "*.jpg")) files.extend(glob.glob(osp.join(img_dir, "*.JPG"))) dep_files = glob.glob(osp.join(img_dir, "*.png")) dep_files.extend(glob.glob(osp.join(img_dir, "*.PNG"))) files.sort() dep_files.sort() if len(files) == len(dep_files) is False: raise Exception("The number of depth img is not same with RGB img.") for index in range(0, len(files)): img = cv2.imread(files[index]) dep_img = cv2.imread(dep_files[index]) if img.shape[0] < img.shape[1]: img = cv2.transpose(img) dep_img = cv2.transpose(dep_img) cv2.imwrite(img_dir + files[index].split("/")[-1], img) cv2.imwrite(img_dir + dep_files[index].split("/")[-1], dep_img) print("start MTCNN") pb_path = os.path.join(FLAGS.pb_path, "mtcnn_model.pb") names_list, dep_name_list = detect_face_with_mtcnn.detect_with_MTCNN( img_dir, mtcnn_dir, pb_path, "depth" ) print("start detect 86pt 3D lmk") tf.reset_default_graph() pb_path = os.path.join(FLAGS.pb_path, "lmk3D_86_model.pb") detect_3D_landmark.detect_lmk86( img_dir, mtcnn_dir, lmk3D_ori_txt_path, names_list, pb_path ) print("start detect 68pt 2D lmk") # need to transfer RGB in the function tf.reset_default_graph() pb_path = os.path.join(FLAGS.pb_path, "lmk2D_68_model.pb") detect_2Dlmk_all_imgs( pb_path, img_dir, lmk3D_ori_txt_path, lmk2D_ori_txt_path ) # N x 68 x 2 print("start crop by 3D lmk") crop_img_by_lmk( lmk2D_ori_txt_path, lmk3D_ori_txt_path, lmk2D_crop_txt_path, lmk3D_crop_txt_path, img_dir, crop_dir, orig=False, ) print("start face seg") # need to transfer RGB in the function tf.reset_default_graph() pb_path = os.path.join(FLAGS.pb_path, "faceseg_model.pb") face_seg(pb_path, lmk3D_crop_txt_path, crop_dir, seg_dir) fopen = open(out_img_names_path, "w") for i in range(0, len(names_list)): text = names_list[i] + "," + dep_name_list[i] + "\n" fopen.write(text) fopen.close() print("finish RGBD data preparation")