predict.py

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))