custom_nodes/ComfyUI-WanAnimatePreprocess/retarget_pose.py

# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import numpy as np
from tqdm import tqdm
import math 
from typing import NamedTuple
import copy
from .pose_utils.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