|
|
|
|
|
""" |
|
utility functions and classes to handle feature extraction and model loading |
|
""" |
|
|
|
import os |
|
import os.path as osp |
|
import cv2 |
|
import torch |
|
import yaml |
|
import argparse |
|
import locale |
|
import numpy as np |
|
from PIL import Image |
|
from rich.console import Console |
|
from collections import OrderedDict |
|
|
|
from ..live_portrait.spade_generator import SPADEDecoder |
|
from ..live_portrait.warping_network import WarpingNetwork |
|
from ..live_portrait.motion_extractor import MotionExtractor |
|
from ..live_portrait.appearance_feature_extractor import AppearanceFeatureExtractor |
|
from ..live_portrait.stitching_retargeting_network import StitchingRetargetingNetwork |
|
from .rprint import rlog as log |
|
|
|
|
|
def suffix(filename): |
|
"""a.jpg -> jpg""" |
|
pos = filename.rfind(".") |
|
if pos == -1: |
|
return "" |
|
return filename[pos + 1:] |
|
|
|
|
|
def prefix(filename): |
|
"""a.jpg -> a""" |
|
pos = filename.rfind(".") |
|
if pos == -1: |
|
return filename |
|
return filename[:pos] |
|
|
|
|
|
def basename(filename): |
|
"""a/b/c.jpg -> c""" |
|
return prefix(osp.basename(filename)) |
|
|
|
|
|
def is_video(file_path): |
|
if file_path.lower().endswith((".mp4", ".mov", ".avi", ".webm")) or osp.isdir(file_path): |
|
return True |
|
return False |
|
|
|
def is_template(file_path): |
|
if file_path.endswith(".pkl"): |
|
return True |
|
return False |
|
|
|
|
|
def mkdir(d, log=False): |
|
|
|
if not osp.exists(d): |
|
os.makedirs(d, exist_ok=True) |
|
if log: |
|
print(f"Make dir: {d}") |
|
return d |
|
|
|
|
|
def squeeze_tensor_to_numpy(tensor): |
|
out = tensor.data.squeeze(0).cpu().numpy() |
|
return out |
|
|
|
|
|
def dct2cuda(dct: dict, device_id: int): |
|
for key in dct: |
|
dct[key] = torch.tensor(dct[key]).cuda(device_id) |
|
return dct |
|
|
|
|
|
def concat_feat(kp_source: torch.Tensor, kp_driving: torch.Tensor) -> torch.Tensor: |
|
""" |
|
kp_source: (bs, k, 3) |
|
kp_driving: (bs, k, 3) |
|
Return: (bs, 2k*3) |
|
""" |
|
bs_src = kp_source.shape[0] |
|
bs_dri = kp_driving.shape[0] |
|
assert bs_src == bs_dri, 'batch size must be equal' |
|
|
|
feat = torch.cat([kp_source.view(bs_src, -1), kp_driving.view(bs_dri, -1)], dim=1) |
|
return feat |
|
|
|
|
|
|
|
def calculate_transformation(config, s_kp_info, t_0_kp_info, t_i_kp_info, R_s, R_t_0, R_t_i): |
|
if config.relative: |
|
new_rotation = (R_t_i @ R_t_0.permute(0, 2, 1)) @ R_s |
|
new_expression = s_kp_info['exp'] + (t_i_kp_info['exp'] - t_0_kp_info['exp']) |
|
else: |
|
new_rotation = R_t_i |
|
new_expression = t_i_kp_info['exp'] |
|
new_translation = s_kp_info['t'] + (t_i_kp_info['t'] - t_0_kp_info['t']) |
|
new_translation[..., 2].fill_(0) |
|
new_scale = s_kp_info['scale'] * (t_i_kp_info['scale'] / t_0_kp_info['scale']) |
|
return new_rotation, new_expression, new_translation, new_scale |
|
|
|
def load_description(fp): |
|
with open(fp, 'r', encoding='utf-8') as f: |
|
content = f.read() |
|
return content |
|
|
|
|
|
def resize_to_limit(img, max_dim=1280, n=2): |
|
h, w = img.shape[:2] |
|
if max_dim > 0 and max(h, w) > max_dim: |
|
if h > w: |
|
new_h = max_dim |
|
new_w = int(w * (max_dim / h)) |
|
else: |
|
new_w = max_dim |
|
new_h = int(h * (max_dim / w)) |
|
img = cv2.resize(img, (new_w, new_h)) |
|
n = max(n, 1) |
|
new_h = img.shape[0] - (img.shape[0] % n) |
|
new_w = img.shape[1] - (img.shape[1] % n) |
|
if new_h == 0 or new_w == 0: |
|
return img |
|
if new_h != img.shape[0] or new_w != img.shape[1]: |
|
img = img[:new_h, :new_w] |
|
return img |
|
|
|
|
|
def load_yaml(file_path): |
|
encoding = locale.getpreferredencoding(False) |
|
with open(file_path, 'r', encoding=encoding) as file: |
|
data = yaml.safe_load(file) |
|
return data |
|
|
|
|
|
def str2bool(v): |
|
if isinstance(v, bool): |
|
return v |
|
if v.lower() in ('yes', 'true', 't', 'y', '1'): |
|
return True |
|
elif v.lower() in ('no', 'false', 'f', 'n', '0'): |
|
return False |
|
else: |
|
raise argparse.ArgumentTypeError('Boolean value expected.') |
|
|
