Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import numpy as np | |
| import pickle | |
| from scipy.io import loadmat | |
| from pytorch3d.transforms import so3_exponential_map | |
| class BFMModule(nn.Module): | |
| def __init__(self, batch_size): | |
| super(BFMModule, self).__init__() | |
| self.id_dims = 80 | |
| self.exp_dims = 64 | |
| model_dict = loadmat('assets/BFM/BFM09_model_info.mat') | |
| self.register_buffer('skinmask', torch.tensor(model_dict['skinmask'])) | |
| kp_inds = torch.tensor(model_dict['keypoints']-1).squeeze().long() | |
| #kp_inds = torch.cat([kp_inds[0:48], kp_inds[49:54], kp_inds[55:68]]) | |
| self.register_buffer('kp_inds', kp_inds) | |
| self.register_buffer('meanshape', torch.tensor(model_dict['meanshape']).float()) | |
| self.register_buffer('idBase', torch.tensor(model_dict['idBase']).float()) | |
| self.register_buffer('expBase', torch.tensor(model_dict['exBase']).float()) | |
| self.register_buffer('faces', torch.tensor(model_dict['tri']-1).long()) | |
| self.batch_size = batch_size | |
| self.id_coeff = nn.Parameter(torch.zeros(1, self.id_dims).float()) | |
| self.exp_coeff = nn.Parameter(torch.zeros(self.batch_size, self.exp_dims).float()) | |
| self.scale = nn.Parameter(torch.ones(1).float() * 0.15) | |
| self.pose = nn.Parameter(torch.zeros(self.batch_size, 6).float()) | |
| def get_lms(self, vs): | |
| lms = vs[:, self.kp_inds, :] | |
| return lms | |
| def get_vs(self, id_coeff, exp_coeff): | |
| n_b = id_coeff.size(0) | |
| face_shape = torch.einsum('ij,aj->ai', self.idBase, id_coeff) + \ | |
| torch.einsum('ij,aj->ai', self.expBase, exp_coeff) + self.meanshape | |
| face_shape = face_shape.view(n_b, -1, 3) | |
| face_shape = face_shape - \ | |
| self.meanshape.view(1, -1, 3).mean(dim=1, keepdim=True) | |
| return face_shape | |
| def forward(self): | |
| id_coeff = self.id_coeff.repeat(self.batch_size, 1) | |
| vertices = self.get_vs(id_coeff, self.exp_coeff) | |
| R = so3_exponential_map(self.pose[:, :3]) | |
| T = self.pose[:, 3:] | |
| vertices = torch.bmm(vertices * self.scale, R.permute(0,2,1)) + T[:, None, :] | |
| landmarks = self.get_lms(vertices) | |
| return vertices, landmarks | |
| def reg_loss(self, id_weight, exp_weight): | |
| id_reg_loss = (self.id_coeff ** 2).sum() | |
| exp_reg_loss = (self.exp_coeff ** 2).sum(-1).mean() | |
| return id_reg_loss * id_weight + exp_reg_loss * exp_weight | |
| def save(self, path, batch_id=-1): | |
| if batch_id < 0: | |
| id_coeff = self.id_coeff.detach().cpu().numpy() | |
| exp_coeff = self.exp_coeff.detach().cpu().numpy() | |
| scale = self.scale.detach().cpu().numpy() | |
| pose = self.pose.detach().cpu().numpy() | |
| np.savez(path, id_coeff=id_coeff, exp_coeff=exp_coeff, scale=scale, pose=pose) | |
| else: | |
| id_coeff = self.id_coeff.detach().cpu().numpy() | |
| exp_coeff = self.exp_coeff[batch_id:batch_id+1].detach().cpu().numpy() | |
| scale = self.scale.detach().cpu().numpy() | |
| pose = self.pose[batch_id:batch_id+1].detach().cpu().numpy() | |
| np.savez(path, id_coeff=id_coeff, exp_coeff=exp_coeff, scale=scale, pose=pose) |