96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
# coding: utf-8
|
|
|
|
"""
|
|
utility functions and classes to handle feature extraction and model loading
|
|
"""
|
|
|
|
import cv2
|
|
import torch
|
|
import numpy as np
|
|
from typing import Union
|
|
from collections import OrderedDict
|
|
from scipy.spatial import ConvexHull # pylint: disable=E0401,E0611
|
|
|
|
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]).to(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 remove_ddp_dumplicate_key(state_dict):
|
|
state_dict_new = OrderedDict()
|
|
for key in state_dict.keys():
|
|
state_dict_new[key.replace('module.', '')] = state_dict[key]
|
|
return state_dict_new
|
|
|
|
# get coefficients of Eqn. 7
|
|
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) # Keep the z-axis unchanged
|
|
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 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 tensor_to_numpy(data: Union[np.ndarray, torch.Tensor]) -> np.ndarray:
|
|
"""transform torch.Tensor into numpy.ndarray"""
|
|
if isinstance(data, torch.Tensor):
|
|
return data.data.cpu().numpy()
|
|
return data
|
|
|
|
def calc_motion_multiplier(
|
|
kp_source: Union[np.ndarray, torch.Tensor],
|
|
kp_driving_initial: Union[np.ndarray, torch.Tensor]
|
|
) -> float:
|
|
"""calculate motion_multiplier based on the source image and the first driving frame"""
|
|
kp_source_np = tensor_to_numpy(kp_source)
|
|
kp_driving_initial_np = tensor_to_numpy(kp_driving_initial)
|
|
|
|
source_area = ConvexHull(kp_source_np.squeeze(0)).volume
|
|
driving_area = ConvexHull(kp_driving_initial_np.squeeze(0)).volume
|
|
motion_multiplier = np.sqrt(source_area) / np.sqrt(driving_area)
|
|
# motion_multiplier = np.cbrt(source_area) / np.cbrt(driving_area)
|
|
|
|
return motion_multiplier |