123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309 |
- import os
- import numpy as np
- import torch
- import torch.nn as nn
- from scipy.io import loadmat
- from deep_3drecon.deep_3drecon_models.bfm import perspective_projection
- class Face3DHelper(nn.Module):
- def __init__(self, bfm_dir='deep_3drecon/BFM', keypoint_mode='lm68', use_gpu=True):
- super().__init__()
- self.keypoint_mode = keypoint_mode # lm68 | mediapipe
- self.bfm_dir = bfm_dir
- self.load_3dmm()
- if use_gpu: self.to("cuda")
-
- def load_3dmm(self):
- model = loadmat(os.path.join(self.bfm_dir, "BFM_model_front.mat"))
- self.register_buffer('mean_shape',torch.from_numpy(model['meanshape'].transpose()).float()) # mean face shape. [3*N, 1], N=35709, xyz=3, ==> 3*N=107127
- mean_shape = self.mean_shape.reshape([-1, 3])
- # re-center
- mean_shape = mean_shape - torch.mean(mean_shape, dim=0, keepdims=True)
- self.mean_shape = mean_shape.reshape([-1, 1])
- self.register_buffer('id_base',torch.from_numpy(model['idBase']).float()) # identity basis. [3*N,80], we have 80 eigen faces for identity
- self.register_buffer('exp_base',torch.from_numpy(model['exBase']).float()) # expression basis. [3*N,64], we have 64 eigen faces for expression
-
- self.register_buffer('mean_texure',torch.from_numpy(model['meantex'].transpose()).float()) # mean face texture. [3*N,1] (0-255)
- self.register_buffer('tex_base',torch.from_numpy(model['texBase']).float()) # texture basis. [3*N,80], rgb=3
-
- self.register_buffer('point_buf',torch.from_numpy(model['point_buf']).float()) # triangle indices for each vertex that lies in. starts from 1. [N,8] (1-F)
- self.register_buffer('face_buf',torch.from_numpy(model['tri']).float()) # vertex indices in each triangle. starts from 1. [F,3] (1-N)
- if self.keypoint_mode == 'mediapipe':
- self.register_buffer('key_points', torch.from_numpy(np.load("deep_3drecon/BFM/index_mp468_from_mesh35709.npy").astype(np.int64)))
- unmatch_mask = self.key_points < 0
- self.key_points[unmatch_mask] = 0
- else:
- self.register_buffer('key_points',torch.from_numpy(model['keypoints'].squeeze().astype(np.int_)).long()) # vertex indices of 68 facial landmarks. starts from 1. [68,1]
-
- self.register_buffer('key_mean_shape',self.mean_shape.reshape([-1,3])[self.key_points,:])
- self.register_buffer('key_id_base', self.id_base.reshape([-1,3,80])[self.key_points, :, :].reshape([-1,80]))
- self.register_buffer('key_exp_base', self.exp_base.reshape([-1,3,64])[self.key_points, :, :].reshape([-1,64]))
- self.key_id_base_np = self.key_id_base.cpu().numpy()
- self.key_exp_base_np = self.key_exp_base.cpu().numpy()
- self.register_buffer('persc_proj', torch.tensor(perspective_projection(focal=1015, center=112)))
- def split_coeff(self, coeff):
- """
- coeff: Tensor[B, T, c=257] or [T, c=257]
- """
- ret_dict = {
- 'identity': coeff[..., :80], # identity, [b, t, c=80]
- 'expression': coeff[..., 80:144], # expression, [b, t, c=80]
- 'texture': coeff[..., 144:224], # texture, [b, t, c=80]
- 'euler': coeff[..., 224:227], # euler euler for pose, [b, t, c=3]
- 'translation': coeff[..., 254:257], # translation, [b, t, c=3]
- 'gamma': coeff[..., 227:254] # lighting, [b, t, c=27]
- }
- return ret_dict
-
- def reconstruct_face_mesh(self, id_coeff, exp_coeff):
- """
- Generate a pose-independent 3D face mesh!
- id_coeff: Tensor[T, c=80]
- exp_coeff: Tensor[T, c=64]
- """
- id_coeff = id_coeff.to(self.key_id_base.device)
- exp_coeff = exp_coeff.to(self.key_id_base.device)
- mean_face = self.mean_shape.squeeze().reshape([1, -1]) # [3N, 1] ==> [1, 3N]
- id_base, exp_base = self.id_base, self.exp_base # [3*N, C]
- identity_diff_face = torch.matmul(id_coeff, id_base.transpose(0,1)) # [t,c],[c,3N] ==> [t,3N]
- expression_diff_face = torch.matmul(exp_coeff, exp_base.transpose(0,1)) # [t,c],[c,3N] ==> [t,3N]
-
- face = mean_face + identity_diff_face + expression_diff_face # [t,3N]
- face = face.reshape([face.shape[0], -1, 3]) # [t,N,3]
- # re-centering the face with mean_xyz, so the face will be in [-1, 1]
- # mean_xyz = self.mean_shape.squeeze().reshape([-1,3]).mean(dim=0) # [1, 3]
- # face_mesh = face - mean_xyz.unsqueeze(0) # [t,N,3]
- return face
- def reconstruct_cano_lm3d(self, id_coeff, exp_coeff):
- """
- Generate 3D landmark with keypoint base!
- id_coeff: Tensor[T, c=80]
- exp_coeff: Tensor[T, c=64]
- """
- id_coeff = id_coeff.to(self.key_id_base.device)
- exp_coeff = exp_coeff.to(self.key_id_base.device)
- mean_face = self.key_mean_shape.squeeze().reshape([1, -1]) # [3*68, 1] ==> [1, 3*68]
- id_base, exp_base = self.key_id_base, self.key_exp_base # [3*68, C]
- identity_diff_face = torch.matmul(id_coeff, id_base.transpose(0,1)) # [t,c],[c,3*68] ==> [t,3*68]
- expression_diff_face = torch.matmul(exp_coeff, exp_base.transpose(0,1)) # [t,c],[c,3*68] ==> [t,3*68]
-
- face = mean_face + identity_diff_face + expression_diff_face # [t,3N]
- face = face.reshape([face.shape[0], -1, 3]) # [t,N,3]
- # re-centering the face with mean_xyz, so the face will be in [-1, 1]
- # mean_xyz = self.key_mean_shape.squeeze().reshape([-1,3]).mean(dim=0) # [1, 3]
- # lm3d = face - mean_xyz.unsqueeze(0) # [t,N,3]
- return face
- def reconstruct_lm3d(self, id_coeff, exp_coeff, euler, trans, to_camera=True):
- """
- Generate 3D landmark with keypoint base!
- id_coeff: Tensor[T, c=80]
- exp_coeff: Tensor[T, c=64]
- """
- id_coeff = id_coeff.to(self.key_id_base.device)
- exp_coeff = exp_coeff.to(self.key_id_base.device)
- mean_face = self.key_mean_shape.squeeze().reshape([1, -1]) # [3*68, 1] ==> [1, 3*68]
- id_base, exp_base = self.key_id_base, self.key_exp_base # [3*68, C]
- identity_diff_face = torch.matmul(id_coeff, id_base.transpose(0,1)) # [t,c],[c,3*68] ==> [t,3*68]
- expression_diff_face = torch.matmul(exp_coeff, exp_base.transpose(0,1)) # [t,c],[c,3*68] ==> [t,3*68]
-
- face = mean_face + identity_diff_face + expression_diff_face # [t,3N]
- face = face.reshape([face.shape[0], -1, 3]) # [t,N,3]
- # re-centering the face with mean_xyz, so the face will be in [-1, 1]
- rot = self.compute_rotation(euler)
- # transform
- lm3d = face @ rot + trans.unsqueeze(1) # [t, N, 3]
- # to camera
- if to_camera:
- lm3d[...,-1] = 10 - lm3d[...,-1]
- return lm3d
- def reconstruct_lm2d_nerf(self, id_coeff, exp_coeff, euler, trans):
- lm2d = self.reconstruct_lm2d(id_coeff, exp_coeff, euler, trans, to_camera=False)
- lm2d[..., 0] = 1 - lm2d[..., 0]
- lm2d[..., 1] = 1 - lm2d[..., 1]
- return lm2d
- def reconstruct_lm2d(self, id_coeff, exp_coeff, euler, trans, to_camera=True):
- """
- Generate 3D landmark with keypoint base!
- id_coeff: Tensor[T, c=80]
- exp_coeff: Tensor[T, c=64]
- """
- is_btc_flag = True if id_coeff.ndim == 3 else False
- if is_btc_flag:
- b,t,_ = id_coeff.shape
- id_coeff = id_coeff.reshape([b*t,-1])
- exp_coeff = exp_coeff.reshape([b*t,-1])
- euler = euler.reshape([b*t,-1])
- trans = trans.reshape([b*t,-1])
- id_coeff = id_coeff.to(self.key_id_base.device)
- exp_coeff = exp_coeff.to(self.key_id_base.device)
- mean_face = self.key_mean_shape.squeeze().reshape([1, -1]) # [3*68, 1] ==> [1, 3*68]
- id_base, exp_base = self.key_id_base, self.key_exp_base # [3*68, C]
- identity_diff_face = torch.matmul(id_coeff, id_base.transpose(0,1)) # [t,c],[c,3*68] ==> [t,3*68]
- expression_diff_face = torch.matmul(exp_coeff, exp_base.transpose(0,1)) # [t,c],[c,3*68] ==> [t,3*68]
-
- face = mean_face + identity_diff_face + expression_diff_face # [t,3N]
- face = face.reshape([face.shape[0], -1, 3]) # [t,N,3]
- # re-centering the face with mean_xyz, so the face will be in [-1, 1]
- rot = self.compute_rotation(euler)
- # transform
- lm3d = face @ rot + trans.unsqueeze(1) # [t, N, 3]
- # to camera
- if to_camera:
- lm3d[...,-1] = 10 - lm3d[...,-1]
- # to image_plane
- lm3d = lm3d @ self.persc_proj
- lm2d = lm3d[..., :2] / lm3d[..., 2:]
- # flip
- lm2d[..., 1] = 224 - lm2d[..., 1]
- lm2d /= 224
- if is_btc_flag:
- return lm2d.reshape([b,t,-1,2])
- return lm2d
-
- def compute_rotation(self, euler):
- """
- Return:
- rot -- torch.tensor, size (B, 3, 3) pts @ trans_mat
- Parameters:
- euler -- torch.tensor, size (B, 3), radian
- """
- batch_size = euler.shape[0]
- euler = euler.to(self.key_id_base.device)
- ones = torch.ones([batch_size, 1]).to(self.key_id_base.device)
- zeros = torch.zeros([batch_size, 1]).to(self.key_id_base.device)
- x, y, z = euler[:, :1], euler[:, 1:2], euler[:, 2:],
-
- rot_x = torch.cat([
- ones, zeros, zeros,
- zeros, torch.cos(x), -torch.sin(x),
- zeros, torch.sin(x), torch.cos(x)
- ], dim=1).reshape([batch_size, 3, 3])
-
- rot_y = torch.cat([
- torch.cos(y), zeros, torch.sin(y),
- zeros, ones, zeros,
- -torch.sin(y), zeros, torch.cos(y)
- ], dim=1).reshape([batch_size, 3, 3])
- rot_z = torch.cat([
- torch.cos(z), -torch.sin(z), zeros,
- torch.sin(z), torch.cos(z), zeros,
- zeros, zeros, ones
- ], dim=1).reshape([batch_size, 3, 3])
- rot = rot_z @ rot_y @ rot_x
- return rot.permute(0, 2, 1)
-
- def reconstruct_idexp_lm3d(self, id_coeff, exp_coeff):
- """
- Generate 3D landmark with keypoint base!
- id_coeff: Tensor[T, c=80]
- exp_coeff: Tensor[T, c=64]
- """
- id_coeff = id_coeff.to(self.key_id_base.device)
- exp_coeff = exp_coeff.to(self.key_id_base.device)
- id_base, exp_base = self.key_id_base, self.key_exp_base # [3*68, C]
- identity_diff_face = torch.matmul(id_coeff, id_base.transpose(0,1)) # [t,c],[c,3*68] ==> [t,3*68]
- expression_diff_face = torch.matmul(exp_coeff, exp_base.transpose(0,1)) # [t,c],[c,3*68] ==> [t,3*68]
-
- face = identity_diff_face + expression_diff_face # [t,3N]
- face = face.reshape([face.shape[0], -1, 3]) # [t,N,3]
- lm3d = face * 10
- return lm3d
-
- def reconstruct_idexp_lm3d_np(self, id_coeff, exp_coeff):
- """
- Generate 3D landmark with keypoint base!
- id_coeff: Tensor[T, c=80]
- exp_coeff: Tensor[T, c=64]
- """
- id_base, exp_base = self.key_id_base_np, self.key_exp_base_np # [3*68, C]
- identity_diff_face = np.dot(id_coeff, id_base.T) # [t,c],[c,3*68] ==> [t,3*68]
- expression_diff_face = np.dot(exp_coeff, exp_base.T) # [t,c],[c,3*68] ==> [t,3*68]
-
- face = identity_diff_face + expression_diff_face # [t,3N]
- face = face.reshape([face.shape[0], -1, 3]) # [t,N,3]
- lm3d = face * 10
- return lm3d
-
- def get_eye_mouth_lm_from_lm3d(self, lm3d):
- eye_lm = lm3d[:, 17:48] # [T, 31, 3]
- mouth_lm = lm3d[:, 48:68] # [T, 20, 3]
- return eye_lm, mouth_lm
-
- def get_eye_mouth_lm_from_lm3d_batch(self, lm3d):
- eye_lm = lm3d[:, :, 17:48] # [T, 31, 3]
- mouth_lm = lm3d[:, :, 48:68] # [T, 20, 3]
- return eye_lm, mouth_lm
-
- def close_mouth_for_idexp_lm3d(self, idexp_lm3d, freeze_as_first_frame=True):
- idexp_lm3d = idexp_lm3d.reshape([-1, 68,3])
- num_frames = idexp_lm3d.shape[0]
- eps = 0.0
- # [n_landmarks=68,xyz=3], x 代表左右,y代表上下,z代表深度
- idexp_lm3d[:,49:54, 1] = (idexp_lm3d[:,49:54, 1] + idexp_lm3d[:,range(59,54,-1), 1])/2 + eps * 2
- idexp_lm3d[:,range(59,54,-1), 1] = (idexp_lm3d[:,49:54, 1] + idexp_lm3d[:,range(59,54,-1), 1])/2 - eps * 2
- idexp_lm3d[:,61:64, 1] = (idexp_lm3d[:,61:64, 1] + idexp_lm3d[:,range(67,64,-1), 1])/2 + eps
- idexp_lm3d[:,range(67,64,-1), 1] = (idexp_lm3d[:,61:64, 1] + idexp_lm3d[:,range(67,64,-1), 1])/2 - eps
- idexp_lm3d[:,49:54, 1] += (0.03 - idexp_lm3d[:,49:54, 1].mean(dim=1) + idexp_lm3d[:,61:64, 1].mean(dim=1)).unsqueeze(1).repeat([1,5])
- idexp_lm3d[:,range(59,54,-1), 1] += (-0.03 - idexp_lm3d[:,range(59,54,-1), 1].mean(dim=1) + idexp_lm3d[:,range(67,64,-1), 1].mean(dim=1)).unsqueeze(1).repeat([1,5])
- if freeze_as_first_frame:
- idexp_lm3d[:, 48:68,] = idexp_lm3d[0, 48:68].unsqueeze(0).clone().repeat([num_frames, 1,1])*0
- return idexp_lm3d.cpu()
- def close_eyes_for_idexp_lm3d(self, idexp_lm3d):
- idexp_lm3d = idexp_lm3d.reshape([-1, 68,3])
- eps = 0.003
- idexp_lm3d[:,37:39, 1] = (idexp_lm3d[:,37:39, 1] + idexp_lm3d[:,range(41,39,-1), 1])/2 + eps
- idexp_lm3d[:,range(41,39,-1), 1] = (idexp_lm3d[:,37:39, 1] + idexp_lm3d[:,range(41,39,-1), 1])/2 - eps
- idexp_lm3d[:,43:45, 1] = (idexp_lm3d[:,43:45, 1] + idexp_lm3d[:,range(47,45,-1), 1])/2 + eps
- idexp_lm3d[:,range(47,45,-1), 1] = (idexp_lm3d[:,43:45, 1] + idexp_lm3d[:,range(47,45,-1), 1])/2 - eps
-
- return idexp_lm3d
- if __name__ == '__main__':
- import cv2
-
- font = cv2.FONT_HERSHEY_SIMPLEX
- face_mesh_helper = Face3DHelper('deep_3drecon/BFM')
- coeff_npy = 'data/coeff_fit_mp/crop_nana_003_coeff_fit_mp.npy'
- coeff_dict = np.load(coeff_npy, allow_pickle=True).tolist()
- lm3d = face_mesh_helper.reconstruct_lm2d(torch.tensor(coeff_dict['id']).cuda(), torch.tensor(coeff_dict['exp']).cuda(), torch.tensor(coeff_dict['euler']).cuda(), torch.tensor(coeff_dict['trans']).cuda() )
- WH = 512
- lm3d = (lm3d * WH).cpu().int().numpy()
- eye_idx = list(range(36,48))
- mouth_idx = list(range(48,68))
- import imageio
- debug_name = 'debug_lm3d.mp4'
- writer = imageio.get_writer(debug_name, fps=25)
- for i_img in range(len(lm3d)):
- lm2d = lm3d[i_img ,:, :2] # [68, 2]
- img = np.ones([WH, WH, 3], dtype=np.uint8) * 255
- for i in range(len(lm2d)):
- x, y = lm2d[i]
- if i in eye_idx:
- color = (0,0,255)
- elif i in mouth_idx:
- color = (0,255,0)
- else:
- color = (255,0,0)
- img = cv2.circle(img, center=(x,y), radius=3, color=color, thickness=-1)
- img = cv2.putText(img, f"{i}", org=(x,y), fontFace=font, fontScale=0.3, color=(255,0,0))
- writer.append_data(img)
- writer.close()
|