| import numpy as np
|
| from typing import Tuple
|
|
|
|
|
| def get_keypoints(joints, landmarks_name):
|
| frames_keypoints = joints[landmarks_name]
|
| frames_keypoints = frames_keypoints[:, :, :2]
|
|
|
| return frames_keypoints
|
|
|
|
|
| def output_keypoints(joints, valid_frames, frames_keypoints):
|
| frames_names = np.array(list(joints.keys()))[valid_frames]
|
| frames_keypoints = frames_keypoints[valid_frames]
|
|
|
| return dict(zip(frames_names, frames_keypoints))
|
|
|
|
|
| def safe_divide(a, b):
|
| return np.divide(a, b, out=np.zeros_like(a), where=b != 0)
|
|
|
|
|
| def local_keypoint_normalization(joints: dict, landmarks: str, select_idx: list = [], padding: float = 0.1) -> dict:
|
| frames_keypoints = get_keypoints(joints, landmarks)
|
|
|
| if select_idx:
|
| frames_keypoints = frames_keypoints[:, select_idx, :]
|
|
|
|
|
| xmin = np.min(frames_keypoints[:, :, 0], axis=1)
|
| ymin = np.min(frames_keypoints[:, :, 1], axis=1)
|
|
|
| frames_keypoints[:, :, 0] -= xmin[:, np.newaxis]
|
| frames_keypoints[:, :, 1] -= ymin[:, np.newaxis]
|
|
|
|
|
| xmax = np.max(frames_keypoints[:, :, 0], axis=1)
|
| ymax = np.max(frames_keypoints[:, :, 1], axis=1)
|
|
|
| dif_full = np.abs(xmax - ymax)
|
| dif = np.floor(dif_full / 2)
|
|
|
| for i in range(len(dif)):
|
| if xmax[i] > ymax[i]:
|
| ymax[i] += dif_full[i]
|
| frames_keypoints[i, :, 1] += dif[i]
|
| else:
|
| xmax[i] += dif_full[i]
|
| frames_keypoints[i, :, 0] += dif[i]
|
|
|
|
|
| side_size = np.max([xmax, ymax], axis=0)
|
| padding = side_size * padding
|
|
|
| frames_keypoints += padding[:, np.newaxis, np.newaxis]
|
| xmax += padding * 2
|
| ymax += padding * 2
|
|
|
|
|
| frames_keypoints = safe_divide(frames_keypoints, xmax[:, np.newaxis, np.newaxis])
|
|
|
| frames_keypoints = frames_keypoints * 2 - 1
|
|
|
| return frames_keypoints
|
|
|
|
|
| def global_keypoint_normalization(
|
| joints: dict,
|
| landmarks: str,
|
| add_landmarks_names: list,
|
| face_select_idx: list = [],
|
| sign_area_size: tuple = (1.5, 1.5),
|
| l_shoulder_idx: int = 11,
|
| r_shoulder_idx: int = 12) -> Tuple[dict, dict]:
|
| frames_keypoints = get_keypoints(joints, landmarks)
|
|
|
|
|
| l_shoulder_points = frames_keypoints[:, l_shoulder_idx, :]
|
| r_shoulder_points = frames_keypoints[:, r_shoulder_idx, :]
|
| distance = np.sqrt((l_shoulder_points[:, 0] - r_shoulder_points[:, 0]) ** 2 + (
|
| l_shoulder_points[:, 1] - r_shoulder_points[:, 1]) ** 2)
|
|
|
|
|
| center_x = np.abs(l_shoulder_points[:, 0] - r_shoulder_points[:, 0]) / 2 + np.min(
|
| [l_shoulder_points[:, 0], r_shoulder_points[:, 0]], 0)
|
| center_y = np.abs(l_shoulder_points[:, 1] - r_shoulder_points[:, 1]) / 2 + np.min(
|
| [l_shoulder_points[:, 1], r_shoulder_points[:, 1]], 0)
|
| sign_area_size = np.array(sign_area_size) * distance[:, np.newaxis]
|
|
|
|
|
| frames_keypoints[:, :, 0] -= center_x[:, np.newaxis]
|
| frames_keypoints[:, :, 1] -= center_y[:, np.newaxis]
|
|
|
|
|
|
|
| frames_keypoints[:, :, 0] = safe_divide(frames_keypoints[:, :, 0], sign_area_size[:, 0, np.newaxis])
|
| frames_keypoints[:, :, 1] = safe_divide(frames_keypoints[:, :, 1], sign_area_size[:, 1, np.newaxis])
|
|
|
|
|
| add_landmarks = {}
|
| for add_landmarks_name in add_landmarks_names:
|
| add_frames_keypoints = get_keypoints(joints, add_landmarks_name)
|
|
|
| if face_select_idx and add_landmarks_name == "face_landmarks":
|
| add_frames_keypoints = add_frames_keypoints[:, face_select_idx, :]
|
|
|
| add_frames_keypoints[:, :, 0] -= center_x[:, np.newaxis]
|
| add_frames_keypoints[:, :, 1] -= center_y[:, np.newaxis]
|
|
|
|
|
|
|
| add_frames_keypoints[:, :, 0] = safe_divide(add_frames_keypoints[:, :, 0], sign_area_size[:, 0, np.newaxis])
|
| add_frames_keypoints[:, :, 1] = safe_divide(add_frames_keypoints[:, :, 1], sign_area_size[:, 1, np.newaxis])
|
|
|
| add_landmarks[add_landmarks_name] = add_frames_keypoints
|
|
|
| return frames_keypoints, add_landmarks
|
|
|
