Shellbrady/LivePortrait5
0
1# coding: utf-82 3"""4Pipeline of LivePortrait5"""6 7import torch8torch.backends.cudnn.benchmark = True # disable CUDNN_BACKEND_EXECUTION_PLAN_DESCRIPTOR warning9 10import cv211import numpy as np12import pickle13import os14import os.path as osp15from rich.progress import track16 17from .config.argument_config import ArgumentConfig18from .config.inference_config import InferenceConfig19from .config.crop_config import CropConfig20from .utils.cropper import Cropper21from .utils.camera import get_rotation_matrix22from .utils.video import images2video, concat_frames, get_fps, add_audio_to_video, has_audio_stream23from .utils.crop import _transform_img, prepare_paste_back, paste_back24from .utils.retargeting_utils import calc_lip_close_ratio25from .utils.io import load_image_rgb, load_driving_info, resize_to_limit26from .utils.helper import mkdir, basename, dct2cuda, is_video, is_template27from .utils.rprint import rlog as log28from .live_portrait_wrapper import LivePortraitWrapper29 30 31def make_abs_path(fn):32 return osp.join(osp.dirname(osp.realpath(__file__)), fn)33 34 35class LivePortraitPipeline(object):36 37 def __init__(self, inference_cfg: InferenceConfig, crop_cfg: CropConfig):38 self.live_portrait_wrapper: LivePortraitWrapper = LivePortraitWrapper(cfg=inference_cfg)39 self.cropper = Cropper(crop_cfg=crop_cfg)40 41 def execute(self, args: ArgumentConfig):42 inference_cfg = self.live_portrait_wrapper.cfg # for convenience43 ######## process source portrait ########44 img_rgb = load_image_rgb(args.source_image)45 img_rgb = resize_to_limit(img_rgb, inference_cfg.ref_max_shape, inference_cfg.ref_shape_n)46 log(f"Load source image from {args.source_image}")47 crop_info = self.cropper.crop_single_image(img_rgb)48 source_lmk = crop_info['lmk_crop']49 img_crop, img_crop_256x256 = crop_info['img_crop'], crop_info['img_crop_256x256']50 if inference_cfg.flag_do_crop:51 I_s = self.live_portrait_wrapper.prepare_source(img_crop_256x256)52 else:53 I_s = self.live_portrait_wrapper.prepare_source(img_rgb)54 x_s_info = self.live_portrait_wrapper.get_kp_info(I_s)55 x_c_s = x_s_info['kp']56 R_s = get_rotation_matrix(x_s_info['pitch'], x_s_info['yaw'], x_s_info['roll'])57 f_s = self.live_portrait_wrapper.extract_feature_3d(I_s)58 x_s = self.live_portrait_wrapper.transform_keypoint(x_s_info)59 60 if inference_cfg.flag_lip_zero:61 # let lip-open scalar to be 0 at first62 c_d_lip_before_animation = [0.]63 combined_lip_ratio_tensor_before_animation = self.live_portrait_wrapper.calc_combined_lip_ratio(c_d_lip_before_animation, source_lmk)64 if combined_lip_ratio_tensor_before_animation[0][0] < inference_cfg.lip_zero_threshold:65 inference_cfg.flag_lip_zero = False66 else:67 lip_delta_before_animation = self.live_portrait_wrapper.retarget_lip(x_s, combined_lip_ratio_tensor_before_animation)68 ############################################69 70 ######## process driving info ########71 output_fps = 30 # default fps72 if is_video(args.driving_info):73 log(f"Load from video file (mp4 mov avi etc...): {args.driving_info}")74 output_fps = int(get_fps(args.driving_info))75 log(f'The FPS of {args.driving_info} is: {output_fps}')76 77 # TODO: 这里track一下驱动视频 -> 构建模板78 driving_rgb_lst = load_driving_info(args.driving_info)79 driving_rgb_lst_256 = [cv2.resize(_, (256, 256)) for _ in driving_rgb_lst]80 I_d_lst = self.live_portrait_wrapper.prepare_driving_videos(driving_rgb_lst_256)81 n_frames = I_d_lst.shape[0]82 if inference_cfg.flag_eye_retargeting or inference_cfg.flag_lip_retargeting:83 driving_lmk_lst = self.cropper.get_retargeting_lmk_info(driving_rgb_lst)84 input_eye_ratio_lst, input_lip_ratio_lst = self.live_portrait_wrapper.calc_retargeting_ratio(source_lmk, driving_lmk_lst)85 elif is_template(args.driving_info):86 log(f"Load from video templates {args.driving_info}")87 with open(args.driving_info, 'rb') as f:88 template_lst, driving_lmk_lst = pickle.load(f)89 n_frames = template_lst[0]['n_frames']90 input_eye_ratio_lst, input_lip_ratio_lst = self.live_portrait_wrapper.calc_retargeting_ratio(source_lmk, driving_lmk_lst)91 else:92 raise Exception("Unsupported driving types!")93 #########################################94 95 ######## prepare for pasteback ########96 if inference_cfg.flag_pasteback:97 mask_ori = prepare_paste_back(inference_cfg.mask_crop, crop_info['M_c2o'], dsize=(img_rgb.shape[1], img_rgb.shape[0]))98 I_p_paste_lst = []99 #########################################100 101 I_p_lst = []102 R_d_0, x_d_0_info = None, None103 for i in track(range(n_frames), description='Animating...', total=n_frames):104 if is_video(args.driving_info):105 # extract kp info by M106 I_d_i = I_d_lst[i]107 x_d_i_info = self.live_portrait_wrapper.get_kp_info(I_d_i)108 R_d_i = get_rotation_matrix(x_d_i_info['pitch'], x_d_i_info['yaw'], x_d_i_info['roll'])109 else:110 # from template111 x_d_i_info = template_lst[i]112 x_d_i_info = dct2cuda(x_d_i_info, inference_cfg.device_id)113 R_d_i = x_d_i_info['R_d']114 115 if i == 0:116 R_d_0 = R_d_i117 x_d_0_info = x_d_i_info118 119 if inference_cfg.flag_relative:120 R_new = (R_d_i @ R_d_0.permute(0, 2, 1)) @ R_s121 delta_new = x_s_info['exp'] + (x_d_i_info['exp'] - x_d_0_info['exp'])122 scale_new = x_s_info['scale'] * (x_d_i_info['scale'] / x_d_0_info['scale'])123 t_new = x_s_info['t'] + (x_d_i_info['t'] - x_d_0_info['t'])124 else:125 R_new = R_d_i126 delta_new = x_d_i_info['exp']127 scale_new = x_s_info['scale']128 t_new = x_d_i_info['t']129 130 t_new[..., 2].fill_(0) # zero tz131 x_d_i_new = scale_new * (x_c_s @ R_new + delta_new) + t_new132 133 # Algorithm 1:134 if not inference_cfg.flag_stitching and not inference_cfg.flag_eye_retargeting and not inference_cfg.flag_lip_retargeting:135 # without stitching or retargeting136 if inference_cfg.flag_lip_zero:137 x_d_i_new += lip_delta_before_animation.reshape(-1, x_s.shape[1], 3)138 else:139 pass140 elif inference_cfg.flag_stitching and not inference_cfg.flag_eye_retargeting and not inference_cfg.flag_lip_retargeting:141 # with stitching and without retargeting142 if inference_cfg.flag_lip_zero:143 x_d_i_new = self.live_portrait_wrapper.stitching(x_s, x_d_i_new) + lip_delta_before_animation.reshape(-1, x_s.shape[1], 3)144 else:145 x_d_i_new = self.live_portrait_wrapper.stitching(x_s, x_d_i_new)146 else:147 eyes_delta, lip_delta = None, None148 if inference_cfg.flag_eye_retargeting:149 c_d_eyes_i = input_eye_ratio_lst[i]150 combined_eye_ratio_tensor = self.live_portrait_wrapper.calc_combined_eye_ratio(c_d_eyes_i, source_lmk)151 # ∆_eyes,i = R_eyes(x_s; c_s,eyes, c_d,eyes,i)152 eyes_delta = self.live_portrait_wrapper.retarget_eye(x_s, combined_eye_ratio_tensor)153 if inference_cfg.flag_lip_retargeting:154 c_d_lip_i = input_lip_ratio_lst[i]155 combined_lip_ratio_tensor = self.live_portrait_wrapper.calc_combined_lip_ratio(c_d_lip_i, source_lmk)156 # ∆_lip,i = R_lip(x_s; c_s,lip, c_d,lip,i)157 lip_delta = self.live_portrait_wrapper.retarget_lip(x_s, combined_lip_ratio_tensor)158 159 if inference_cfg.flag_relative: # use x_s160 x_d_i_new = x_s + \161 (eyes_delta.reshape(-1, x_s.shape[1], 3) if eyes_delta is not None else 0) + \162 (lip_delta.reshape(-1, x_s.shape[1], 3) if lip_delta is not None else 0)163 else: # use x_d,i164 x_d_i_new = x_d_i_new + \165 (eyes_delta.reshape(-1, x_s.shape[1], 3) if eyes_delta is not None else 0) + \166 (lip_delta.reshape(-1, x_s.shape[1], 3) if lip_delta is not None else 0)167 168 if inference_cfg.flag_stitching:169 x_d_i_new = self.live_portrait_wrapper.stitching(x_s, x_d_i_new)170 171 out = self.live_portrait_wrapper.warp_decode(f_s, x_s, x_d_i_new)172 I_p_i = self.live_portrait_wrapper.parse_output(out['out'])[0]173 I_p_lst.append(I_p_i)174 175 if inference_cfg.flag_pasteback:176 I_p_i_to_ori_blend = paste_back(I_p_i, crop_info['M_c2o'], img_rgb, mask_ori)177 I_p_paste_lst.append(I_p_i_to_ori_blend)178 179 mkdir(args.output_dir)180 wfp_concat = None181 flag_has_audio = has_audio_stream(args.driving_info)182 183 if is_video(args.driving_info):184 frames_concatenated = concat_frames(I_p_lst, driving_rgb_lst, img_crop_256x256)185 # save (driving frames, source image, drived frames) result186 wfp_concat = osp.join(args.output_dir, f'{basename(args.source_image)}--{basename(args.driving_info)}_concat.mp4')187 images2video(frames_concatenated, wfp=wfp_concat, fps=output_fps)188 if flag_has_audio:189 # final result with concat190 wfp_concat_with_audio = osp.join(args.output_dir, f'{basename(args.source_image)}--{basename(args.driving_info)}_concat_with_audio.mp4')191 add_audio_to_video(wfp_concat, args.driving_info, wfp_concat_with_audio)192 os.replace(wfp_concat_with_audio, wfp_concat)193 log(f"Replace {wfp_concat} with {wfp_concat_with_audio}")194 195 # save drived result196 wfp = osp.join(args.output_dir, f'{basename(args.source_image)}--{basename(args.driving_info)}.mp4')197 if inference_cfg.flag_pasteback:198 images2video(I_p_paste_lst, wfp=wfp, fps=output_fps)199 else:200 images2video(I_p_lst, wfp=wfp, fps=output_fps)201 202 ######### build final result #########203 if flag_has_audio:204 wfp_with_audio = osp.join(args.output_dir, f'{basename(args.source_image)}--{basename(args.driving_info)}_with_audio.mp4')205 add_audio_to_video(wfp, args.driving_info, wfp_with_audio)206 os.replace(wfp_with_audio, wfp)207 log(f"Replace {wfp} with {wfp_with_audio}")208 209 return wfp, wfp_concat210