Spaces:
Running on Zero
Running on Zero
| import shutil | |
| import os | |
| import json | |
| from tqdm import tqdm | |
| import numpy as np | |
| from scipy.io import loadmat | |
| from scipy.ndimage import gaussian_filter1d | |
| from preprocess_faceverse import make_cam_dataset_FFHQ, render_orth_mp | |
| def make_faceverse_labels_FFHQ(tracking_dir, root_dir, fv2fl_T_path, focal, need_render=False, save_uv=True, save_mesh=False, save_name=None, | |
| render_normal_uv=False, num_thread=1, use_smooth=False, test_data=False, skip=False): | |
| save_dir = os.path.join(root_dir, 'dataset') | |
| fv2fl_T = np.load(fv2fl_T_path).astype(np.float32) | |
| orth_scale, orth_shift, box_warp = 5.00, np.asarray([0, 0.005, 0.], dtype=np.float32), 2. | |
| face_model_dir = 'lib/data_preprocess/FaceVerse/v3' | |
| save_render_dir = os.path.join(save_dir, 'orthRender256x256_face_eye' if save_name is None else save_name) | |
| save_mesh_dir = None if not save_mesh else os.path.join(save_dir, 'FVmeshes512x512') | |
| save_uv_dir = None if not save_uv else os.path.join(save_dir, 'uvRender256x256') | |
| render_orth_mp(tracking_dir, save_render_dir, face_model_dir, fv2fl_T, {'scale': orth_scale, 'shift': orth_shift}, focal, render_vis=need_render, | |
| save_mesh_dir=save_mesh_dir, save_uv_dir=save_uv_dir, render_normal_uv=render_normal_uv, skip=skip, | |
| num_thread=num_thread, crop_param=[128, 114, 256, 256], save_coeff=True) | |
| normalizeFL_T = np.eye(4, dtype=np.float32) | |
| scale_T = (orth_scale / box_warp) * np.eye(3, dtype=np.float32) | |
| shift_T = scale_T.dot(orth_shift.reshape(3, 1)) | |
| normalizeFL_T[:3, :3], normalizeFL_T[:3, 3:] = scale_T, shift_T | |
| fv2fl_T = np.dot(normalizeFL_T, fv2fl_T) | |
| cam_params, cond_cam_params, fv_exp_eye_params = make_cam_dataset_FFHQ(tracking_dir, fv2fl_T, focal, test_data=test_data) | |
| if test_data: | |
| for prefix in cam_params.keys(): | |
| save_json_name = 'dataset_%s_realcam.json' % prefix | |
| with open(os.path.join(save_dir, 'images512x512', save_json_name), "w") as f: | |
| json.dump({"labels": cam_params[prefix]}, f, indent=4) | |
| else: | |
| save_json_name = 'dataset_realcam.json' | |
| if use_smooth: | |
| new_json = [] | |
| for sub_name in os.listdir(os.path.join(save_dir, 'images512x512')): | |
| if not os.path.isdir(os.path.join(save_dir, 'images512x512', sub_name)): continue | |
| sub_json = [case for case in cam_params if case[0].split('/')[0] == sub_name] | |
| sub_json.sort(key=lambda x: int(x[0].split('/')[1].split('.')[0])) | |
| coeff_seq = np.asarray([x[1] for x in sub_json], dtype=np.float32) | |
| coeff_seq = gaussian_filter1d(coeff_seq, sigma=1.5, axis=0) | |
| new_json.extend([[x[0], coeff_seq[idx].tolist()] for idx, x in enumerate(sub_json)]) | |
| cam_params = new_json | |
| with open(os.path.join(save_dir, 'images512x512', save_json_name), "w") as f: | |
| json.dump({"labels": cam_params}, f, indent=4) | |
| make_coeff_dataset_FFHQ(tracking_dir, os.path.join(save_dir, 'coeffs'), smooth=use_smooth) | |
| def make_coeff_dataset_FFHQ(tracking_dir, save_dir, smooth=False): | |
| for prefix in tqdm(os.listdir(tracking_dir)): | |
| if not os.path.isdir(os.path.join(tracking_dir, prefix)): | |
| continue | |
| sub_dir = os.path.join(tracking_dir, prefix) | |
| fname_ls = [name for name in os.listdir(sub_dir) if os.path.exists(os.path.join(sub_dir, name, 'finish'))] | |
| fname_ls.sort(key=lambda x: int(x)) | |
| try: | |
| coeff_seq = np.stack([np.load(os.path.join(sub_dir, fname, 'coeffs.npy')) for fname in fname_ls], axis=0) | |
| if smooth: coeff_seq = gaussian_filter1d(coeff_seq, sigma=0.5, axis=0) | |
| os.makedirs(os.path.join(save_dir, prefix), exist_ok=True) | |
| for idx, fname in enumerate(fname_ls): | |
| dst_path = os.path.join(save_dir, prefix, fname + '.npy') | |
| np.save(dst_path, coeff_seq[idx]) | |
| except: | |
| continue | |
| if __name__ == "__main__": | |
| import argparse | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--tracking_dir', type=str, default=None) | |
| parser.add_argument('--root_dir', type=str, default=None) | |
| parser.add_argument('--test_data', action='store_true', default=False) | |
| parser.add_argument('--skip', action='store_true', default=False) | |
| args = parser.parse_args() | |
| make_faceverse_labels_FFHQ(tracking_dir=args.tracking_dir, root_dir=args.root_dir, | |
| fv2fl_T_path='FaceVerse/v3/fv2fl_30.npy', need_render=False, | |
| save_mesh=False, focal=4.2647, num_thread=1 if args.test_data else 8, test_data=args.test_data, skip=args.skip) | |