# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. import copy import math from typing import NamedTuple import numpy as np from pose2d_utils import AAPoseMeta from tqdm import tqdm # load skeleton name and bone lines keypoint_list = [ "Nose", "Neck", "RShoulder", "RElbow", "RWrist", # No.4 "LShoulder", "LElbow", "LWrist", # No.7 "RHip", "RKnee", "RAnkle", # No.10 "LHip", "LKnee", "LAnkle", # No.13 "REye", "LEye", "REar", "LEar", "LToe", "RToe", ] limbSeq = [ [2, 3], [2, 6], # shoulders [3, 4], [4, 5], # left arm [6, 7], [7, 8], # right arm [2, 9], [9, 10], [10, 11], # right leg [2, 12], [12, 13], [13, 14], # left leg [2, 1], [1, 15], [15, 17], [1, 16], [16, 18], # face (nose, eyes, ears) [14, 19], # left foot [11, 20], # right foot ] eps = 0.01 class Keypoint(NamedTuple): x: float y: float score: float = 1.0 id: int = -1 # for each limb, calculate src & dst bone's length # and calculate their ratios def get_length(skeleton, limb): k1_index, k2_index = limb H, W = skeleton["height"], skeleton["width"] keypoints = skeleton["keypoints_body"] keypoint1 = keypoints[k1_index - 1] keypoint2 = keypoints[k2_index - 1] if keypoint1 is None or keypoint2 is None: return None, None, None X = np.array([keypoint1[0], keypoint2[0]]) * float(W) Y = np.array([keypoint1[1], keypoint2[1]]) * float(H) length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5 return X, Y, length def get_handpose_meta(keypoints, delta, src_H, src_W): new_keypoints = [] for idx, keypoint in enumerate(keypoints): if keypoint is None: new_keypoints.append(None) continue if keypoint.score == 0: new_keypoints.append(None) continue x, y = keypoint.x, keypoint.y x = int(x * src_W + delta[0]) y = int(y * src_H + delta[1]) new_keypoints.append( Keypoint( x=x, y=y, score=keypoint.score, ) ) return new_keypoints def deal_hand_keypoints(hand_res, r_ratio, l_ratio, hand_score_th=0.5): left_hand = [] right_hand = [] left_delta_x = hand_res["left"][0][0] * (l_ratio - 1) left_delta_y = hand_res["left"][0][1] * (l_ratio - 1) right_delta_x = hand_res["right"][0][0] * (r_ratio - 1) right_delta_y = hand_res["right"][0][1] * (r_ratio - 1) length = len(hand_res["left"]) for i in range(length): # left hand if hand_res["left"][i][2] < hand_score_th: left_hand.append( Keypoint( x=-1, y=-1, score=0, ) ) else: left_hand.append(Keypoint(x=hand_res["left"][i][0] * l_ratio - left_delta_x, y=hand_res["left"][i][1] * l_ratio - left_delta_y, score=hand_res["left"][i][2])) # right hand if hand_res["right"][i][2] < hand_score_th: right_hand.append( Keypoint( x=-1, y=-1, score=0, ) ) else: right_hand.append(Keypoint(x=hand_res["right"][i][0] * r_ratio - right_delta_x, y=hand_res["right"][i][1] * r_ratio - right_delta_y, score=hand_res["right"][i][2])) return right_hand, left_hand def get_scaled_pose(canvas, src_canvas, keypoints, keypoints_hand, bone_ratio_list, delta_ground_x, delta_ground_y, rescaled_src_ground_x, body_flag, id, scale_min, threshold=0.4): H, W = canvas src_H, src_W = src_canvas new_length_list = [] angle_list = [] # keypoints from 0-1 to H/W range for idx in range(len(keypoints)): if keypoints[idx] is None or len(keypoints[idx]) == 0: continue keypoints[idx] = [keypoints[idx][0] * src_W, keypoints[idx][1] * src_H, keypoints[idx][2]] # first traverse, get new_length_list and angle_list for idx, (k1_index, k2_index) in enumerate(limbSeq): keypoint1 = keypoints[k1_index - 1] keypoint2 = keypoints[k2_index - 1] if keypoint1 is None or keypoint2 is None or len(keypoint1) == 0 or len(keypoint2) == 0: new_length_list.append(None) angle_list.append(None) continue Y = np.array([keypoint1[0], keypoint2[0]]) # * float(W) X = np.array([keypoint1[1], keypoint2[1]]) # * float(H) length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5 new_length = length * bone_ratio_list[idx] angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1])) new_length_list.append(new_length) angle_list.append(angle) # Keep foot length within 0.5x calf length foot_lower_leg_ratio = 0.5 if new_length_list[8] != None and new_length_list[18] != None: # noqa if new_length_list[18] > new_length_list[8] * foot_lower_leg_ratio: new_length_list[18] = new_length_list[8] * foot_lower_leg_ratio if new_length_list[11] != None and new_length_list[17] != None: # noqa if new_length_list[17] > new_length_list[11] * foot_lower_leg_ratio: new_length_list[17] = new_length_list[11] * foot_lower_leg_ratio # second traverse, calculate new keypoints rescale_keypoints = keypoints.copy() for idx, (k1_index, k2_index) in enumerate(limbSeq): # update dst_keypoints start_keypoint = rescale_keypoints[k1_index - 1] new_length = new_length_list[idx] angle = angle_list[idx] if rescale_keypoints[k1_index - 1] is None or rescale_keypoints[k2_index - 1] is None or len(rescale_keypoints[k1_index - 1]) == 0 or len(rescale_keypoints[k2_index - 1]) == 0: continue # calculate end_keypoint delta_x = new_length * math.cos(math.radians(angle)) delta_y = new_length * math.sin(math.radians(angle)) end_keypoint_x = start_keypoint[0] - delta_x end_keypoint_y = start_keypoint[1] - delta_y # update keypoints rescale_keypoints[k2_index - 1] = [end_keypoint_x, end_keypoint_y, rescale_keypoints[k2_index - 1][2]] if id == 0: if body_flag == "full_body" and rescale_keypoints[8] != None and rescale_keypoints[11] != None: # noqa delta_ground_x_offset_first_frame = (rescale_keypoints[8][0] + rescale_keypoints[11][0]) / 2 - rescaled_src_ground_x delta_ground_x += delta_ground_x_offset_first_frame elif body_flag == "half_body" and rescale_keypoints[1] != None: # noqa delta_ground_x_offset_first_frame = rescale_keypoints[1][0] - rescaled_src_ground_x delta_ground_x += delta_ground_x_offset_first_frame # offset all keypoints for idx in range(len(rescale_keypoints)): if rescale_keypoints[idx] is None or len(rescale_keypoints[idx]) == 0: continue rescale_keypoints[idx][0] -= delta_ground_x rescale_keypoints[idx][1] -= delta_ground_y # rescale keypoints to original size rescale_keypoints[idx][0] /= scale_min rescale_keypoints[idx][1] /= scale_min # Scale hand proportions based on body skeletal ratios r_ratio = max(bone_ratio_list[0], bone_ratio_list[1]) / scale_min l_ratio = max(bone_ratio_list[0], bone_ratio_list[1]) / scale_min left_hand, right_hand = deal_hand_keypoints(keypoints_hand, r_ratio, l_ratio, hand_score_th=threshold) left_hand_new = left_hand.copy() right_hand_new = right_hand.copy() if rescale_keypoints[4] == None and rescale_keypoints[7] == None: # noqa pass elif rescale_keypoints[4] == None and rescale_keypoints[7] != None: # noqa right_hand_delta = np.array(rescale_keypoints[7][:2]) - np.array(keypoints[7][:2]) right_hand_new = get_handpose_meta(right_hand, right_hand_delta, src_H, src_W) elif rescale_keypoints[4] != None and rescale_keypoints[7] == None: # noqa left_hand_delta = np.array(rescale_keypoints[4][:2]) - np.array(keypoints[4][:2]) left_hand_new = get_handpose_meta(left_hand, left_hand_delta, src_H, src_W) else: # get left_hand and right_hand offset left_hand_delta = np.array(rescale_keypoints[4][:2]) - np.array(keypoints[4][:2]) right_hand_delta = np.array(rescale_keypoints[7][:2]) - np.array(keypoints[7][:2]) if keypoints[4][0] != None and left_hand[0].x != -1: # noqa left_hand_root_offset = np.array((keypoints[4][0] - left_hand[0].x * src_W, keypoints[4][1] - left_hand[0].y * src_H)) left_hand_delta += left_hand_root_offset if keypoints[7][0] != None and right_hand[0].x != -1: # noqa right_hand_root_offset = np.array((keypoints[7][0] - right_hand[0].x * src_W, keypoints[7][1] - right_hand[0].y * src_H)) right_hand_delta += right_hand_root_offset dis_left_hand = ((keypoints[4][0] - left_hand[0].x * src_W) ** 2 + (keypoints[4][1] - left_hand[0].y * src_H) ** 2) ** 0.5 dis_right_hand = ((keypoints[7][0] - left_hand[0].x * src_W) ** 2 + (keypoints[7][1] - left_hand[0].y * src_H) ** 2) ** 0.5 if dis_left_hand > dis_right_hand: right_hand_new = get_handpose_meta(left_hand, right_hand_delta, src_H, src_W) left_hand_new = get_handpose_meta(right_hand, left_hand_delta, src_H, src_W) else: left_hand_new = get_handpose_meta(left_hand, left_hand_delta, src_H, src_W) right_hand_new = get_handpose_meta(right_hand, right_hand_delta, src_H, src_W) # get normalized keypoints_body norm_body_keypoints = [] for body_keypoint in rescale_keypoints: if body_keypoint != None: # noqa norm_body_keypoints.append([body_keypoint[0] / W, body_keypoint[1] / H, body_keypoint[2]]) else: norm_body_keypoints.append(None) frame_info = { "height": H, "width": W, "keypoints_body": norm_body_keypoints, "keypoints_left_hand": left_hand_new, "keypoints_right_hand": right_hand_new, } return frame_info def rescale_skeleton(H, W, keypoints, bone_ratio_list): rescale_keypoints = keypoints.copy() new_length_list = [] angle_list = [] # keypoints from 0-1 to H/W range for idx in range(len(rescale_keypoints)): if rescale_keypoints[idx] is None or len(rescale_keypoints[idx]) == 0: continue rescale_keypoints[idx] = [rescale_keypoints[idx][0] * W, rescale_keypoints[idx][1] * H] # first traverse, get new_length_list and angle_list for idx, (k1_index, k2_index) in enumerate(limbSeq): keypoint1 = rescale_keypoints[k1_index - 1] keypoint2 = rescale_keypoints[k2_index - 1] if keypoint1 is None or keypoint2 is None or len(keypoint1) == 0 or len(keypoint2) == 0: new_length_list.append(None) angle_list.append(None) continue Y = np.array([keypoint1[0], keypoint2[0]]) # * float(W) X = np.array([keypoint1[1], keypoint2[1]]) # * float(H) length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5 new_length = length * bone_ratio_list[idx] angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1])) new_length_list.append(new_length) angle_list.append(angle) # # second traverse, calculate new keypoints for idx, (k1_index, k2_index) in enumerate(limbSeq): # update dst_keypoints start_keypoint = rescale_keypoints[k1_index - 1] new_length = new_length_list[idx] angle = angle_list[idx] if rescale_keypoints[k1_index - 1] is None or rescale_keypoints[k2_index - 1] is None or len(rescale_keypoints[k1_index - 1]) == 0 or len(rescale_keypoints[k2_index - 1]) == 0: continue # calculate end_keypoint delta_x = new_length * math.cos(math.radians(angle)) delta_y = new_length * math.sin(math.radians(angle)) end_keypoint_x = start_keypoint[0] - delta_x end_keypoint_y = start_keypoint[1] - delta_y # update keypoints rescale_keypoints[k2_index - 1] = [end_keypoint_x, end_keypoint_y] return rescale_keypoints def fix_lack_keypoints_use_sym(skeleton): keypoints = skeleton["keypoints_body"] H, W = skeleton["height"], skeleton["width"] limb_points_list = [ [3, 4, 5], [6, 7, 8], [12, 13, 14, 19], [9, 10, 11, 20], ] for limb_points in limb_points_list: miss_flag = False for point in limb_points: if keypoints[point - 1] is None: miss_flag = True continue if miss_flag: skeleton["keypoints_body"][point - 1] = None repair_limb_seq_left = [ [3, 4], [4, 5], # left arm [12, 13], [13, 14], # left leg [14, 19], # left foot ] repair_limb_seq_right = [ [6, 7], [7, 8], # right arm [9, 10], [10, 11], # right leg [11, 20], # right foot ] repair_limb_seq = [repair_limb_seq_left, repair_limb_seq_right] for idx_part, part in enumerate(repair_limb_seq): for idx, limb in enumerate(part): k1_index, k2_index = limb keypoint1 = keypoints[k1_index - 1] keypoint2 = keypoints[k2_index - 1] if keypoint1 != None and keypoint2 is None: # noqa # reference to symmetric limb sym_limb = repair_limb_seq[1 - idx_part][idx] k1_index_sym, k2_index_sym = sym_limb keypoint1_sym = keypoints[k1_index_sym - 1] keypoint2_sym = keypoints[k2_index_sym - 1] ref_length = 0 if keypoint1_sym != None and keypoint2_sym != None: # noqa X = np.array([keypoint1_sym[0], keypoint2_sym[0]]) * float(W) Y = np.array([keypoint1_sym[1], keypoint2_sym[1]]) * float(H) ref_length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5 else: ref_length_left, ref_length_right = 0, 0 if keypoints[1] != None and keypoints[8] != None: # noqa X = np.array([keypoints[1][0], keypoints[8][0]]) * float(W) Y = np.array([keypoints[1][1], keypoints[8][1]]) * float(H) ref_length_left = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5 if idx <= 1: # arms ref_length_left /= 2 if keypoints[1] != None and keypoints[11] != None: # noqa X = np.array([keypoints[1][0], keypoints[11][0]]) * float(W) Y = np.array([keypoints[1][1], keypoints[11][1]]) * float(H) ref_length_right = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5 if idx <= 1: # arms ref_length_right /= 2 elif idx == 4: # foot ref_length_right /= 5 ref_length = max(ref_length_left, ref_length_right) if ref_length != 0: skeleton["keypoints_body"][k2_index - 1] = [0, 0] # init skeleton["keypoints_body"][k2_index - 1][0] = skeleton["keypoints_body"][k1_index - 1][0] skeleton["keypoints_body"][k2_index - 1][1] = skeleton["keypoints_body"][k1_index - 1][1] + ref_length / H return skeleton def rescale_shorten_skeleton(ratio_list, src_length_list, dst_length_list): modify_bone_list = [[0, 1], [2, 4], [3, 5], [6, 9], [7, 10], [8, 11], [17, 18]] for modify_bone in modify_bone_list: new_ratio = max(ratio_list[modify_bone[0]], ratio_list[modify_bone[1]]) ratio_list[modify_bone[0]] = new_ratio ratio_list[modify_bone[1]] = new_ratio if ratio_list[13] != None and ratio_list[15] != None: # noqa ratio_eye_avg = (ratio_list[13] + ratio_list[15]) / 2 ratio_list[13] = ratio_eye_avg ratio_list[15] = ratio_eye_avg if ratio_list[14] != None and ratio_list[16] != None: # noqa ratio_eye_avg = (ratio_list[14] + ratio_list[16]) / 2 ratio_list[14] = ratio_eye_avg ratio_list[16] = ratio_eye_avg return ratio_list, src_length_list, dst_length_list def check_full_body(keypoints, threshold=0.4): body_flag = "half_body" # 1. If ankle points exist, confidence is greater than the threshold, and points do not exceed the frame, return full_body if keypoints[10] != None and keypoints[13] != None and keypoints[8] != None and keypoints[11] != None: # noqa if ( (keypoints[10][1] <= 1 and keypoints[13][1] <= 1) and (keypoints[10][2] >= threshold and keypoints[13][2] >= threshold) and (keypoints[8][1] <= 1 and keypoints[11][1] <= 1) and (keypoints[8][2] >= threshold and keypoints[11][2] >= threshold) ): body_flag = "full_body" return body_flag # 2. If hip points exist, return three_quarter_body if keypoints[8] != None and keypoints[11] != None: # noqa if (keypoints[8][1] <= 1 and keypoints[11][1] <= 1) and (keypoints[8][2] >= threshold and keypoints[11][2] >= threshold): body_flag = "three_quarter_body" return body_flag return body_flag def check_full_body_both(flag1, flag2): body_flag_dict = {"full_body": 2, "three_quarter_body": 1, "half_body": 0} body_flag_dict_reverse = {2: "full_body", 1: "three_quarter_body", 0: "half_body"} flag1_num = body_flag_dict[flag1] flag2_num = body_flag_dict[flag2] flag_both_num = min(flag1_num, flag2_num) return body_flag_dict_reverse[flag_both_num] def write_to_poses(data_to_json, none_idx, dst_shape, bone_ratio_list, delta_ground_x, delta_ground_y, rescaled_src_ground_x, body_flag, scale_min): outputs = [] length = len(data_to_json) for id in tqdm(range(length)): src_height, src_width = data_to_json[id]["height"], data_to_json[id]["width"] width, height = dst_shape keypoints = data_to_json[id]["keypoints_body"] for idx in range(len(keypoints)): if idx in none_idx: keypoints[idx] = None new_keypoints = keypoints.copy() # get hand keypoints keypoints_hand = {"left": data_to_json[id]["keypoints_left_hand"], "right": data_to_json[id]["keypoints_right_hand"]} # Normalize hand coordinates to 0-1 range for hand_idx in range(len(data_to_json[id]["keypoints_left_hand"])): data_to_json[id]["keypoints_left_hand"][hand_idx][0] = data_to_json[id]["keypoints_left_hand"][hand_idx][0] / src_width data_to_json[id]["keypoints_left_hand"][hand_idx][1] = data_to_json[id]["keypoints_left_hand"][hand_idx][1] / src_height for hand_idx in range(len(data_to_json[id]["keypoints_right_hand"])): data_to_json[id]["keypoints_right_hand"][hand_idx][0] = data_to_json[id]["keypoints_right_hand"][hand_idx][0] / src_width data_to_json[id]["keypoints_right_hand"][hand_idx][1] = data_to_json[id]["keypoints_right_hand"][hand_idx][1] / src_height frame_info = get_scaled_pose( (height, width), (src_height, src_width), new_keypoints, keypoints_hand, bone_ratio_list, delta_ground_x, delta_ground_y, rescaled_src_ground_x, body_flag, id, scale_min ) outputs.append(frame_info) return outputs def calculate_scale_ratio(skeleton, skeleton_edit, scale_ratio_flag): if scale_ratio_flag: headw = max(skeleton["keypoints_body"][0][0], skeleton["keypoints_body"][14][0], skeleton["keypoints_body"][15][0], skeleton["keypoints_body"][16][0], skeleton["keypoints_body"][17][0]) - min( skeleton["keypoints_body"][0][0], skeleton["keypoints_body"][14][0], skeleton["keypoints_body"][15][0], skeleton["keypoints_body"][16][0], skeleton["keypoints_body"][17][0] ) headw_edit = max( skeleton_edit["keypoints_body"][0][0], skeleton_edit["keypoints_body"][14][0], skeleton_edit["keypoints_body"][15][0], skeleton_edit["keypoints_body"][16][0], skeleton_edit["keypoints_body"][17][0], ) - min( skeleton_edit["keypoints_body"][0][0], skeleton_edit["keypoints_body"][14][0], skeleton_edit["keypoints_body"][15][0], skeleton_edit["keypoints_body"][16][0], skeleton_edit["keypoints_body"][17][0], ) headw_ratio = headw / headw_edit _, _, shoulder = get_length(skeleton, [6, 3]) _, _, shoulder_edit = get_length(skeleton_edit, [6, 3]) shoulder_ratio = shoulder / shoulder_edit return max(headw_ratio, shoulder_ratio) else: return 1 def retarget_pose(src_skeleton, dst_skeleton, all_src_skeleton, src_skeleton_edit, dst_skeleton_edit, threshold=0.4): if src_skeleton_edit is not None and dst_skeleton_edit is not None: # noqa use_edit_for_base = True else: use_edit_for_base = False src_skeleton_ori = copy.deepcopy(src_skeleton) dst_skeleton_ori_h, dst_skeleton_ori_w = dst_skeleton["height"], dst_skeleton["width"] if ( src_skeleton["keypoints_body"][0] != None # noqa and src_skeleton["keypoints_body"][10] != None # noqa and src_skeleton["keypoints_body"][13] != None # noqa and dst_skeleton["keypoints_body"][0] != None # noqa and dst_skeleton["keypoints_body"][10] != None # noqa and dst_skeleton["keypoints_body"][13] != None # noqa and src_skeleton["keypoints_body"][0][2] > 0.5 and src_skeleton["keypoints_body"][10][2] > 0.5 and src_skeleton["keypoints_body"][13][2] > 0.5 and dst_skeleton["keypoints_body"][0][2] > 0.5 and dst_skeleton["keypoints_body"][10][2] > 0.5 and dst_skeleton["keypoints_body"][13][2] > 0.5 ): src_height = src_skeleton["height"] * abs((src_skeleton["keypoints_body"][10][1] + src_skeleton["keypoints_body"][13][1]) / 2 - src_skeleton["keypoints_body"][0][1]) dst_height = dst_skeleton["height"] * abs((dst_skeleton["keypoints_body"][10][1] + dst_skeleton["keypoints_body"][13][1]) / 2 - dst_skeleton["keypoints_body"][0][1]) scale_min = 1.0 * src_height / dst_height elif ( src_skeleton["keypoints_body"][0] != None # noqa and src_skeleton["keypoints_body"][8] != None # noqa and src_skeleton["keypoints_body"][11] != None # noqa and dst_skeleton["keypoints_body"][0] != None # noqa and dst_skeleton["keypoints_body"][8] != None # noqa and dst_skeleton["keypoints_body"][11] != None # noqa and src_skeleton["keypoints_body"][0][2] > 0.5 and src_skeleton["keypoints_body"][8][2] > 0.5 and src_skeleton["keypoints_body"][11][2] > 0.5 and dst_skeleton["keypoints_body"][0][2] > 0.5 and dst_skeleton["keypoints_body"][8][2] > 0.5 and dst_skeleton["keypoints_body"][11][2] > 0.5 ): src_height = src_skeleton["height"] * abs((src_skeleton["keypoints_body"][8][1] + src_skeleton["keypoints_body"][11][1]) / 2 - src_skeleton["keypoints_body"][0][1]) dst_height = dst_skeleton["height"] * abs((dst_skeleton["keypoints_body"][8][1] + dst_skeleton["keypoints_body"][11][1]) / 2 - dst_skeleton["keypoints_body"][0][1]) scale_min = 1.0 * src_height / dst_height else: scale_min = np.sqrt(src_skeleton["height"] * src_skeleton["width"]) / np.sqrt(dst_skeleton["height"] * dst_skeleton["width"]) if use_edit_for_base: scale_ratio_flag = False if ( src_skeleton_edit["keypoints_body"][0] != None # noqa and src_skeleton_edit["keypoints_body"][10] != None # noqa and src_skeleton_edit["keypoints_body"][13] != None # noqa and dst_skeleton_edit["keypoints_body"][0] != None # noqa and dst_skeleton_edit["keypoints_body"][10] != None # noqa and dst_skeleton_edit["keypoints_body"][13] != None # noqa and src_skeleton_edit["keypoints_body"][0][2] > 0.5 and src_skeleton_edit["keypoints_body"][10][2] > 0.5 and src_skeleton_edit["keypoints_body"][13][2] > 0.5 and dst_skeleton_edit["keypoints_body"][0][2] > 0.5 and dst_skeleton_edit["keypoints_body"][10][2] > 0.5 and dst_skeleton_edit["keypoints_body"][13][2] > 0.5 ): src_height_edit = src_skeleton_edit["height"] * abs( (src_skeleton_edit["keypoints_body"][10][1] + src_skeleton_edit["keypoints_body"][13][1]) / 2 - src_skeleton_edit["keypoints_body"][0][1] ) dst_height_edit = dst_skeleton_edit["height"] * abs( (dst_skeleton_edit["keypoints_body"][10][1] + dst_skeleton_edit["keypoints_body"][13][1]) / 2 - dst_skeleton_edit["keypoints_body"][0][1] ) scale_min_edit = 1.0 * src_height_edit / dst_height_edit elif ( src_skeleton_edit["keypoints_body"][0] != None # noqa and src_skeleton_edit["keypoints_body"][8] != None # noqa and src_skeleton_edit["keypoints_body"][11] != None # noqa and dst_skeleton_edit["keypoints_body"][0] != None # noqa and dst_skeleton_edit["keypoints_body"][8] != None # noqa and dst_skeleton_edit["keypoints_body"][11] != None # noqa and src_skeleton_edit["keypoints_body"][0][2] > 0.5 and src_skeleton_edit["keypoints_body"][8][2] > 0.5 and src_skeleton_edit["keypoints_body"][11][2] > 0.5 and dst_skeleton_edit["keypoints_body"][0][2] > 0.5 and dst_skeleton_edit["keypoints_body"][8][2] > 0.5 and dst_skeleton_edit["keypoints_body"][11][2] > 0.5 ): src_height_edit = src_skeleton_edit["height"] * abs( (src_skeleton_edit["keypoints_body"][8][1] + src_skeleton_edit["keypoints_body"][11][1]) / 2 - src_skeleton_edit["keypoints_body"][0][1] ) dst_height_edit = dst_skeleton_edit["height"] * abs( (dst_skeleton_edit["keypoints_body"][8][1] + dst_skeleton_edit["keypoints_body"][11][1]) / 2 - dst_skeleton_edit["keypoints_body"][0][1] ) scale_min_edit = 1.0 * src_height_edit / dst_height_edit else: scale_min_edit = np.sqrt(src_skeleton_edit["height"] * src_skeleton_edit["width"]) / np.sqrt(dst_skeleton_edit["height"] * dst_skeleton_edit["width"]) scale_ratio_flag = True # Flux may change the scale, compensate for it here ratio_src = calculate_scale_ratio(src_skeleton, src_skeleton_edit, scale_ratio_flag) ratio_dst = calculate_scale_ratio(dst_skeleton, dst_skeleton_edit, scale_ratio_flag) dst_skeleton_edit["height"] = int(dst_skeleton_edit["height"] * scale_min_edit) dst_skeleton_edit["width"] = int(dst_skeleton_edit["width"] * scale_min_edit) for idx in range(len(dst_skeleton_edit["keypoints_left_hand"])): dst_skeleton_edit["keypoints_left_hand"][idx][0] *= scale_min_edit dst_skeleton_edit["keypoints_left_hand"][idx][1] *= scale_min_edit for idx in range(len(dst_skeleton_edit["keypoints_right_hand"])): dst_skeleton_edit["keypoints_right_hand"][idx][0] *= scale_min_edit dst_skeleton_edit["keypoints_right_hand"][idx][1] *= scale_min_edit dst_skeleton["height"] = int(dst_skeleton["height"] * scale_min) dst_skeleton["width"] = int(dst_skeleton["width"] * scale_min) for idx in range(len(dst_skeleton["keypoints_left_hand"])): dst_skeleton["keypoints_left_hand"][idx][0] *= scale_min dst_skeleton["keypoints_left_hand"][idx][1] *= scale_min for idx in range(len(dst_skeleton["keypoints_right_hand"])): dst_skeleton["keypoints_right_hand"][idx][0] *= scale_min dst_skeleton["keypoints_right_hand"][idx][1] *= scale_min dst_body_flag = check_full_body(dst_skeleton["keypoints_body"], threshold) src_body_flag = check_full_body(src_skeleton_ori["keypoints_body"], threshold) body_flag = check_full_body_both(dst_body_flag, src_body_flag) # print('body_flag: ', body_flag) if use_edit_for_base: src_skeleton_edit = fix_lack_keypoints_use_sym(src_skeleton_edit) dst_skeleton_edit = fix_lack_keypoints_use_sym(dst_skeleton_edit) else: src_skeleton = fix_lack_keypoints_use_sym(src_skeleton) dst_skeleton = fix_lack_keypoints_use_sym(dst_skeleton) none_idx = [] for idx in range(len(dst_skeleton["keypoints_body"])): if dst_skeleton["keypoints_body"][idx] == None or src_skeleton["keypoints_body"][idx] == None: # noqa src_skeleton["keypoints_body"][idx] = None dst_skeleton["keypoints_body"][idx] = None none_idx.append(idx) # get bone ratio list ratio_list, src_length_list, dst_length_list = [], [], [] for idx, limb in enumerate(limbSeq): if use_edit_for_base: src_X, src_Y, src_length = get_length(src_skeleton_edit, limb) dst_X, dst_Y, dst_length = get_length(dst_skeleton_edit, limb) if src_X is None or src_Y is None or dst_X is None or dst_Y is None: # noqa ratio = -1 else: ratio = 1.0 * dst_length * ratio_dst / src_length / ratio_src else: src_X, src_Y, src_length = get_length(src_skeleton, limb) dst_X, dst_Y, dst_length = get_length(dst_skeleton, limb) if src_X is None or src_Y is None or dst_X is None or dst_Y is None: # noqa ratio = -1 else: ratio = 1.0 * dst_length / src_length ratio_list.append(ratio) src_length_list.append(src_length) dst_length_list.append(dst_length) for idx, ratio in enumerate(ratio_list): if ratio == -1: if ratio_list[0] != -1 and ratio_list[1] != -1: ratio_list[idx] = (ratio_list[0] + ratio_list[1]) / 2 # Consider adding constraints when Flux fails to correct head pose, causing neck issues. # if ratio_list[12] > (ratio_list[0]+ratio_list[1])/2*1.25: # ratio_list[12] = (ratio_list[0]+ratio_list[1])/2*1.25 ratio_list, src_length_list, dst_length_list = rescale_shorten_skeleton(ratio_list, src_length_list, dst_length_list) rescaled_src_skeleton_ori = rescale_skeleton(src_skeleton_ori["height"], src_skeleton_ori["width"], src_skeleton_ori["keypoints_body"], ratio_list) # get global translation offset_x and offset_y if body_flag == "full_body": # print('use foot mark.') dst_ground_y = max(dst_skeleton["keypoints_body"][10][1], dst_skeleton["keypoints_body"][13][1]) * dst_skeleton["height"] # The midpoint between toe and ankle if dst_skeleton["keypoints_body"][18] != None and dst_skeleton["keypoints_body"][19] != None: # noqa right_foot_mid = (dst_skeleton["keypoints_body"][10][1] + dst_skeleton["keypoints_body"][19][1]) / 2 left_foot_mid = (dst_skeleton["keypoints_body"][13][1] + dst_skeleton["keypoints_body"][18][1]) / 2 dst_ground_y = max(left_foot_mid, right_foot_mid) * dst_skeleton["height"] rescaled_src_ground_y = max(rescaled_src_skeleton_ori[10][1], rescaled_src_skeleton_ori[13][1]) delta_ground_y = rescaled_src_ground_y - dst_ground_y dst_ground_x = (dst_skeleton["keypoints_body"][8][0] + dst_skeleton["keypoints_body"][11][0]) * dst_skeleton["width"] / 2 rescaled_src_ground_x = (rescaled_src_skeleton_ori[8][0] + rescaled_src_skeleton_ori[11][0]) / 2 delta_ground_x = rescaled_src_ground_x - dst_ground_x delta_x, delta_y = delta_ground_x, delta_ground_y else: # print('use neck mark.') # use neck keypoint as mark src_neck_y = rescaled_src_skeleton_ori[1][1] dst_neck_y = dst_skeleton["keypoints_body"][1][1] delta_neck_y = src_neck_y - dst_neck_y * dst_skeleton["height"] src_neck_x = rescaled_src_skeleton_ori[1][0] dst_neck_x = dst_skeleton["keypoints_body"][1][0] delta_neck_x = src_neck_x - dst_neck_x * dst_skeleton["width"] delta_x, delta_y = delta_neck_x, delta_neck_y rescaled_src_ground_x = src_neck_x dst_shape = (dst_skeleton_ori_w, dst_skeleton_ori_h) output = write_to_poses(all_src_skeleton, none_idx, dst_shape, ratio_list, delta_x, delta_y, rescaled_src_ground_x, body_flag, scale_min) return output def get_retarget_pose(tpl_pose_meta0, refer_pose_meta, tpl_pose_metas, tql_edit_pose_meta0, refer_edit_pose_meta): for key, value in tpl_pose_meta0.items(): if type(value) is np.ndarray: if key in ["keypoints_left_hand", "keypoints_right_hand"]: value = value * np.array([[tpl_pose_meta0["width"], tpl_pose_meta0["height"], 1.0]]) if not isinstance(value, list): value = value.tolist() tpl_pose_meta0[key] = value for key, value in refer_pose_meta.items(): if type(value) is np.ndarray: if key in ["keypoints_left_hand", "keypoints_right_hand"]: value = value * np.array([[refer_pose_meta["width"], refer_pose_meta["height"], 1.0]]) if not isinstance(value, list): value = value.tolist() refer_pose_meta[key] = value tpl_pose_metas_new = [] for meta in tpl_pose_metas: for key, value in meta.items(): if type(value) is np.ndarray: if key in ["keypoints_left_hand", "keypoints_right_hand"]: value = value * np.array([[meta["width"], meta["height"], 1.0]]) if not isinstance(value, list): value = value.tolist() meta[key] = value tpl_pose_metas_new.append(meta) if tql_edit_pose_meta0 is not None: for key, value in tql_edit_pose_meta0.items(): if type(value) is np.ndarray: if key in ["keypoints_left_hand", "keypoints_right_hand"]: value = value * np.array([[tql_edit_pose_meta0["width"], tql_edit_pose_meta0["height"], 1.0]]) if not isinstance(value, list): value = value.tolist() tql_edit_pose_meta0[key] = value if refer_edit_pose_meta is not None: for key, value in refer_edit_pose_meta.items(): if type(value) is np.ndarray: if key in ["keypoints_left_hand", "keypoints_right_hand"]: value = value * np.array([[refer_edit_pose_meta["width"], refer_edit_pose_meta["height"], 1.0]]) if not isinstance(value, list): value = value.tolist() refer_edit_pose_meta[key] = value retarget_tpl_pose_metas = retarget_pose(tpl_pose_meta0, refer_pose_meta, tpl_pose_metas_new, tql_edit_pose_meta0, refer_edit_pose_meta) pose_metas = [] for meta in retarget_tpl_pose_metas: pose_meta = AAPoseMeta() width, height = meta["width"], meta["height"] pose_meta.width = width pose_meta.height = height pose_meta.kps_body = np.array(meta["keypoints_body"])[:, :2] * (width, height) pose_meta.kps_body_p = np.array(meta["keypoints_body"])[:, 2] kps_lhand = [] kps_lhand_p = [] for each_kps_lhand in meta["keypoints_left_hand"]: if each_kps_lhand is not None: kps_lhand.append([each_kps_lhand.x, each_kps_lhand.y]) kps_lhand_p.append(each_kps_lhand.score) else: kps_lhand.append([None, None]) kps_lhand_p.append(0.0) pose_meta.kps_lhand = np.array(kps_lhand) pose_meta.kps_lhand_p = np.array(kps_lhand_p) kps_rhand = [] kps_rhand_p = [] for each_kps_rhand in meta["keypoints_right_hand"]: if each_kps_rhand is not None: kps_rhand.append([each_kps_rhand.x, each_kps_rhand.y]) kps_rhand_p.append(each_kps_rhand.score) else: kps_rhand.append([None, None]) kps_rhand_p.append(0.0) pose_meta.kps_rhand = np.array(kps_rhand) pose_meta.kps_rhand_p = np.array(kps_rhand_p) pose_metas.append(pose_meta) return pose_metas