pipelines/utils.py

import torch
import math
import pathlib

import cv2
import numpy as np
import os

from imageio_ffmpeg import get_ffmpeg_exe
from scipy.ndimage import median_filter


tensor_interpolation = None


def get_tensor_interpolation_method():
    return tensor_interpolation


def set_tensor_interpolation_method(is_slerp):
    global tensor_interpolation
    tensor_interpolation = slerp if is_slerp else linear


def linear(v1, v2, t):
    return (1.0 - t) * v1 + t * v2


def slerp(
        v0: torch.Tensor, v1: torch.Tensor, t: float, DOT_THRESHOLD: float = 0.9995
) -> torch.Tensor:
    u0 = v0 / v0.norm()
    u1 = v1 / v1.norm()
    dot = (u0 * u1).sum()
    if dot.abs() > DOT_THRESHOLD:
        # logger.info(f'warning: v0 and v1 close to parallel, using linear interpolation instead.')
        return (1.0 - t) * v0 + t * v1
    omega = dot.acos()
    return (((1.0 - t) * omega).sin() * v0 + (t * omega).sin() * v1) / omega.sin()


def draw_kps_image(image, kps, color_list=[(255, 0, 0), (0, 255, 0), (0, 0, 255)]):
    stick_width = 4
    limb_seq = np.array([[0, 2], [1, 2]])
    kps = np.array(kps)

    canvas = image

    for i in range(len(limb_seq)):
        index = limb_seq[i]
        color = color_list[index[0]]

        x = kps[index][:, 0]
        y = kps[index][:, 1]
        length = ((x[0] - x[1]) ** 2 + (y[0] - y[1]) ** 2) ** 0.5
        angle = int(math.degrees(math.atan2(y[0] - y[1], x[0] - x[1])))
        polygon = cv2.ellipse2Poly((int(np.mean(x)), int(np.mean(y))), (int(length / 2), stick_width), angle, 0, 360, 1)
        cv2.fillConvexPoly(canvas, polygon, [int(float(c) * 0.6) for c in color])

    for idx_kp, kp in enumerate(kps):
        color = color_list[idx_kp]
        x, y = kp
        cv2.circle(canvas, (int(x), int(y)), 4, color, -1)

    return canvas


def save_video(video_tensor, audio_path, output_path, fps=30.0):
    pathlib.Path(output_path).parent.mkdir(exist_ok=True, parents=True)

    video_tensor = video_tensor[0, ...]
    _, num_frames, height, width = video_tensor.shape

    video_tensor = video_tensor.permute(1, 2, 3, 0)
    video_np = (video_tensor * 255).numpy().astype(np.uint8)
    video_np_filtered = median_filter(video_np, size=(3, 3, 3, 1))

    output_name = pathlib.Path(output_path).stem
    temp_output_path = output_path.replace(output_name, output_name + '-temp')
    video_writer = cv2.VideoWriter(temp_output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))

    for i in range(num_frames):
        frame_image = video_np_filtered[i]
        frame_image = cv2.cvtColor(frame_image, cv2.COLOR_RGB2BGR)
        video_writer.write(frame_image)
    video_writer.release()

    cmd = (f'{get_ffmpeg_exe()} -i "{temp_output_path}" -i "{audio_path}" '
           f'-map 0:v -map 1:a -c:v h264 -shortest -y "{output_path}" -loglevel quiet')
    os.system(cmd)
    os.system(f'rm -rf "{temp_output_path}"')


def compute_dist(x1, y1, x2, y2):
    return math.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2)


def compute_ratio(kps):
    l_eye_x, l_eye_y = kps[0][0], kps[0][1]
    r_eye_x, r_eye_y = kps[1][0], kps[1][1]
    nose_x, nose_y = kps[2][0], kps[2][1]
    d_left = compute_dist(l_eye_x, l_eye_y, nose_x, nose_y)
    d_right = compute_dist(r_eye_x, r_eye_y, nose_x, nose_y)
    ratio = d_left / (d_right + 1e-6)
    return ratio


def point_to_line_dist(point, line_points):
    point = np.array(point)
    line_points = np.array(line_points)
    line_vec = line_points[1] - line_points[0]
    point_vec = point - line_points[0]
    line_norm = line_vec / np.sqrt(np.sum(line_vec ** 2))
    point_vec_scaled = point_vec * 1.0 / np.sqrt(np.sum(line_vec ** 2))
    t = np.dot(line_norm, point_vec_scaled)
    if t < 0.0:
        t = 0.0
    elif t > 1.0:
        t = 1.0
    nearest = line_points[0] + t * line_vec
    dist = np.sqrt(np.sum((point - nearest) ** 2))
    return dist


def get_face_size(kps):
    # 0: left eye, 1: right eye, 2: nose
    A = kps[0, :]
    B = kps[1, :]
    C = kps[2, :]

    AB_dist = math.sqrt((A[0] - B[0])**2 + (A[1] - B[1])**2)
    C_AB_dist = point_to_line_dist(C, [A, B])
    return AB_dist, C_AB_dist


def get_rescale_params(kps_ref, kps_target):
    kps_ref = np.array(kps_ref)
    kps_target = np.array(kps_target)

    ref_AB_dist, ref_C_AB_dist = get_face_size(kps_ref)
    target_AB_dist, target_C_AB_dist = get_face_size(kps_target)

    scale_width = ref_AB_dist / target_AB_dist
    scale_height = ref_C_AB_dist / target_C_AB_dist

    return scale_width, scale_height


def retarget_kps(ref_kps, tgt_kps_list, only_offset=True):
    ref_kps = np.array(ref_kps)
    tgt_kps_list = np.array(tgt_kps_list)

    ref_ratio = compute_ratio(ref_kps)

    ratio_delta = 10000
    selected_tgt_kps_idx = None
    for idx, tgt_kps in enumerate(tgt_kps_list):
        tgt_ratio = compute_ratio(tgt_kps)
        if math.fabs(tgt_ratio - ref_ratio) < ratio_delta:
            selected_tgt_kps_idx = idx
            ratio_delta = tgt_ratio

    scale_width, scale_height = get_rescale_params(
        kps_ref=ref_kps,
        kps_target=tgt_kps_list[selected_tgt_kps_idx],
    )

    rescaled_tgt_kps_list = np.array(tgt_kps_list)
    rescaled_tgt_kps_list[:, :, 0] *= scale_width
    rescaled_tgt_kps_list[:, :, 1] *= scale_height

    if only_offset:
        nose_offset = rescaled_tgt_kps_list[:, 2, :] - rescaled_tgt_kps_list[0, 2, :]
        nose_offset = nose_offset[:, np.newaxis, :]
        ref_kps_repeat = np.tile(ref_kps, (tgt_kps_list.shape[0], 1, 1))

        ref_kps_repeat[:, :, :] -= (nose_offset / 2.0)
        rescaled_tgt_kps_list = ref_kps_repeat
    else:
        nose_offset_x = rescaled_tgt_kps_list[0, 2, 0] - ref_kps[2][0]
        nose_offset_y = rescaled_tgt_kps_list[0, 2, 1] - ref_kps[2][1]

        rescaled_tgt_kps_list[:, :, 0] -= nose_offset_x
        rescaled_tgt_kps_list[:, :, 1] -= nose_offset_y

    return rescaled_tgt_kps_list