RabbitRUI/ruispace

1import os 2import torch3import numpy as np4from scipy.io import savemat, loadmat5from yacs.config import CfgNode as CN6from scipy.signal import savgol_filter7 8from src.audio2pose_models.audio2pose import Audio2Pose9from src.audio2exp_models.networks import SimpleWrapperV2 10from src.audio2exp_models.audio2exp import Audio2Exp  11 12def load_cpk(checkpoint_path, model=None, optimizer=None, device="cpu"):13    checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))14    if model is not None:15        model.load_state_dict(checkpoint['model'])16    if optimizer is not None:17        optimizer.load_state_dict(checkpoint['optimizer'])18 19    return checkpoint['epoch']20 21class Audio2Coeff():22 23    def __init__(self, audio2pose_checkpoint, audio2pose_yaml_path, 24                        audio2exp_checkpoint, audio2exp_yaml_path, 25                        wav2lip_checkpoint, device):26        #load config27        fcfg_pose = open(audio2pose_yaml_path)28        cfg_pose = CN.load_cfg(fcfg_pose)29        cfg_pose.freeze()30        fcfg_exp = open(audio2exp_yaml_path)31        cfg_exp = CN.load_cfg(fcfg_exp)32        cfg_exp.freeze()33 34        # load audio2pose_model35        self.audio2pose_model = Audio2Pose(cfg_pose, wav2lip_checkpoint, device=device)36        self.audio2pose_model = self.audio2pose_model.to(device)37        self.audio2pose_model.eval()38        for param in self.audio2pose_model.parameters():39            param.requires_grad = False 40        try:41            load_cpk(audio2pose_checkpoint, model=self.audio2pose_model, device=device)42        except:43            raise Exception("Failed in loading audio2pose_checkpoint")44 45        # load audio2exp_model46        netG = SimpleWrapperV2()47        netG = netG.to(device)48        for param in netG.parameters():49            netG.requires_grad = False50        netG.eval()51        try:52            load_cpk(audio2exp_checkpoint, model=netG, device=device)53        except:54            raise Exception("Failed in loading audio2exp_checkpoint")55        self.audio2exp_model = Audio2Exp(netG, cfg_exp, device=device, prepare_training_loss=False)56        self.audio2exp_model = self.audio2exp_model.to(device)57        for param in self.audio2exp_model.parameters():58            param.requires_grad = False59        self.audio2exp_model.eval()60 61        self.device = device62 63    def generate(self, batch, coeff_save_dir, pose_style, ref_pose_coeff_path=None):64 65        with torch.no_grad():66            #test67            results_dict_exp= self.audio2exp_model.test(batch)68            exp_pred = results_dict_exp['exp_coeff_pred']                         #bs T 6469 70            #for class_id in  range(1):71            #class_id = 0#(i+10)%4572            #class_id = random.randint(0,46)                                   #46 styles can be selected 73            batch['class'] = torch.LongTensor([pose_style]).to(self.device)74            results_dict_pose = self.audio2pose_model.test(batch) 75            pose_pred = results_dict_pose['pose_pred']                        #bs T 676 77            pose_len = pose_pred.shape[1]78            if pose_len<13: 79                pose_len = int((pose_len-1)/2)*2+180                pose_pred = torch.Tensor(savgol_filter(np.array(pose_pred.cpu()), pose_len, 2, axis=1)).to(self.device)81            else:82                pose_pred = torch.Tensor(savgol_filter(np.array(pose_pred.cpu()), 13, 2, axis=1)).to(self.device) 83            84            coeffs_pred = torch.cat((exp_pred, pose_pred), dim=-1)            #bs T 7085 86            coeffs_pred_numpy = coeffs_pred[0].clone().detach().cpu().numpy() 87 88            89            if ref_pose_coeff_path is not None: 90                 coeffs_pred_numpy = self.using_refpose(coeffs_pred_numpy, ref_pose_coeff_path)91        92            savemat(os.path.join(coeff_save_dir, '%s##%s.mat'%(batch['pic_name'], batch['audio_name'])),  93                    {'coeff_3dmm': coeffs_pred_numpy})94 95            return os.path.join(coeff_save_dir, '%s##%s.mat'%(batch['pic_name'], batch['audio_name']))96    97    def using_refpose(self, coeffs_pred_numpy, ref_pose_coeff_path):98        num_frames = coeffs_pred_numpy.shape[0]99        refpose_coeff_dict = loadmat(ref_pose_coeff_path)100        refpose_coeff = refpose_coeff_dict['coeff_3dmm'][:,64:70]101        refpose_num_frames = refpose_coeff.shape[0]102        if refpose_num_frames<num_frames:103            div = num_frames//refpose_num_frames104            re = num_frames%refpose_num_frames105            refpose_coeff_list = [refpose_coeff for i in range(div)]106            refpose_coeff_list.append(refpose_coeff[:re, :])107            refpose_coeff = np.concatenate(refpose_coeff_list, axis=0)108 109        coeffs_pred_numpy[:, 64:70] = refpose_coeff[:num_frames, :] 110        return coeffs_pred_numpy111 112 113