RabbitRUI/ruispace
0
1"""This script defines the parametric 3d face model for Deep3DFaceRecon_pytorch2"""3 4import numpy as np5import torch6import torch.nn.functional as F7from scipy.io import loadmat8from src.face3d.util.load_mats import transferBFM099import os10 11def perspective_projection(focal, center):12 # return p.T (N, 3) @ (3, 3) 13 return np.array([14 focal, 0, center,15 0, focal, center,16 0, 0, 117 ]).reshape([3, 3]).astype(np.float32).transpose()18 19class SH:20 def __init__(self):21 self.a = [np.pi, 2 * np.pi / np.sqrt(3.), 2 * np.pi / np.sqrt(8.)]22 self.c = [1/np.sqrt(4 * np.pi), np.sqrt(3.) / np.sqrt(4 * np.pi), 3 * np.sqrt(5.) / np.sqrt(12 * np.pi)]23 24 25 26class ParametricFaceModel:27 def __init__(self, 28 bfm_folder='./BFM', 29 recenter=True,30 camera_distance=10.,31 init_lit=np.array([32 0.8, 0, 0, 0, 0, 0, 0, 0, 033 ]),34 focal=1015.,35 center=112.,36 is_train=True,37 default_name='BFM_model_front.mat'):38 39 if not os.path.isfile(os.path.join(bfm_folder, default_name)):40 transferBFM09(bfm_folder)41 42 model = loadmat(os.path.join(bfm_folder, default_name))43 # mean face shape. [3*N,1]44 self.mean_shape = model['meanshape'].astype(np.float32)45 # identity basis. [3*N,80]46 self.id_base = model['idBase'].astype(np.float32)47 # expression basis. [3*N,64]48 self.exp_base = model['exBase'].astype(np.float32)49 # mean face texture. [3*N,1] (0-255)50 self.mean_tex = model['meantex'].astype(np.float32)51 # texture basis. [3*N,80]52 self.tex_base = model['texBase'].astype(np.float32)53 # face indices for each vertex that lies in. starts from 0. [N,8]54 self.point_buf = model['point_buf'].astype(np.int64) - 155 # vertex indices for each face. starts from 0. [F,3]56 self.face_buf = model['tri'].astype(np.int64) - 157 # vertex indices for 68 landmarks. starts from 0. [68,1]58 self.keypoints = np.squeeze(model['keypoints']).astype(np.int64) - 159 60 if is_train:61 # vertex indices for small face region to compute photometric error. starts from 0.62 self.front_mask = np.squeeze(model['frontmask2_idx']).astype(np.int64) - 163 # vertex indices for each face from small face region. starts from 0. [f,3]64 self.front_face_buf = model['tri_mask2'].astype(np.int64) - 165 # vertex indices for pre-defined skin region to compute reflectance loss66 self.skin_mask = np.squeeze(model['skinmask'])67 68 if recenter:69 mean_shape = self.mean_shape.reshape([-1, 3])70 mean_shape = mean_shape - np.mean(mean_shape, axis=0, keepdims=True)71 self.mean_shape = mean_shape.reshape([-1, 1])72 73 self.persc_proj = perspective_projection(focal, center)74 self.device = 'cpu'75 self.camera_distance = camera_distance76 self.SH = SH()77 self.init_lit = init_lit.reshape([1, 1, -1]).astype(np.float32)78 79 80 def to(self, device):81 self.device = device82 for key, value in self.__dict__.items():83 if type(value).__module__ == np.__name__:84 setattr(self, key, torch.tensor(value).to(device))85 86 87 def compute_shape(self, id_coeff, exp_coeff):88 """89 Return:90 face_shape -- torch.tensor, size (B, N, 3)91 92 Parameters:93 id_coeff -- torch.tensor, size (B, 80), identity coeffs94 exp_coeff -- torch.tensor, size (B, 64), expression coeffs95 """96 batch_size = id_coeff.shape[0]97 id_part = torch.einsum('ij,aj->ai', self.id_base, id_coeff)98 exp_part = torch.einsum('ij,aj->ai', self.exp_base, exp_coeff)99 face_shape = id_part + exp_part + self.mean_shape.reshape([1, -1])100 return face_shape.reshape([batch_size, -1, 3])101 102 103 def compute_texture(self, tex_coeff, normalize=True):104 """105 Return:106 face_texture -- torch.tensor, size (B, N, 3), in RGB order, range (0, 1.)107 108 Parameters:109 tex_coeff -- torch.tensor, size (B, 80)110 """111 batch_size = tex_coeff.shape[0]112 face_texture = torch.einsum('ij,aj->ai', self.tex_base, tex_coeff) + self.mean_tex113 if normalize:114 face_texture = face_texture / 255.115 return face_texture.reshape([batch_size, -1, 3])116 117 118 def compute_norm(self, face_shape):119 """120 Return:121 vertex_norm -- torch.tensor, size (B, N, 3)122 123 Parameters:124 face_shape -- torch.tensor, size (B, N, 3)125 """126 127 v1 = face_shape[:, self.face_buf[:, 0]]128 v2 = face_shape[:, self.face_buf[:, 1]]129 v3 = face_shape[:, self.face_buf[:, 2]]130 e1 = v1 - v2131 e2 = v2 - v3132 face_norm = torch.cross(e1, e2, dim=-1)133 face_norm = F.normalize(face_norm, dim=-1, p=2)134 face_norm = torch.cat([face_norm, torch.zeros(face_norm.shape[0], 1, 3).to(self.device)], dim=1)135 136 vertex_norm = torch.sum(face_norm[:, self.point_buf], dim=2)137 vertex_norm = F.normalize(vertex_norm, dim=-1, p=2)138 return vertex_norm139 140 141 def compute_color(self, face_texture, face_norm, gamma):142 """143 Return:144 face_color -- torch.tensor, size (B, N, 3), range (0, 1.)145 146 Parameters:147 face_texture -- torch.tensor, size (B, N, 3), from texture model, range (0, 1.)148 face_norm -- torch.tensor, size (B, N, 3), rotated face normal149 gamma -- torch.tensor, size (B, 27), SH coeffs150 """151 batch_size = gamma.shape[0]152 v_num = face_texture.shape[1]153 a, c = self.SH.a, self.SH.c154 gamma = gamma.reshape([batch_size, 3, 9])155 gamma = gamma + self.init_lit156 gamma = gamma.permute(0, 2, 1)157 Y = torch.cat([158 a[0] * c[0] * torch.ones_like(face_norm[..., :1]).to(self.device),159 -a[1] * c[1] * face_norm[..., 1:2],160 a[1] * c[1] * face_norm[..., 2:],161 -a[1] * c[1] * face_norm[..., :1],162 a[2] * c[2] * face_norm[..., :1] * face_norm[..., 1:2],163 -a[2] * c[2] * face_norm[..., 1:2] * face_norm[..., 2:],164 0.5 * a[2] * c[2] / np.sqrt(3.) * (3 * face_norm[..., 2:] ** 2 - 1),165 -a[2] * c[2] * face_norm[..., :1] * face_norm[..., 2:],166 0.5 * a[2] * c[2] * (face_norm[..., :1] ** 2 - face_norm[..., 1:2] ** 2)167 ], dim=-1)168 r = Y @ gamma[..., :1]169 g = Y @ gamma[..., 1:2]170 b = Y @ gamma[..., 2:]171 face_color = torch.cat([r, g, b], dim=-1) * face_texture172 return face_color173 174 175 def compute_rotation(self, angles):176 """177 Return:178 rot -- torch.tensor, size (B, 3, 3) pts @ trans_mat179 180 Parameters:181 angles -- torch.tensor, size (B, 3), radian182 """183 184 batch_size = angles.shape[0]185 ones = torch.ones([batch_size, 1]).to(self.device)186 zeros = torch.zeros([batch_size, 1]).to(self.device)187 x, y, z = angles[:, :1], angles[:, 1:2], angles[:, 2:],188 189 rot_x = torch.cat([190 ones, zeros, zeros,191 zeros, torch.cos(x), -torch.sin(x), 192 zeros, torch.sin(x), torch.cos(x)193 ], dim=1).reshape([batch_size, 3, 3])194 195 rot_y = torch.cat([196 torch.cos(y), zeros, torch.sin(y),197 zeros, ones, zeros,198 -torch.sin(y), zeros, torch.cos(y)199 ], dim=1).reshape([batch_size, 3, 3])200 201 rot_z = torch.cat([202 torch.cos(z), -torch.sin(z), zeros,203 torch.sin(z), torch.cos(z), zeros,204 zeros, zeros, ones205 ], dim=1).reshape([batch_size, 3, 3])206 207 rot = rot_z @ rot_y @ rot_x208 return rot.permute(0, 2, 1)209 210 211 def to_camera(self, face_shape):212 face_shape[..., -1] = self.camera_distance - face_shape[..., -1]213 return face_shape214 215 def to_image(self, face_shape):216 """217 Return:218 face_proj -- torch.tensor, size (B, N, 2), y direction is opposite to v direction219 220 Parameters:221 face_shape -- torch.tensor, size (B, N, 3)222 """223 # to image_plane224 face_proj = face_shape @ self.persc_proj225 face_proj = face_proj[..., :2] / face_proj[..., 2:]226 227 return face_proj228 229 230 def transform(self, face_shape, rot, trans):231 """232 Return:233 face_shape -- torch.tensor, size (B, N, 3) pts @ rot + trans234 235 Parameters:236 face_shape -- torch.tensor, size (B, N, 3)237 rot -- torch.tensor, size (B, 3, 3)238 trans -- torch.tensor, size (B, 3)239 """240 return face_shape @ rot + trans.unsqueeze(1)241 242 243 def get_landmarks(self, face_proj):244 """245 Return:246 face_lms -- torch.tensor, size (B, 68, 2)247 248 Parameters:249 face_proj -- torch.tensor, size (B, N, 2)250 """ 251 return face_proj[:, self.keypoints]252 253 def split_coeff(self, coeffs):254 """255 Return:256 coeffs_dict -- a dict of torch.tensors257 258 Parameters:259 coeffs -- torch.tensor, size (B, 256)260 """261 id_coeffs = coeffs[:, :80]262 exp_coeffs = coeffs[:, 80: 144]263 tex_coeffs = coeffs[:, 144: 224]264 angles = coeffs[:, 224: 227]265 gammas = coeffs[:, 227: 254]266 translations = coeffs[:, 254:]267 return {268 'id': id_coeffs,269 'exp': exp_coeffs,270 'tex': tex_coeffs,271 'angle': angles,272 'gamma': gammas,273 'trans': translations274 }275 def compute_for_render(self, coeffs):276 """277 Return:278 face_vertex -- torch.tensor, size (B, N, 3), in camera coordinate279 face_color -- torch.tensor, size (B, N, 3), in RGB order280 landmark -- torch.tensor, size (B, 68, 2), y direction is opposite to v direction281 Parameters:282 coeffs -- torch.tensor, size (B, 257)283 """284 coef_dict = self.split_coeff(coeffs)285 face_shape = self.compute_shape(coef_dict['id'], coef_dict['exp'])286 rotation = self.compute_rotation(coef_dict['angle'])287 288 289 face_shape_transformed = self.transform(face_shape, rotation, coef_dict['trans'])290 face_vertex = self.to_camera(face_shape_transformed)291 292 face_proj = self.to_image(face_vertex)293 landmark = self.get_landmarks(face_proj)294 295 face_texture = self.compute_texture(coef_dict['tex'])296 face_norm = self.compute_norm(face_shape)297 face_norm_roted = face_norm @ rotation298 face_color = self.compute_color(face_texture, face_norm_roted, coef_dict['gamma'])299 300 return face_vertex, face_texture, face_color, landmark301 302 def compute_for_render_woRotation(self, coeffs):303 """304 Return:305 face_vertex -- torch.tensor, size (B, N, 3), in camera coordinate306 face_color -- torch.tensor, size (B, N, 3), in RGB order307 landmark -- torch.tensor, size (B, 68, 2), y direction is opposite to v direction308 Parameters:309 coeffs -- torch.tensor, size (B, 257)310 """311 coef_dict = self.split_coeff(coeffs)312 face_shape = self.compute_shape(coef_dict['id'], coef_dict['exp'])313 #rotation = self.compute_rotation(coef_dict['angle'])314 315 316 #face_shape_transformed = self.transform(face_shape, rotation, coef_dict['trans'])317 face_vertex = self.to_camera(face_shape)318 319 face_proj = self.to_image(face_vertex)320 landmark = self.get_landmarks(face_proj)321 322 face_texture = self.compute_texture(coef_dict['tex'])323 face_norm = self.compute_norm(face_shape)324 face_norm_roted = face_norm # @ rotation325 face_color = self.compute_color(face_texture, face_norm_roted, coef_dict['gamma'])326 327 return face_vertex, face_texture, face_color, landmark328 329 330if __name__ == '__main__':331 transferBFM09()