from scipy.spatial import ConvexHull import torch import torch.nn.functional as F import numpy as np from tqdm import tqdm import time def normalize_kp(kp_source, kp_driving, kp_driving_initial, adapt_movement_scale=False, use_relative_movement=False, use_relative_jacobian=False): if adapt_movement_scale: source_area = ConvexHull(kp_source['value'][0].data.cpu().numpy()).volume driving_area = ConvexHull(kp_driving_initial['value'][0].data.cpu().numpy()).volume adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area) else: adapt_movement_scale = 1 kp_new = {k: v for k, v in kp_driving.items()} if use_relative_movement: kp_value_diff = (kp_driving['value'] - kp_driving_initial['value']) kp_value_diff *= adapt_movement_scale kp_new['value'] = kp_value_diff + kp_source['value'] if use_relative_jacobian: jacobian_diff = torch.matmul(kp_driving['jacobian'], torch.inverse(kp_driving_initial['jacobian'])) kp_new['jacobian'] = torch.matmul(jacobian_diff, kp_source['jacobian']) return kp_new def headpose_pred_to_degree(pred): device = pred.device idx_tensor = [idx for idx in range(66)] idx_tensor = torch.FloatTensor(idx_tensor).to(device) pred = F.softmax(pred) degree = torch.sum(pred*idx_tensor, 1) * 3 - 99 return degree def get_rotation_matrix(yaw, pitch, roll): yaw = yaw / 180 * 3.14 pitch = pitch / 180 * 3.14 roll = roll / 180 * 3.14 roll = roll.unsqueeze(1) pitch = pitch.unsqueeze(1) yaw = yaw.unsqueeze(1) pitch_mat = torch.cat([torch.ones_like(pitch), torch.zeros_like(pitch), torch.zeros_like(pitch), torch.zeros_like(pitch), torch.cos(pitch), -torch.sin(pitch), torch.zeros_like(pitch), torch.sin(pitch), torch.cos(pitch)], dim=1) pitch_mat = pitch_mat.view(pitch_mat.shape[0], 3, 3) yaw_mat = torch.cat([torch.cos(yaw), torch.zeros_like(yaw), torch.sin(yaw), torch.zeros_like(yaw), torch.ones_like(yaw), torch.zeros_like(yaw), -torch.sin(yaw), torch.zeros_like(yaw), torch.cos(yaw)], dim=1) yaw_mat = yaw_mat.view(yaw_mat.shape[0], 3, 3) roll_mat = torch.cat([torch.cos(roll), -torch.sin(roll), torch.zeros_like(roll), torch.sin(roll), torch.cos(roll), torch.zeros_like(roll), torch.zeros_like(roll), torch.zeros_like(roll), torch.ones_like(roll)], dim=1) roll_mat = roll_mat.view(roll_mat.shape[0], 3, 3) rot_mat = torch.einsum('bij,bjk,bkm->bim', pitch_mat, yaw_mat, roll_mat) return rot_mat def keypoint_transformation(kp_canonical, he, wo_exp=False): kp = kp_canonical['value'] # (bs, k, 3) yaw, pitch, roll= he['yaw'], he['pitch'], he['roll'] yaw = headpose_pred_to_degree(yaw) pitch = headpose_pred_to_degree(pitch) roll = headpose_pred_to_degree(roll) if 'yaw_in' in he: yaw = he['yaw_in'] if 'pitch_in' in he: pitch = he['pitch_in'] if 'roll_in' in he: roll = he['roll_in'] rot_mat = get_rotation_matrix(yaw, pitch, roll) # (bs, 3, 3) t, exp = he['t'], he['exp'] if wo_exp: exp = exp*0 # keypoint rotation kp_rotated = torch.einsum('bmp,bkp->bkm', rot_mat, kp) # keypoint translation t[:, 0] = t[:, 0]*0 t[:, 2] = t[:, 2]*0 t = t.unsqueeze(1).repeat(1, kp.shape[1], 1) kp_t = kp_rotated + t # add expression deviation exp = exp.view(exp.shape[0], -1, 3) kp_transformed = kp_t + exp return {'value': kp_transformed} # def make_animation(source_image, source_semantics, target_semantics, # generator, kp_detector, he_estimator, mapping, # yaw_c_seq=None, pitch_c_seq=None, roll_c_seq=None, # use_exp=True): # with torch.no_grad(): # predictions = [] # kp_canonical = kp_detector(source_image) # he_source = mapping(source_semantics) # kp_source = keypoint_transformation(kp_canonical, he_source) # for frame_idx in tqdm(range(target_semantics.shape[1]), 'Face Renderer:'): # target_semantics_frame = target_semantics[:, frame_idx] # he_driving = mapping(target_semantics_frame) # if yaw_c_seq is not None: # he_driving['yaw_in'] = yaw_c_seq[:, frame_idx] # if pitch_c_seq is not None: # he_driving['pitch_in'] = pitch_c_seq[:, frame_idx] # if roll_c_seq is not None: # he_driving['roll_in'] = roll_c_seq[:, frame_idx] # kp_driving = keypoint_transformation(kp_canonical, he_driving) # #kp_norm = normalize_kp(kp_source=kp_source, kp_driving=kp_driving, # #kp_driving_initial=kp_driving_initial) # kp_norm = kp_driving # out = generator(source_image, kp_source=kp_source, kp_driving=kp_norm) # ''' # source_image_new = out['prediction'].squeeze(1) # kp_canonical_new = kp_detector(source_image_new) # he_source_new = he_estimator(source_image_new) # kp_source_new = keypoint_transformation(kp_canonical_new, he_source_new, wo_exp=True) # kp_driving_new = keypoint_transformation(kp_canonical_new, he_driving, wo_exp=True) # out = generator(source_image_new, kp_source=kp_source_new, kp_driving=kp_driving_new) # ''' # predictions.append(out['prediction']) # torch.cuda.empty_cache() # predictions_ts = torch.stack(predictions, dim=1) # return predictions_ts import torch from torch.cuda.amp import autocast def make_animation(source_image, source_semantics, target_semantics, generator, kp_detector, he_estimator, mapping, yaw_c_seq=None, pitch_c_seq=None, roll_c_seq=None, use_exp=True): device='cuda' # Move inputs to GPU source_image = source_image.to(device) source_semantics = source_semantics.to(device) target_semantics = target_semantics.to(device) with torch.no_grad(): # No gradients needed predictions = [] start = time.time() kp_canonical = kp_detector(source_image) he_source = mapping(source_semantics) kp_source = keypoint_transformation(kp_canonical, he_source) end = time.time() diff = end-start print("time taken for 3 values within make animation:", diff) with autocast(): for frame_idx in tqdm(range(target_semantics.shape[1]), desc='Face Renderer:', unit='frame'): target_semantics_frame = target_semantics[:, frame_idx] he_driving = mapping(target_semantics_frame) if yaw_c_seq is not None: he_driving['yaw_in'] = yaw_c_seq[:, frame_idx] if pitch_c_seq is not None: he_driving['pitch_in'] = pitch_c_seq[:, frame_idx] if roll_c_seq is not None: he_driving['roll_in'] = roll_c_seq[:, frame_idx] kp_driving = keypoint_transformation(kp_canonical, he_driving) kp_norm = kp_driving out = generator(source_image, kp_source=kp_source, kp_driving=kp_norm) predictions.append(out['prediction']) torch.cuda.synchronize() # Stack predictions into a single tensor predictions_ts = torch.stack(predictions, dim=1) return predictions_ts class AnimateModel(torch.nn.Module): """ Merge all generator related updates into single model for better multi-gpu usage """ def __init__(self, generator, kp_extractor, mapping): super(AnimateModel, self).__init__() self.kp_extractor = kp_extractor self.generator = generator self.mapping = mapping self.kp_extractor.eval() self.generator.eval() self.mapping.eval() def forward(self, x): source_image = x['source_image'] source_semantics = x['source_semantics'] target_semantics = x['target_semantics'] yaw_c_seq = x['yaw_c_seq'] pitch_c_seq = x['pitch_c_seq'] roll_c_seq = x['roll_c_seq'] predictions_video = make_animation(source_image, source_semantics, target_semantics, self.generator, self.kp_extractor, self.mapping, use_exp = True, yaw_c_seq=yaw_c_seq, pitch_c_seq=pitch_c_seq, roll_c_seq=roll_c_seq) return predictions_video