import os import subprocess from os.path import join import yaml import tempfile import argparse from skimage.io import imread import numpy as np import librosa from util import util from tqdm import tqdm import torch from collections import OrderedDict import cv2 from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_subclip from cog import BasePredictor, Input, Path import scipy.io as sio import albumentations as A from options.test_audio2feature_options import TestOptions as FeatureOptions from options.test_audio2headpose_options import TestOptions as HeadposeOptions from options.test_feature2face_options import TestOptions as RenderOptions from datasets import create_dataset from models import create_model from models.networks import APC_encoder from util.visualizer import Visualizer from funcs import utils, audio_funcs from demo import write_video_with_audio import warnings warnings.filterwarnings("ignore") class Predictor(BasePredictor): def setup(self): self.parser = argparse.ArgumentParser() self.parser.add_argument('--id', default='May', help="person name, e.g. Obama1, Obama2, May, Nadella, McStay") self.parser.add_argument('--driving_audio', default='data/Input/00083.wav', help="path to driving audio") self.parser.add_argument('--save_intermediates', default=0, help="whether to save intermediate results") def predict(self, driving_audio: Path = Input(description='driving audio, if the file is more than 20 seconds, only the first 20 seconds will be processed for video generation'), talking_head: str = Input(description="choose a talking head", choices=['May', 'Obama1', 'Obama2', 'Nadella', 'McStay'], default='May') ) -> Path: ############################### I/O Settings ############################## # load config files opt = self.parser.parse_args('') opt.driving_audio = str(driving_audio) opt.id = talking_head with open(join('config', opt.id + '.yaml')) as f: config = yaml.safe_load(f) data_root = join('data', opt.id) ############################ Hyper Parameters ############################# h, w, sr, FPS = 512, 512, 16000, 60 mouth_indices = np.concatenate([np.arange(4, 11), np.arange(46, 64)]) eye_brow_indices = [27, 65, 28, 68, 29, 67, 30, 66, 31, 72, 32, 69, 33, 70, 34, 71] eye_brow_indices = np.array(eye_brow_indices, np.int32) ############################ Pre-defined Data ############################# mean_pts3d = np.load(join(data_root, 'mean_pts3d.npy')) fit_data = np.load(config['dataset_params']['fit_data_path']) pts3d = np.load(config['dataset_params']['pts3d_path']) - mean_pts3d trans = fit_data['trans'][:, :, 0].astype(np.float32) mean_translation = trans.mean(axis=0) candidate_eye_brow = pts3d[10:, eye_brow_indices] std_mean_pts3d = np.load(config['dataset_params']['pts3d_path']).mean(axis=0) # candidates images img_candidates = [] for j in range(4): output = imread(join(data_root, 'candidates', f'normalized_full_{j}.jpg')) output = A.pytorch.transforms.ToTensor(normalize={'mean': (0.5, 0.5, 0.5), 'std': (0.5, 0.5, 0.5)})(image=output)['image'] img_candidates.append(output) img_candidates = torch.cat(img_candidates).unsqueeze(0).cuda() # shoulders shoulders = np.load(join(data_root, 'normalized_shoulder_points.npy')) shoulder3D = np.load(join(data_root, 'shoulder_points3D.npy'))[1] ref_trans = trans[1] # camera matrix, we always use training set intrinsic parameters. camera = utils.camera() camera_intrinsic = np.load(join(data_root, 'camera_intrinsic.npy')).astype(np.float32) APC_feat_database = np.load(join(data_root, 'APC_feature_base.npy')) # load reconstruction data scale = sio.loadmat(join(data_root, 'id_scale.mat'))['scale'][0, 0] Audio2Mel_torch = audio_funcs.Audio2Mel(n_fft=512, hop_length=int(16000 / 120), win_length=int(16000 / 60), sampling_rate=16000, n_mel_channels=80, mel_fmin=90, mel_fmax=7600.0).cuda() ########################### Experiment Settings ########################### #### user config use_LLE = config['model_params']['APC']['use_LLE'] Knear = config['model_params']['APC']['Knear'] LLE_percent = config['model_params']['APC']['LLE_percent'] headpose_sigma = config['model_params']['Headpose']['sigma'] Feat_smooth_sigma = config['model_params']['Audio2Mouth']['smooth'] Head_smooth_sigma = config['model_params']['Headpose']['smooth'] Feat_center_smooth_sigma, Head_center_smooth_sigma = 0, 0 AMP_method = config['model_params']['Audio2Mouth']['AMP'][0] Feat_AMPs = config['model_params']['Audio2Mouth']['AMP'][1:] rot_AMP, trans_AMP = config['model_params']['Headpose']['AMP'] shoulder_AMP = config['model_params']['Headpose']['shoulder_AMP'] save_feature_maps = config['model_params']['Image2Image']['save_input'] #### common settings Featopt = FeatureOptions().parse() Headopt = HeadposeOptions().parse() Renderopt = RenderOptions().parse() Featopt.load_epoch = config['model_params']['Audio2Mouth']['ckp_path'] Headopt.load_epoch = config['model_params']['Headpose']['ckp_path'] Renderopt.dataroot = config['dataset_params']['root'] Renderopt.load_epoch = config['model_params']['Image2Image']['ckp_path'] Renderopt.size = config['model_params']['Image2Image']['size'] ############################# Load Models ################################# print('---------- Loading Model: APC-------------') APC_model = APC_encoder(config['model_params']['APC']['mel_dim'], config['model_params']['APC']['hidden_size'], config['model_params']['APC']['num_layers'], config['model_params']['APC']['residual']) # load all 5 here? APC_model.load_state_dict(torch.load(config['model_params']['APC']['ckp_path']), strict=False) APC_model.cuda() APC_model.eval() print('---------- Loading Model: {} -------------'.format(Featopt.task)) Audio2Feature = create_model(Featopt) Audio2Feature.setup(Featopt) Audio2Feature.eval() print('---------- Loading Model: {} -------------'.format(Headopt.task)) Audio2Headpose = create_model(Headopt) Audio2Headpose.setup(Headopt) Audio2Headpose.eval() if Headopt.feature_decoder == 'WaveNet': Headopt.A2H_receptive_field = Audio2Headpose.Audio2Headpose.module.WaveNet.receptive_field print('---------- Loading Model: {} -------------'.format(Renderopt.task)) facedataset = create_dataset(Renderopt) Feature2Face = create_model(Renderopt) Feature2Face.setup(Renderopt) Feature2Face.eval() visualizer = Visualizer(Renderopt) # check audio duration and trim audio extension_name = os.path.basename(opt.driving_audio).split('.')[-1] audio_threshold = 10 duration = librosa.get_duration(filename=opt.driving_audio) if duration > audio_threshold: print(f'audio file is longer than {audio_threshold} seconds, trimming the first {audio_threshold} seconds ' f'for further processing') ffmpeg_extract_subclip(opt.driving_audio, 0, audio_threshold, targetname=f'shorter_input.{extension_name}') opt.driving_audio = f'shorter_input.{extension_name}' # create the results folder audio_name = os.path.basename(opt.driving_audio).split('.')[0] save_root = join('results', opt.id, audio_name) os.makedirs(save_root, exist_ok=True) clean_folder(save_root) out_path = Path(tempfile.mkdtemp()) / "out.mp4" ############################## Inference ################################## print('Processing audio: {} ...'.format(audio_name)) # read audio audio, _ = librosa.load(opt.driving_audio, sr=sr) total_frames = np.int32(audio.shape[0] / sr * FPS) #### 1. compute APC features print('1. Computing APC features...') mel80 = utils.compute_mel_one_sequence(audio) mel_nframe = mel80.shape[0] with torch.no_grad(): length = torch.Tensor([mel_nframe]) mel80_torch = torch.from_numpy(mel80.astype(np.float32)).cuda().unsqueeze(0) hidden_reps = APC_model.forward(mel80_torch, length)[0] # [mel_nframe, 512] hidden_reps = hidden_reps.cpu().numpy() audio_feats = hidden_reps #### 2. manifold projection if use_LLE: print('2. Manifold projection...') ind = utils.KNN_with_torch(audio_feats, APC_feat_database, K=Knear) weights, feat_fuse = utils.compute_LLE_projection_all_frame(audio_feats, APC_feat_database, ind, audio_feats.shape[0]) audio_feats = audio_feats * (1 - LLE_percent) + feat_fuse * LLE_percent #### 3. Audio2Mouth print('3. Audio2Mouth inference...') pred_Feat = Audio2Feature.generate_sequences(audio_feats, sr, FPS, fill_zero=True, opt=Featopt) #### 4. Audio2Headpose print('4. Headpose inference...') # set history headposes as zero pre_headpose = np.zeros(Headopt.A2H_wavenet_input_channels, np.float32) pred_Head = Audio2Headpose.generate_sequences(audio_feats, pre_headpose, fill_zero=True, sigma_scale=0.3, opt=Headopt) #### 5. Post-Processing print('5. Post-processing...') nframe = min(pred_Feat.shape[0], pred_Head.shape[0]) pred_pts3d = np.zeros([nframe, 73, 3]) pred_pts3d[:, mouth_indices] = pred_Feat.reshape(-1, 25, 3)[:nframe] ## mouth pred_pts3d = utils.landmark_smooth_3d(pred_pts3d, Feat_smooth_sigma, area='only_mouth') pred_pts3d = utils.mouth_pts_AMP(pred_pts3d, True, AMP_method, Feat_AMPs) pred_pts3d = pred_pts3d + mean_pts3d pred_pts3d = utils.solve_intersect_mouth(pred_pts3d) # solve intersect lips if exist ## headpose pred_Head[:, 0:3] *= rot_AMP pred_Head[:, 3:6] *= trans_AMP pred_headpose = utils.headpose_smooth(pred_Head[:, :6], Head_smooth_sigma).astype(np.float32) pred_headpose[:, 3:] += mean_translation pred_headpose[:, 0] += 180 ## compute projected landmarks pred_landmarks = np.zeros([nframe, 73, 2], dtype=np.float32) final_pts3d = np.zeros([nframe, 73, 3], dtype=np.float32) final_pts3d[:] = std_mean_pts3d.copy() final_pts3d[:, 46:64] = pred_pts3d[:nframe, 46:64] for k in tqdm(range(nframe)): ind = k % candidate_eye_brow.shape[0] final_pts3d[k, eye_brow_indices] = candidate_eye_brow[ind] + mean_pts3d[eye_brow_indices] pred_landmarks[k], _, _ = utils.project_landmarks(camera_intrinsic, camera.relative_rotation, camera.relative_translation, scale, pred_headpose[k], final_pts3d[k]) ## Upper Body Motion pred_shoulders = np.zeros([nframe, 18, 2], dtype=np.float32) pred_shoulders3D = np.zeros([nframe, 18, 3], dtype=np.float32) for k in range(nframe): diff_trans = pred_headpose[k][3:] - ref_trans pred_shoulders3D[k] = shoulder3D + diff_trans * shoulder_AMP # project project = camera_intrinsic.dot(pred_shoulders3D[k].T) project[:2, :] /= project[2, :] # divide z pred_shoulders[k] = project[:2, :].T #### 6. Image2Image translation & Save resuls print('6. Image2Image translation & Saving results...') for ind in tqdm(range(0, nframe), desc='Image2Image translation inference'): # feature_map: [input_nc, h, w] current_pred_feature_map = facedataset.dataset.get_data_test_mode(pred_landmarks[ind], pred_shoulders[ind], facedataset.dataset.image_pad) input_feature_maps = current_pred_feature_map.unsqueeze(0).cuda() pred_fake = Feature2Face.inference(input_feature_maps, img_candidates) # save results visual_list = [('pred', util.tensor2im(pred_fake[0]))] if save_feature_maps: visual_list += [('input', np.uint8(current_pred_feature_map[0].cpu().numpy() * 255))] visuals = OrderedDict(visual_list) visualizer.save_images(save_root, visuals, str(ind + 1)) ## make videos # generate corresponding audio, reused for all results tmp_audio_path = join(save_root, 'tmp.wav') tmp_audio_clip = audio[: np.int32(nframe * sr / FPS)] librosa.output.write_wav(tmp_audio_path, tmp_audio_clip, sr) def write_video_with_audio(audio_path, output_path, prefix='pred_'): fps, fourcc = 60, cv2.VideoWriter_fourcc(*'DIVX') video_tmp_path = join(save_root, 'tmp.avi') out = cv2.VideoWriter(video_tmp_path, fourcc, fps, (Renderopt.loadSize, Renderopt.loadSize)) for j in tqdm(range(nframe), position=0, desc='writing video'): img = cv2.imread(join(save_root, prefix + str(j + 1) + '.jpg')) out.write(img) out.release() cmd = 'ffmpeg -i "' + video_tmp_path + '" -i "' + audio_path + '" -codec copy -shortest "' + output_path + '"' subprocess.call(cmd, shell=True) os.remove(video_tmp_path) # remove the template video temp_out = 'temp_video.avi' write_video_with_audio(tmp_audio_path, temp_out, 'pred_') # convert to mp4 cmd = ("ffmpeg -i " + temp_out + " -strict -2 " + str(out_path) ) subprocess.call(cmd, shell=True) if os.path.exists(tmp_audio_path): os.remove(tmp_audio_path) if os.path.exists(temp_out): os.remove(temp_out) if os.path.exists(f'shorter_input.{extension_name}'): os.remove(f'shorter_input.{extension_name}') if not opt.save_intermediates: _img_paths = list(map(lambda x: str(x), list(Path(save_root).glob('*.jpg')))) for i in tqdm(range(len(_img_paths)), desc='deleting intermediate images'): os.remove(_img_paths[i]) print('Finish!') return out_path def clean_folder(folder): for filename in os.listdir(folder): file_path = os.path.join(folder, filename) try: if os.path.isfile(file_path) or os.path.islink(file_path): os.unlink(file_path) elif os.path.isdir(file_path): shutil.rmtree(file_path) except Exception as e: print('Failed to delete %s. Reason: %s' % (file_path, e))