66 lines
2.4 KiB
Python
66 lines
2.4 KiB
Python
# coding: utf-8
|
|
|
|
"""
|
|
Make video template
|
|
"""
|
|
|
|
import os
|
|
import cv2
|
|
import numpy as np
|
|
import pickle
|
|
from rich.progress import track
|
|
from .utils.cropper import Cropper
|
|
|
|
from .utils.io import load_driving_info
|
|
from .utils.camera import get_rotation_matrix
|
|
from .utils.helper import mkdir, basename
|
|
from .utils.rprint import rlog as log
|
|
from .config.crop_config import CropConfig
|
|
from .config.inference_config import InferenceConfig
|
|
from .live_portrait_wrapper import LivePortraitWrapper
|
|
|
|
class TemplateMaker:
|
|
|
|
def __init__(self, inference_cfg: InferenceConfig, crop_cfg: CropConfig):
|
|
self.live_portrait_wrapper: LivePortraitWrapper = LivePortraitWrapper(cfg=inference_cfg)
|
|
self.cropper = Cropper(crop_cfg=crop_cfg)
|
|
|
|
def make_motion_template(self, video_fp: str, output_path: str, **kwargs):
|
|
""" make video template (.pkl format)
|
|
video_fp: driving video file path
|
|
output_path: where to save the pickle file
|
|
"""
|
|
|
|
driving_rgb_lst = load_driving_info(video_fp)
|
|
driving_rgb_lst = [cv2.resize(_, (256, 256)) for _ in driving_rgb_lst]
|
|
driving_lmk_lst = self.cropper.get_retargeting_lmk_info(driving_rgb_lst)
|
|
I_d_lst = self.live_portrait_wrapper.prepare_driving_videos(driving_rgb_lst)
|
|
|
|
n_frames = I_d_lst.shape[0]
|
|
|
|
templates = []
|
|
|
|
|
|
for i in track(range(n_frames), description='Making templates...', total=n_frames):
|
|
I_d_i = I_d_lst[i]
|
|
x_d_i_info = self.live_portrait_wrapper.get_kp_info(I_d_i)
|
|
R_d_i = get_rotation_matrix(x_d_i_info['pitch'], x_d_i_info['yaw'], x_d_i_info['roll'])
|
|
# collect s_d, R_d, δ_d and t_d for inference
|
|
template_dct = {
|
|
'n_frames': n_frames,
|
|
'frames_index': i,
|
|
}
|
|
template_dct['scale'] = x_d_i_info['scale'].cpu().numpy().astype(np.float32)
|
|
template_dct['R_d'] = R_d_i.cpu().numpy().astype(np.float32)
|
|
template_dct['exp'] = x_d_i_info['exp'].cpu().numpy().astype(np.float32)
|
|
template_dct['t'] = x_d_i_info['t'].cpu().numpy().astype(np.float32)
|
|
|
|
templates.append(template_dct)
|
|
|
|
mkdir(output_path)
|
|
# Save the dictionary as a pickle file
|
|
pickle_fp = os.path.join(output_path, f'{basename(video_fp)}.pkl')
|
|
with open(pickle_fp, 'wb') as f:
|
|
pickle.dump([templates, driving_lmk_lst], f)
|
|
log(f"Template saved at {pickle_fp}")
|