cnywt/SyncTalk

1import torch2import torch.nn as nn3import numpy as np4import os5 6 7class Face_3DMM(nn.Module):8    def __init__(self, modelpath, id_dim, exp_dim, tex_dim, point_num):9        super(Face_3DMM, self).__init__()10        # id_dim = 10011        # exp_dim = 7912        # tex_dim = 10013        self.point_num = point_num14        DMM_info = np.load(15            os.path.join(modelpath, "3DMM_info.npy"), allow_pickle=True16        ).item()17        base_id = DMM_info["b_shape"][:id_dim, :]18        mu_id = DMM_info["mu_shape"]19        base_exp = DMM_info["b_exp"][:exp_dim, :]20        mu_exp = DMM_info["mu_exp"]21        mu = mu_id + mu_exp22        mu = mu.reshape(-1, 3)23        for i in range(3):24            mu[:, i] -= np.mean(mu[:, i])25        mu = mu.reshape(-1)26        self.base_id = torch.as_tensor(base_id).cuda() /1000.027        self.base_exp = torch.as_tensor(base_exp).cuda() /1000.028        self.mu = torch.as_tensor(mu).cuda() /1000.029        base_tex = DMM_info["b_tex"][:tex_dim, :]30        mu_tex = DMM_info["mu_tex"]31        self.base_tex = torch.as_tensor(base_tex).cuda()32        self.mu_tex = torch.as_tensor(mu_tex).cuda()33        sig_id = DMM_info["sig_shape"][:id_dim]34        sig_tex = DMM_info["sig_tex"][:tex_dim]35        sig_exp = DMM_info["sig_exp"][:exp_dim]36        self.sig_id = torch.as_tensor(sig_id).cuda()37        self.sig_tex = torch.as_tensor(sig_tex).cuda()38        self.sig_exp = torch.as_tensor(sig_exp).cuda()39 40    def forward_geo_sub(self, id_para, exp_para, sub_index):41        id_para = id_para*self.sig_id42        exp_para = exp_para*self.sig_exp43        sel_index = torch.cat((3*sub_index.unsqueeze(1), 3*sub_index.unsqueeze(1)+1,44                               3*sub_index.unsqueeze(1)+2), dim=1).reshape(-1)45        geometry = torch.mm(id_para, self.base_id[:, sel_index]) + \46            torch.mm(exp_para, self.base_exp[:,47                                             sel_index]) + self.mu[sel_index]48        return geometry.reshape(-1, sub_index.shape[0], 3)49 50    def forward_geo(self, id_para, exp_para):51        id_para = id_para*self.sig_id52        exp_para = exp_para*self.sig_exp53        geometry = torch.mm(id_para, self.base_id) + \54            torch.mm(exp_para, self.base_exp) + self.mu55        return geometry.reshape(-1, self.point_num, 3)56 57    def forward_tex(self, tex_para):58        tex_para = tex_para*self.sig_tex59        texture = torch.mm(tex_para, self.base_tex) + self.mu_tex60        return texture.reshape(-1, self.point_num, 3)61