Shellbrady/LivePortrait5

1# coding: utf-82 3"""4Make video template5"""6 7import os8import cv29import numpy as np10import pickle11from rich.progress import track12from .utils.cropper import Cropper13 14from .utils.io import load_driving_info15from .utils.camera import get_rotation_matrix16from .utils.helper import mkdir, basename17from .utils.rprint import rlog as log18from .config.crop_config import CropConfig19from .config.inference_config import InferenceConfig20from .live_portrait_wrapper import LivePortraitWrapper21 22class TemplateMaker:23 24    def __init__(self, inference_cfg: InferenceConfig, crop_cfg: CropConfig):25        self.live_portrait_wrapper: LivePortraitWrapper = LivePortraitWrapper(cfg=inference_cfg)26        self.cropper = Cropper(crop_cfg=crop_cfg)27 28    def make_motion_template(self, video_fp: str, output_path: str, **kwargs):29        """ make video template (.pkl format)30        video_fp: driving video file path31        output_path: where to save the pickle file32        """33 34        driving_rgb_lst = load_driving_info(video_fp)35        driving_rgb_lst = [cv2.resize(_, (256, 256)) for _ in driving_rgb_lst]36        driving_lmk_lst = self.cropper.get_retargeting_lmk_info(driving_rgb_lst)37        I_d_lst = self.live_portrait_wrapper.prepare_driving_videos(driving_rgb_lst)38 39        n_frames = I_d_lst.shape[0]40 41        templates = []42 43 44        for i in track(range(n_frames), description='Making templates...', total=n_frames):45            I_d_i = I_d_lst[i]46            x_d_i_info = self.live_portrait_wrapper.get_kp_info(I_d_i)47            R_d_i = get_rotation_matrix(x_d_i_info['pitch'], x_d_i_info['yaw'], x_d_i_info['roll'])48            # collect s_d, R_d, δ_d and t_d for inference49            template_dct = {50                'n_frames': n_frames,51                'frames_index': i,52            }53            template_dct['scale'] = x_d_i_info['scale'].cpu().numpy().astype(np.float32)54            template_dct['R_d'] = R_d_i.cpu().numpy().astype(np.float32)55            template_dct['exp'] = x_d_i_info['exp'].cpu().numpy().astype(np.float32)56            template_dct['t'] = x_d_i_info['t'].cpu().numpy().astype(np.float32)57 58            templates.append(template_dct)59 60        mkdir(output_path)61        # Save the dictionary as a pickle file62        pickle_fp = os.path.join(output_path, f'{basename(video_fp)}.pkl')63        with open(pickle_fp, 'wb') as f:64            pickle.dump([templates, driving_lmk_lst], f)65        log(f"Template saved at {pickle_fp}")66