import copy import math import time, random import torch import numpy as np import cv2 import smplx import pickle import trimesh import os, glob import argparse from .lib.networks.faceverse_torch import FaceVerseModel def estimate_rigid_transformation(src, tgt): src = src.transpose() tgt = tgt.transpose() mu1, mu2 = src.mean(axis=1, keepdims=True), tgt.mean(axis=1, keepdims=True) X1, X2 = src - mu1, tgt - mu2 K = X1.dot(X2.T) U, s, Vh = np.linalg.svd(K) V = Vh.T Z = np.eye(U.shape[0]) Z[-1, -1] *= np.sign(np.linalg.det(U.dot(V.T))) R = V.dot(Z.dot(U.T)) t = mu2 - R.dot(mu1) orient, _ = cv2.Rodrigues(R) orient = orient.reshape([-1]) t = t.reshape([-1]) return orient, t def load_smpl_beta(body_fitting_result_fpath): if os.path.isdir(body_fitting_result_fpath): body_fitting_result = torch.load( os.path.join(body_fitting_result_fpath, 'checkpoints/latest.pt'), map_location='cpu') betas = body_fitting_result['betas']['weight'] elif body_fitting_result_fpath.endswith('.pt'): body_fitting_result = torch.load(body_fitting_result_fpath, map_location='cpu') betas = body_fitting_result['beta'].reshape(1, -1) elif body_fitting_result_fpath.endswith('.npz'): body_fitting_result = np.load(body_fitting_result_fpath) betas = body_fitting_result['betas'].reshape(1, -1) betas = torch.from_numpy(betas.astype(np.float32)) else: raise ValueError('Unknown body fitting result file format: {}'.format(body_fitting_result_fpath)) return betas def load_face_id_scale_param(face_fitting_result_fpath): if os.path.isfile(face_fitting_result_fpath): face_fitting_result = dict(np.load(face_fitting_result_fpath)) id_tensor = face_fitting_result['id_coeff'].astype(np.float32) scale_tensor = face_fitting_result['scale'].astype(np.float32) id_tensor = torch.from_numpy(id_tensor).reshape(1, -1) scale_tensor = torch.from_numpy(scale_tensor).reshape(1, -1) else: param_paths = sorted(glob.glob(os.path.join(face_fitting_result_fpath, '*', 'params.npz'))) param = np.load(param_paths[0]) id_tensor = torch.from_numpy(param['id_coeff']).reshape(1, -1) scale_tensor = torch.from_numpy(param['scale']).reshape(1, 1) return id_tensor, scale_tensor def calc_smplx2faceverse(body_fitting_result_fpath, face_fitting_result_fpath, output_dir): device = torch.device('cuda') os.makedirs(output_dir, exist_ok=True) betas = load_smpl_beta(body_fitting_result_fpath) id_tensor, scale_tensor = load_face_id_scale_param(face_fitting_result_fpath) smpl = smplx.SMPLX(model_path='./AnimatableGaussians/smpl_files/smplx', gender='neutral', use_pca=True, num_pca_comps=45, flat_hand_mean=True, batch_size=1) flame = smplx.FLAME(model_path='./AnimatableGaussians/smpl_files/FLAME2019', gender='neutral', batch_size=1) pose = np.zeros([63], dtype=np.float32) pose = torch.from_numpy(pose).unsqueeze(0) smpl_out = smpl(body_pose=pose, betas=betas) verts = smpl_out.vertices.detach().cpu().squeeze(0).numpy() flame_out = flame() verts_flame = flame_out.vertices.detach().cpu().squeeze(0).numpy() smplx2flame_data = np.load('./data/smpl_models/smplx_mano_flame_correspondences/SMPL-X__FLAME_vertex_ids.npy') verts_flame_on_smplx = verts[smplx2flame_data] orient, t = estimate_rigid_transformation(verts_flame, verts_flame_on_smplx) R, _ = cv2.Rodrigues(orient) rel_transf = np.eye(4) rel_transf[:3, :3] = R rel_transf[:3, 3] = t.reshape(-1) np.save('%s/flame_to_smplx.npy' % (output_dir), rel_transf.astype(np.float32)) # TODO: DELETE ME with open('./debug/debug_verts_smplx.obj', 'w') as fp: for v in verts: fp.write('v %f %f %f\n' % (v[0], v[1], v[2])) with open('./debug/debug_verts_flame_in_smplx.obj', 'w') as fp: for v in np.matmul(verts_flame, R.transpose()) + t.reshape([1, 3]): fp.write('v %f %f %f\n' % (v[0], v[1], v[2])) # align Faceverse to T-pose FLAME on SMPL-X faceverse_mesh = trimesh.load_mesh('./data/smpl_models/faceverse2flame/faceverse_icp.obj', process=False) verts_faceverse_ref = np.matmul(np.asarray(faceverse_mesh.vertices), R.transpose()) + t.reshape(1, 3) # TODO: DELETE ME with open('./debug/debug_verts_faceverse_ref.obj', 'w') as fp: for v in verts_faceverse_ref: fp.write('v %f %f %f\n' % (v[0], v[1], v[2])) model_dict = np.load('./data/faceverse_models/faceverse_simple_v2.npy', allow_pickle=True).item() faceverse_model = FaceVerseModel(model_dict, batch_size=1) faceverse_model.init_coeff_tensors() coeffs = faceverse_model.get_packed_tensors() fv_out = faceverse_model.forward(coeffs=coeffs) verts_faceverse = fv_out['v'].squeeze(0).detach().cpu().numpy() # calculate the relative transformation between FLAME in canonical pose and its position on SMPL-X orient, t = estimate_rigid_transformation(verts_faceverse, verts_faceverse_ref) orient = torch.from_numpy(orient.astype(np.float32)).unsqueeze(0).to(device) t = torch.from_numpy(t.astype(np.float32)).unsqueeze(0).to(device) # optimize the Faceverse to fit SMPL-X faceverse_model.init_coeff_tensors( rot_coeff=orient, trans_coeff=t, id_coeff=id_tensor.to(device), scale_coeff=scale_tensor.to(device)) nonrigid_optim_params = [ faceverse_model.get_exp_tensor(), faceverse_model.get_rot_tensor(), faceverse_model.get_trans_tensor(), # faceverse_model.get_scale_tensor(), faceverse_model.get_id_tensor() ] nonrigid_optimizer = torch.optim.Adam(nonrigid_optim_params, lr=1e-1) verts_faceverse_ref = torch.from_numpy(verts_faceverse_ref.astype(np.float32)).to(device).unsqueeze(0) for iter in range(1000): coeffs = faceverse_model.get_packed_tensors() fv_out = faceverse_model.forward(coeffs=coeffs) verts_pred = fv_out['v'] loss = torch.mean(torch.square(verts_pred - verts_faceverse_ref)) if iter % 10 == 0: print(loss.item()) nonrigid_optimizer.zero_grad() loss.backward() nonrigid_optimizer.step() np.savez('%s/faceverse_param_to_smplx.npz' % (output_dir), { 'id': faceverse_model.get_id_tensor().detach().cpu().numpy(), 'exp': faceverse_model.get_exp_tensor().detach().cpu().numpy(), 'rot': faceverse_model.get_rot_tensor().detach().cpu().numpy(), 'transl': faceverse_model.get_trans_tensor().detach().cpu().numpy(), 'scale': faceverse_model.get_scale_tensor().detach().cpu().numpy(), }) # calculate SMPLX to faceverse space transformation (without scale) orient = faceverse_model.get_rot_tensor().detach().cpu().numpy() transl = faceverse_model.get_trans_tensor().detach().cpu().numpy() rotmat, _ = cv2.Rodrigues(orient) transform_mat = np.eye(4) transform_mat[:3, :3] = rotmat transform_mat[:3, 3] = transl transform_mat = np.linalg.inv(transform_mat) np.save('%s/smplx_to_faceverse_space.npy' % (output_dir), transform_mat.astype(np.float32)) # calculate SMPLX to faceverse distance dists = [] verts_faceverse_ref = verts_faceverse_ref.detach().cpu().numpy() for v in verts: dist = np.linalg.norm(v.reshape(1, 3) - verts_faceverse_ref, axis=-1) dist = np.min(dist) dists.append(dist) dists = np.asarray(dists) np.save('%s/smplx_verts_to_faceverse_dist.npy' % (output_dir), dists.astype(np.float32)) # sample nodes on facial area dists_ = np.ones_like(dists) smplx_topo_new = np.load('./data/smpl_models/smplx_topo_new.npy') valid_vids = set(smplx_topo_new.reshape([-1]).tolist()) dists_[np.asarray(list(valid_vids))] = dists[np.asarray(list(valid_vids))] vids_on_face = np.where(dists_ < 0.01)[0] verts_on_face = verts[vids_on_face] geod_dist_mat = np.linalg.norm(np.expand_dims(verts_on_face, axis=0) - np.expand_dims(verts_on_face, axis=1), axis=2) nodes = [0] # nose dist_nodes_to_rest_points = geod_dist_mat[nodes[0]] for i in range(1, 256): idx = np.argmax(dist_nodes_to_rest_points) nodes.append(idx) new_dist = geod_dist_mat[idx] update_flag = new_dist < dist_nodes_to_rest_points dist_nodes_to_rest_points[update_flag] = new_dist[update_flag] # with open('./debug/debug_face_nodes.obj', 'w') as fp: # for n in verts_on_face[np.asarray(nodes)]: # fp.write('v %f %f %f\n' % (n[0], n[1], n[2])) vids_on_faces_sampled = vids_on_face[np.asarray(nodes)] vids_on_faces_sampled = np.ascontiguousarray(vids_on_faces_sampled).astype(np.int32) np.save('%s/vids_on_faces_sampled.npy' % (output_dir), vids_on_faces_sampled) # test SMPLX-to-faceverse space transformation (without scale) verts_smpl_in_faceverse = np.matmul(verts, transform_mat[:3, :3].transpose()) + \ transform_mat[:3, 3].reshape(1, 3) with open('./debug/debug_verts_smpl_in_faceverse.obj', 'w') as fp: for v in verts_smpl_in_faceverse: fp.write('v %f %f %f\n' % (v[0], v[1], v[2])) # save personalized, canonical faceverse model faceverse_model.init_coeff_tensors(id_coeff=id_tensor.to(device), scale_coeff=scale_tensor.to(device)) coeffs = faceverse_model.get_packed_tensors() fv_out = faceverse_model.forward(coeffs=coeffs) verts_faceverse = fv_out['v'].squeeze(0).detach().cpu().numpy() with open('./debug/debug_verts_faceverse.obj', 'w') as fp: for v in verts_faceverse: fp.write('v %f %f %f\n' % (v[0], v[1], v[2])) if __name__ == '__main__': # body_fitting_result_dir = 'D:/UpperBodyAvatar/code/smplfitting_multiview_sequence_smplx/results/Shuangqing/zzr_fullbody_20221130_01_2k/whole.pt' # face_fitting_result_dir = 'D:\\UpperBodyAvatar\\data\\Shuangqing\\zzr_face_20221130_01_2k\\faceverse_params' # # output_dir = './data/faceverse/' # result_suffix = 'shuangqing_zzr' # body_fitting_result_dir = 'D:/Product/FullAppRelease/smplfitting_multiview_sequence_smplx/results/body_data_model_20231224/whole.pt' # face_fitting_result_dir = 'D:/Product/data/HuiyingCenter/20231224_model/model_20231224_face_data/faceverse_params' # # output_dir = './data/faceverse/' # result_suffix = 'huiyin_model20231224' body_fitting_result_dir = 'D:/Product/FullAppRelease/smplfitting_multiview_sequence_smplx/results/body_data_male20230530_betterhand/whole.pt' face_fitting_result_dir = 'D:/Product/data/HuiyingCenter/20230531_models/male20230530_face_data/faceverse_params' output_dir = './data/body_face_stitching/huiyin_male20230530' calc_smplx2faceverse(body_fitting_result_dir, face_fitting_result_dir, output_dir)