render_utils/calc_smplx2faceverse.py

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)