# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. import os import cv2 import numpy as np import json from tqdm import tqdm import math from typing import NamedTuple, List import copy from pose2d_utils import AAPoseMeta # 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: 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: 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: 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: 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: pass elif rescale_keypoints[4] == None and rescale_keypoints[7] != None: 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: 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: 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: 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: 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: # 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: 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: 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: 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: 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: 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: 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): 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: 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 and src_skeleton['keypoints_body'][10] != None and src_skeleton['keypoints_body'][13] != None and \ dst_skeleton['keypoints_body'][0] != None and dst_skeleton['keypoints_body'][10] != None and dst_skeleton['keypoints_body'][13] != None 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 and src_skeleton['keypoints_body'][8] != None and src_skeleton['keypoints_body'][11] != None and \ dst_skeleton['keypoints_body'][0] != None and dst_skeleton['keypoints_body'][8] != None and dst_skeleton['keypoints_body'][11] != None 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 and src_skeleton_edit['keypoints_body'][10] != None and src_skeleton_edit['keypoints_body'][13] != None and \ dst_skeleton_edit['keypoints_body'][0] != None and dst_skeleton_edit['keypoints_body'][10] != None and dst_skeleton_edit['keypoints_body'][13] != None 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 and src_skeleton_edit['keypoints_body'][8] != None and src_skeleton_edit['keypoints_body'][11] != None and \ dst_skeleton_edit['keypoints_body'][0] != None and dst_skeleton_edit['keypoints_body'][8] != None and dst_skeleton_edit['keypoints_body'][11] != None 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: 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: 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: 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: 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