Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,718 Bytes
8ed2f16 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
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)
|