Shellbrady/LivePortrait5

1# coding: utf-82 3"""4Wrapper for LivePortrait core functions5"""6 7import os.path as osp8import numpy as np9import cv210import torch11import yaml12 13from .utils.timer import Timer14from .utils.helper import load_model, concat_feat15from .utils.camera import headpose_pred_to_degree, get_rotation_matrix16from .utils.retargeting_utils import calc_eye_close_ratio, calc_lip_close_ratio17from .config.inference_config import InferenceConfig18from .utils.rprint import rlog as log19 20 21class LivePortraitWrapper(object):22 23    def __init__(self, cfg: InferenceConfig):24 25        model_config = yaml.load(open(cfg.models_config, 'r'), Loader=yaml.SafeLoader)26 27        # init F28        self.appearance_feature_extractor = load_model(cfg.checkpoint_F, model_config, cfg.device_id, 'appearance_feature_extractor')29        log(f'Load appearance_feature_extractor done.')30        # init M31        self.motion_extractor = load_model(cfg.checkpoint_M, model_config, cfg.device_id, 'motion_extractor')32        log(f'Load motion_extractor done.')33        # init W34        self.warping_module = load_model(cfg.checkpoint_W, model_config, cfg.device_id, 'warping_module')35        log(f'Load warping_module done.')36        # init G37        self.spade_generator = load_model(cfg.checkpoint_G, model_config, cfg.device_id, 'spade_generator')38        log(f'Load spade_generator done.')39        # init S and R40        if cfg.checkpoint_S is not None and osp.exists(cfg.checkpoint_S):41            self.stitching_retargeting_module = load_model(cfg.checkpoint_S, model_config, cfg.device_id, 'stitching_retargeting_module')42            log(f'Load stitching_retargeting_module done.')43        else:44            self.stitching_retargeting_module = None45 46        self.cfg = cfg47        self.device_id = cfg.device_id48        self.timer = Timer()49 50    def update_config(self, user_args):51        for k, v in user_args.items():52            if hasattr(self.cfg, k):53                setattr(self.cfg, k, v)54 55    def prepare_source(self, img: np.ndarray) -> torch.Tensor:56        """ construct the input as standard57        img: HxWx3, uint8, 256x25658        """59        h, w = img.shape[:2]60        if h != self.cfg.input_shape[0] or w != self.cfg.input_shape[1]:61            x = cv2.resize(img, (self.cfg.input_shape[0], self.cfg.input_shape[1]))62        else:63            x = img.copy()64 65        if x.ndim == 3:66            x = x[np.newaxis].astype(np.float32) / 255.  # HxWx3 -> 1xHxWx3, normalized to 0~167        elif x.ndim == 4:68            x = x.astype(np.float32) / 255.  # BxHxWx3, normalized to 0~169        else:70            raise ValueError(f'img ndim should be 3 or 4: {x.ndim}')71        x = np.clip(x, 0, 1)  # clip to 0~172        x = torch.from_numpy(x).permute(0, 3, 1, 2)  # 1xHxWx3 -> 1x3xHxW73        x = x.cuda(self.device_id)74        return x75 76    def prepare_driving_videos(self, imgs) -> torch.Tensor:77        """ construct the input as standard78        imgs: NxBxHxWx3, uint879        """80        if isinstance(imgs, list):81            _imgs = np.array(imgs)[..., np.newaxis]  # TxHxWx3x182        elif isinstance(imgs, np.ndarray):83            _imgs = imgs84        else:85            raise ValueError(f'imgs type error: {type(imgs)}')86 87        y = _imgs.astype(np.float32) / 255.88        y = np.clip(y, 0, 1)  # clip to 0~189        y = torch.from_numpy(y).permute(0, 4, 3, 1, 2)  # TxHxWx3x1 -> Tx1x3xHxW90        y = y.cuda(self.device_id)91 92        return y93 94    def extract_feature_3d(self, x: torch.Tensor) -> torch.Tensor:95        """ get the appearance feature of the image by F96        x: Bx3xHxW, normalized to 0~197        """98        with torch.no_grad():99            with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=self.cfg.flag_use_half_precision):100                feature_3d = self.appearance_feature_extractor(x)101 102        return feature_3d.float()103 104    def get_kp_info(self, x: torch.Tensor, **kwargs) -> dict:105        """ get the implicit keypoint information106        x: Bx3xHxW, normalized to 0~1107        flag_refine_info: whether to trandform the pose to degrees and the dimention of the reshape108        return: A dict contains keys: 'pitch', 'yaw', 'roll', 't', 'exp', 'scale', 'kp'109        """110        with torch.no_grad():111            with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=self.cfg.flag_use_half_precision):112                kp_info = self.motion_extractor(x)113 114            if self.cfg.flag_use_half_precision:115                # float the dict116                for k, v in kp_info.items():117                    if isinstance(v, torch.Tensor):118                        kp_info[k] = v.float()119 120        flag_refine_info: bool = kwargs.get('flag_refine_info', True)121        if flag_refine_info:122            bs = kp_info['kp'].shape[0]123            kp_info['pitch'] = headpose_pred_to_degree(kp_info['pitch'])[:, None]  # Bx1124            kp_info['yaw'] = headpose_pred_to_degree(kp_info['yaw'])[:, None]  # Bx1125            kp_info['roll'] = headpose_pred_to_degree(kp_info['roll'])[:, None]  # Bx1126            kp_info['kp'] = kp_info['kp'].reshape(bs, -1, 3)  # BxNx3127            kp_info['exp'] = kp_info['exp'].reshape(bs, -1, 3)  # BxNx3128 129        return kp_info130 131    def get_pose_dct(self, kp_info: dict) -> dict:132        pose_dct = dict(133            pitch=headpose_pred_to_degree(kp_info['pitch']).item(),134            yaw=headpose_pred_to_degree(kp_info['yaw']).item(),135            roll=headpose_pred_to_degree(kp_info['roll']).item(),136        )137        return pose_dct138 139    def get_fs_and_kp_info(self, source_prepared, driving_first_frame):140 141        # get the canonical keypoints of source image by M142        source_kp_info = self.get_kp_info(source_prepared, flag_refine_info=True)143        source_rotation = get_rotation_matrix(source_kp_info['pitch'], source_kp_info['yaw'], source_kp_info['roll'])144 145        # get the canonical keypoints of first driving frame by M146        driving_first_frame_kp_info = self.get_kp_info(driving_first_frame, flag_refine_info=True)147        driving_first_frame_rotation = get_rotation_matrix(148            driving_first_frame_kp_info['pitch'],149            driving_first_frame_kp_info['yaw'],150            driving_first_frame_kp_info['roll']151        )152 153        # get feature volume by F154        source_feature_3d = self.extract_feature_3d(source_prepared)155 156        return source_kp_info, source_rotation, source_feature_3d, driving_first_frame_kp_info, driving_first_frame_rotation157 158    def transform_keypoint(self, kp_info: dict):159        """160        transform the implicit keypoints with the pose, shift, and expression deformation161        kp: BxNx3162        """163        kp = kp_info['kp']    # (bs, k, 3)164        pitch, yaw, roll = kp_info['pitch'], kp_info['yaw'], kp_info['roll']165 166        t, exp = kp_info['t'], kp_info['exp']167        scale = kp_info['scale']168 169        pitch = headpose_pred_to_degree(pitch)170        yaw = headpose_pred_to_degree(yaw)171        roll = headpose_pred_to_degree(roll)172 173        bs = kp.shape[0]174        if kp.ndim == 2:175            num_kp = kp.shape[1] // 3  # Bx(num_kpx3)176        else:177            num_kp = kp.shape[1]  # Bxnum_kpx3178 179        rot_mat = get_rotation_matrix(pitch, yaw, roll)    # (bs, 3, 3)180 181        # Eqn.2: s * (R * x_c,s + exp) + t182        kp_transformed = kp.view(bs, num_kp, 3) @ rot_mat + exp.view(bs, num_kp, 3)183        kp_transformed *= scale[..., None]  # (bs, k, 3) * (bs, 1, 1) = (bs, k, 3)184        kp_transformed[:, :, 0:2] += t[:, None, 0:2]  # remove z, only apply tx ty185 186        return kp_transformed187 188    def retarget_eye(self, kp_source: torch.Tensor, eye_close_ratio: torch.Tensor) -> torch.Tensor:189        """190        kp_source: BxNx3191        eye_close_ratio: Bx3192        Return: Bx(3*num_kp+2)193        """194        feat_eye = concat_feat(kp_source, eye_close_ratio)195 196        with torch.no_grad():197            delta = self.stitching_retargeting_module['eye'](feat_eye)198 199        return delta200 201    def retarget_lip(self, kp_source: torch.Tensor, lip_close_ratio: torch.Tensor) -> torch.Tensor:202        """203        kp_source: BxNx3204        lip_close_ratio: Bx2205        """206        feat_lip = concat_feat(kp_source, lip_close_ratio)207 208        with torch.no_grad():209            delta = self.stitching_retargeting_module['lip'](feat_lip)210 211        return delta212 213    def stitch(self, kp_source: torch.Tensor, kp_driving: torch.Tensor) -> torch.Tensor:214        """215        kp_source: BxNx3216        kp_driving: BxNx3217        Return: Bx(3*num_kp+2)218        """219        feat_stiching = concat_feat(kp_source, kp_driving)220 221        with torch.no_grad():222            delta = self.stitching_retargeting_module['stitching'](feat_stiching)223 224        return delta225 226    def stitching(self, kp_source: torch.Tensor, kp_driving: torch.Tensor) -> torch.Tensor:227        """ conduct the stitching228        kp_source: Bxnum_kpx3229        kp_driving: Bxnum_kpx3230        """231 232        if self.stitching_retargeting_module is not None:233 234            bs, num_kp = kp_source.shape[:2]235 236            kp_driving_new = kp_driving.clone()237            delta = self.stitch(kp_source, kp_driving_new)238 239            delta_exp = delta[..., :3*num_kp].reshape(bs, num_kp, 3)  # 1x20x3240            delta_tx_ty = delta[..., 3*num_kp:3*num_kp+2].reshape(bs, 1, 2)  # 1x1x2241 242            kp_driving_new += delta_exp243            kp_driving_new[..., :2] += delta_tx_ty244 245            return kp_driving_new246 247        return kp_driving248 249    def warp_decode(self, feature_3d: torch.Tensor, kp_source: torch.Tensor, kp_driving: torch.Tensor) -> torch.Tensor:250        """ get the image after the warping of the implicit keypoints251        feature_3d: Bx32x16x64x64, feature volume252        kp_source: BxNx3253        kp_driving: BxNx3254        """255        # The line 18 in Algorithm 1: D(W(f_s; x_s, x′_d,i)）256        with torch.no_grad():257            with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=self.cfg.flag_use_half_precision):258                # get decoder input259                ret_dct = self.warping_module(feature_3d, kp_source=kp_source, kp_driving=kp_driving)260                # decode261                ret_dct['out'] = self.spade_generator(feature=ret_dct['out'])262 263            # float the dict264            if self.cfg.flag_use_half_precision:265                for k, v in ret_dct.items():266                    if isinstance(v, torch.Tensor):267                        ret_dct[k] = v.float()268 269        return ret_dct270 271    def parse_output(self, out: torch.Tensor) -> np.ndarray:272        """ construct the output as standard273        return: 1xHxWx3, uint8274        """275        out = np.transpose(out.data.cpu().numpy(), [0, 2, 3, 1])  # 1x3xHxW -> 1xHxWx3276        out = np.clip(out, 0, 1)  # clip to 0~1277        out = np.clip(out * 255, 0, 255).astype(np.uint8)  # 0~1 -> 0~255278 279        return out280 281    def calc_retargeting_ratio(self, source_lmk, driving_lmk_lst):282        input_eye_ratio_lst = []283        input_lip_ratio_lst = []284        for lmk in driving_lmk_lst:285            # for eyes retargeting286            input_eye_ratio_lst.append(calc_eye_close_ratio(lmk[None]))287            # for lip retargeting288            input_lip_ratio_lst.append(calc_lip_close_ratio(lmk[None]))289        return input_eye_ratio_lst, input_lip_ratio_lst290 291    def calc_combined_eye_ratio(self, input_eye_ratio, source_lmk):292        eye_close_ratio = calc_eye_close_ratio(source_lmk[None])293        eye_close_ratio_tensor = torch.from_numpy(eye_close_ratio).float().cuda(self.device_id)294        input_eye_ratio_tensor = torch.Tensor([input_eye_ratio[0][0]]).reshape(1, 1).cuda(self.device_id)295        # [c_s,eyes, c_d,eyes,i]296        combined_eye_ratio_tensor = torch.cat([eye_close_ratio_tensor, input_eye_ratio_tensor], dim=1)297        return combined_eye_ratio_tensor298 299    def calc_combined_lip_ratio(self, input_lip_ratio, source_lmk):300        lip_close_ratio = calc_lip_close_ratio(source_lmk[None])301        lip_close_ratio_tensor = torch.from_numpy(lip_close_ratio).float().cuda(self.device_id)302        # [c_s,lip, c_d,lip,i]303        input_lip_ratio_tensor = torch.Tensor([input_lip_ratio[0]]).cuda(self.device_id)304        if input_lip_ratio_tensor.shape != [1, 1]:305            input_lip_ratio_tensor = input_lip_ratio_tensor.reshape(1, 1)306        combined_lip_ratio_tensor = torch.cat([lip_close_ratio_tensor, input_lip_ratio_tensor], dim=1)307        return combined_lip_ratio_tensor308