preprocess/pose_align_withdiffaug.py

import os
import sys

this_dir = os.path.dirname(__file__)

mmcv_pkg_root = os.path.join(os.path.dirname(this_dir), "mmcv")
if os.path.exists(mmcv_pkg_root):
    print(f"please make sure you have mmcv package successfully installed in local mmcv folder {mmcv_pkg_root}")
    print(f">>> [check] sys.path before mmcv insert = {sys.path}")
    print(f">>> [check] mmcv_pkg_root = {mmcv_pkg_root}")
    if mmcv_pkg_root in sys.path:
        sys.path.remove(mmcv_pkg_root)
    sys.path.insert(0, mmcv_pkg_root)
    print(f">>> [check] sys.path after mmcv insert = {sys.path}")
else:
    print(f">>> [check] mmcv_pkg_root not exists: {mmcv_pkg_root}")
    print(f"please make sure you have mmcv package successfully installed by 'pip install mmcv' or 'mim install mmcv'")
import mmcv
print(">>> [check] mmcv __file__ =", getattr(mmcv, "__file__", None))
print(">>> [check] mmcv __version__ =", getattr(mmcv, "__version__", None))
assert mmcv.__version__ >= "2.0.0" and mmcv.__version__ < "2.2.0", "mmcv version must be >=2.0.0 and <2.2.0"

import numpy as np
import argparse
import torch
import copy
import cv2
import os
import moviepy.video.io.ImageSequenceClip

from pose.script.dwpose import DWposeDetector, draw_pose
from pose.script.util import size_calculate, warpAffine_kps
from utils_aug import pose_aug_diff



'''
    Detect dwpose from img, then align it by scale parameters
    img: frame from the pose video
    detector: DWpose
    scales: scale parameters
'''
def align_img(img, pose_ori, scales, detect_resolution, image_resolution):

    body_pose = copy.deepcopy(pose_ori['bodies']['candidate'])
    hands = copy.deepcopy(pose_ori['hands'])
    faces = copy.deepcopy(pose_ori['faces'])

    '''
    计算逻辑:
    0. 该函数内进行绝对变换，始终保持人体中心点 body_pose[1] 不变
    1. 先把 ref 和 pose 的高 resize 到一样，且都保持原来的长宽比。
    2. 用点在图中的实际坐标来计算。
    3. 实际计算中，把h的坐标归一化到 [0, 1],  w为[0, W/H]
    4. 由于 dwpose 的输出本来就是归一化的坐标，所以h不需要变，w要乘W/H
    注意：dwpose 输出是 (w, h)
    '''

    # h不变，w缩放到原比例
    H_in, W_in, C_in = img.shape 
    video_ratio = W_in / H_in
    body_pose[:, 0]  = body_pose[:, 0] * video_ratio
    hands[:, :, 0] = hands[:, :, 0] * video_ratio
    faces[:, :, 0] = faces[:, :, 0] * video_ratio

    # scales of 10 body parts 
    scale_neck      = scales["scale_neck"] 
    # scale_face      = scales["scale_face"]
    scale_face_left = scales["scale_face_left"]
    scale_face_right  = scales["scale_face_right"]
    scale_shoulder  = scales["scale_shoulder"]
    scale_arm_upper = scales["scale_arm_upper"]
    scale_arm_lower = scales["scale_arm_lower"]
    scale_hand      = scales["scale_hand"]
    scale_body_len  = scales["scale_body_len"]
    scale_leg_upper = scales["scale_leg_upper"]
    scale_leg_lower = scales["scale_leg_lower"]

    scale_sum = 0
    count = 0
    # scale_list = [scale_neck, scale_face, scale_shoulder, scale_arm_upper, scale_arm_lower, scale_hand, scale_body_len, scale_leg_upper, scale_leg_lower]
    scale_list = [scale_neck, scale_face_left, scale_face_right, scale_shoulder, scale_arm_upper, scale_arm_lower, scale_hand, scale_body_len, scale_leg_upper, scale_leg_lower]
    for i in range(len(scale_list)):
        if not np.isinf(scale_list[i]):
            scale_sum = scale_sum + scale_list[i]
            count = count + 1
    for i in range(len(scale_list)):
        if np.isinf(scale_list[i]):   
            scale_list[i] = scale_sum/count



    # offsets of each part 
    offset = dict()
    # offset["14_15_16_17_to_0"] = body_pose[[14,15,16,17], :] - body_pose[[0], :]
    offset["14_16_to_0"] = body_pose[[14,16], :] - body_pose[[0], :]
    offset["15_17_to_0"] = body_pose[[15,17], :] - body_pose[[0], :]
    offset["3_to_2"] = body_pose[[3], :] - body_pose[[2], :] 
    offset["4_to_3"] = body_pose[[4], :] - body_pose[[3], :] 
    offset["6_to_5"] = body_pose[[6], :] - body_pose[[5], :] 
    offset["7_to_6"] = body_pose[[7], :] - body_pose[[6], :] 
    offset["9_to_8"] = body_pose[[9], :] - body_pose[[8], :] 
    offset["10_to_9"] = body_pose[[10], :] - body_pose[[9], :] 
    offset["12_to_11"] = body_pose[[12], :] - body_pose[[11], :] 
    offset["13_to_12"] = body_pose[[13], :] - body_pose[[12], :] 
    offset["hand_left_to_4"] = hands[1, :, :] - body_pose[[4], :]
    offset["hand_right_to_7"] = hands[0, :, :] - body_pose[[7], :]

    # neck
    c_ = body_pose[1]
    cx = c_[0]
    cy = c_[1]
    M = cv2.getRotationMatrix2D((cx,cy), 0, scale_neck)

    neck = body_pose[[0], :] 
    neck = warpAffine_kps(neck, M)
    body_pose[[0], :] = neck

    # # body_pose_up_shoulder
    # c_ = body_pose[0]
    # cx = c_[0]
    # cy = c_[1]
    # M = cv2.getRotationMatrix2D((cx,cy), 0, scale_face)

    # body_pose_up_shoulder = offset["14_15_16_17_to_0"] + body_pose[[0], :]
    # body_pose_up_shoulder = warpAffine_kps(body_pose_up_shoulder, M)
    # body_pose[[14,15,16,17], :] = body_pose_up_shoulder

    # body_pose_up_shoulder left
    c_ = body_pose[0]
    cx = c_[0]
    cy = c_[1]
    M = cv2.getRotationMatrix2D((cx,cy), 0, scale_face_left)

    body_pose_up_shoulder = offset["14_16_to_0"] + body_pose[[0], :]
    body_pose_up_shoulder = warpAffine_kps(body_pose_up_shoulder, M)
    body_pose[[14,16], :] = body_pose_up_shoulder


    # body_pose_up_shoulder right
    c_ = body_pose[0]
    cx = c_[0]
    cy = c_[1]
    M = cv2.getRotationMatrix2D((cx,cy), 0, scale_face_right)

    body_pose_up_shoulder = offset["15_17_to_0"] + body_pose[[0], :]
    body_pose_up_shoulder = warpAffine_kps(body_pose_up_shoulder, M)
    body_pose[[15,17], :] = body_pose_up_shoulder

    # shoulder 
    c_ = body_pose[1]
    cx = c_[0]
    cy = c_[1]
    M = cv2.getRotationMatrix2D((cx,cy), 0, scale_shoulder)

    body_pose_shoulder = body_pose[[2,5], :] 
    body_pose_shoulder = warpAffine_kps(body_pose_shoulder, M) 
    body_pose[[2,5], :] = body_pose_shoulder

    # arm upper left
    c_ = body_pose[2]
    cx = c_[0]
    cy = c_[1]
    M = cv2.getRotationMatrix2D((cx,cy), 0, scale_arm_upper)
 
    elbow = offset["3_to_2"] + body_pose[[2], :]
    elbow = warpAffine_kps(elbow, M)
    body_pose[[3], :] = elbow

    # arm lower left
    c_ = body_pose[3]
    cx = c_[0]
    cy = c_[1]
    M = cv2.getRotationMatrix2D((cx,cy), 0, scale_arm_lower)
 
    wrist = offset["4_to_3"] + body_pose[[3], :]
    wrist = warpAffine_kps(wrist, M)
    body_pose[[4], :] = wrist

    # hand left
    c_ = body_pose[4]
    cx = c_[0]
    cy = c_[1]
    M = cv2.getRotationMatrix2D((cx,cy), 0, scale_hand)
 
    hand = offset["hand_left_to_4"] + body_pose[[4], :]
    hand = warpAffine_kps(hand, M)
    hands[1, :, :] = hand

    # arm upper right
    c_ = body_pose[5]
    cx = c_[0]
    cy = c_[1]
    M = cv2.getRotationMatrix2D((cx,cy), 0, scale_arm_upper)
 
    elbow = offset["6_to_5"] + body_pose[[5], :]
    elbow = warpAffine_kps(elbow, M)
    body_pose[[6], :] = elbow

    # arm lower right
    c_ = body_pose[6]
    cx = c_[0]
    cy = c_[1]
    M = cv2.getRotationMatrix2D((cx,cy), 0, scale_arm_lower)
 
    wrist = offset["7_to_6"] + body_pose[[6], :]
    wrist = warpAffine_kps(wrist, M)
    body_pose[[7], :] = wrist

    # hand right
    c_ = body_pose[7]
    cx = c_[0]
    cy = c_[1]
    M = cv2.getRotationMatrix2D((cx,cy), 0, scale_hand)
 
    hand = offset["hand_right_to_7"] + body_pose[[7], :]
    hand = warpAffine_kps(hand, M)
    hands[0, :, :] = hand

    # body len
    c_ = body_pose[1]
    cx = c_[0]
    cy = c_[1]
    M = cv2.getRotationMatrix2D((cx,cy), 0, scale_body_len)

    body_len = body_pose[[8,11], :] 
    body_len = warpAffine_kps(body_len, M)
    body_pose[[8,11], :] = body_len

    # leg upper left
    c_ = body_pose[8]
    cx = c_[0]
    cy = c_[1]
    M = cv2.getRotationMatrix2D((cx,cy), 0, scale_leg_upper)
 
    knee = offset["9_to_8"] + body_pose[[8], :]
    knee = warpAffine_kps(knee, M)
    body_pose[[9], :] = knee

    # leg lower left
    c_ = body_pose[9]
    cx = c_[0]
    cy = c_[1]
    M = cv2.getRotationMatrix2D((cx,cy), 0, scale_leg_lower)
 
    ankle = offset["10_to_9"] + body_pose[[9], :]
    ankle = warpAffine_kps(ankle, M)
    body_pose[[10], :] = ankle

    # leg upper right
    c_ = body_pose[11]
    cx = c_[0]
    cy = c_[1]
    M = cv2.getRotationMatrix2D((cx,cy), 0, scale_leg_upper)
 
    knee = offset["12_to_11"] + body_pose[[11], :]
    knee = warpAffine_kps(knee, M)
    body_pose[[12], :] = knee

    # leg lower right
    c_ = body_pose[12]
    cx = c_[0]
    cy = c_[1]
    M = cv2.getRotationMatrix2D((cx,cy), 0, scale_leg_lower)
 
    ankle = offset["13_to_12"] + body_pose[[12], :]
    ankle = warpAffine_kps(ankle, M)
    body_pose[[13], :] = ankle

    # none part
    body_pose_none = pose_ori['bodies']['candidate'] == -1.
    hands_none = pose_ori['hands'] == -1.
    faces_none = pose_ori['faces'] == -1.

    body_pose[body_pose_none] = -1.
    hands[hands_none] = -1. 
    nan = float('nan')
    if len(hands[np.isnan(hands)]) > 0:
        print('nan')
    faces[faces_none] = -1.

    # last check nan -> -1.
    body_pose = np.nan_to_num(body_pose, nan=-1.)
    hands = np.nan_to_num(hands, nan=-1.)
    faces = np.nan_to_num(faces, nan=-1.)

    # return
    pose_align = copy.deepcopy(pose_ori)
    pose_align['bodies']['candidate'] = body_pose
    pose_align['hands'] = hands
    pose_align['faces'] = faces

    return pose_align



def run_align_video_with_filterPose_translate_smooth(args):

    vidfn=args.vidfn
    imgfn_refer=args.imgfn_refer
    outfn=args.outfn
    
    video = cv2.VideoCapture(vidfn)
    width= video.get(cv2.CAP_PROP_FRAME_WIDTH)
    height= video.get(cv2.CAP_PROP_FRAME_HEIGHT)
 
    total_frame= video.get(cv2.CAP_PROP_FRAME_COUNT)
    fps= video.get(cv2.CAP_PROP_FPS)

    print("height:", height)
    print("width:", width)
    print("fps:", fps)

    H_in, W_in  = height, width
    H_out, W_out = size_calculate(H_in,W_in,args.detect_resolution) 
    H_out, W_out = size_calculate(H_out,W_out,args.image_resolution) 

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    detector = DWposeDetector(
        det_config = args.yolox_config, 
        det_ckpt = args.yolox_ckpt,
        pose_config = args.dwpose_config, 
        pose_ckpt = args.dwpose_ckpt, 
        keypoints_only=False
        )    
    detector = detector.to(device)

    #### refer_img 的精确处理 在前边直接处理完毕，就不需要考虑后续 pose 的二次处理 ####
    refer_img = cv2.imread(imgfn_refer)
    ref_height, ref_width, channels = refer_img.shape
    # print("ref_height: ", ref_height)
    # print("ref_width: ", ref_width)
    aspect_ratio = ref_height / ref_width
    # max_area = "832*480"
    max_area = "1024*576"
    # max_area = "1664*960"
    lat_h = round(
        np.sqrt(int(eval(max_area)) * aspect_ratio) // 16)
    lat_w = round(
        np.sqrt(int(eval(max_area)) / aspect_ratio) // 16)
    new_height = lat_h * 16
    new_width = lat_w * 16
    # print("new_height:", new_height)
    # print("new_width:", new_width)
    # resize_height = int(ref_height*(1.0 * new_height/ref_height))
    # resize_width = int(ref_width*(1.0 * new_width/ref_width))
    # resize_height = new_height * 2
    # resize_width = new_width * 2
    resize_height = new_height
    resize_width = new_width
    # print("resize_height:", resize_height)
    # print("resize_width:", resize_width)
    refer_img = cv2.resize(refer_img, (resize_width, resize_height), interpolation=cv2.INTER_CUBIC)
    ref_height, ref_width, channels = refer_img.shape

    output_refer, pose_refer = detector(refer_img,detect_resolution=args.detect_resolution, image_resolution=args.image_resolution, output_type='cv2',return_pose_dict=True)
    body_ref_img  = pose_refer['bodies']['candidate']
    hands_ref_img = pose_refer['hands']
    faces_ref_img = pose_refer['faces']
    output_refer = cv2.cvtColor(output_refer, cv2.COLOR_RGB2BGR)
    

    skip_frames = args.align_frame
    max_frame = args.max_frame
    pose_list, video_frame_buffer, video_pose_buffer = [], [], []


    cap = cv2.VideoCapture('2.mp4')     # 读取视频
    while cap.isOpened():               # 当视频被打开时：
        ret, frame = cap.read()         # 读取视频，读取到的某一帧存储到frame，若是读取成功，ret为True，反之为False
        if ret:                         # 若是读取成功
            cv2.imshow('frame', frame)  # 显示读取到的这一帧画面
            key = cv2.waitKey(25)       # 等待一段时间，并且检测键盘输入
            if key == ord('q'):         # 若是键盘输入'q',则退出，释放视频
                cap.release()           # 释放视频
                break
        else:
            cap.release()
    cv2.destroyAllWindows()             # 关闭所有窗口


    for i in range(max_frame):
        ret, img = video.read()
        if img is None: 
            break 
        else: 
            if i < skip_frames:
                continue           
            video_frame_buffer.append(img)


       
        # estimate scale parameters by the 1st frame in the video
        if i==skip_frames:
            output_1st_img, pose_1st_img = detector(img, args.detect_resolution, args.image_resolution, output_type='cv2', return_pose_dict=True)
            body_1st_img  = pose_1st_img['bodies']['candidate']
            hands_1st_img = pose_1st_img['hands']
            faces_1st_img = pose_1st_img['faces']

            '''
            计算逻辑:
            1. 先把 ref 和 pose 的高 resize 到一样，且都保持原来的长宽比。
            2. 用点在图中的实际坐标来计算。
            3. 实际计算中，把h的坐标归一化到 [0, 1],  w为[0, W/H]
            4. 由于 dwpose 的输出本来就是归一化的坐标，所以h不需要变，w要乘W/H
            注意：dwpose 输出是 (w, h)
            '''
            
            # h不变，w缩放到原比例
            ref_H, ref_W = refer_img.shape[0], refer_img.shape[1]
            ref_ratio = ref_W / ref_H
            body_ref_img[:, 0]  = body_ref_img[:, 0] * ref_ratio
            hands_ref_img[:, :, 0] = hands_ref_img[:, :, 0] * ref_ratio
            faces_ref_img[:, :, 0] = faces_ref_img[:, :, 0] * ref_ratio

            video_ratio = width / height
            body_1st_img[:, 0]  = body_1st_img[:, 0] * video_ratio
            hands_1st_img[:, :, 0] = hands_1st_img[:, :, 0] * video_ratio
            faces_1st_img[:, :, 0] = faces_1st_img[:, :, 0] * video_ratio

            # scale
            align_args = dict()
            
            dist_1st_img = np.linalg.norm(body_1st_img[0]-body_1st_img[1])   # 0.078   
            dist_ref_img = np.linalg.norm(body_ref_img[0]-body_ref_img[1])   # 0.106
            align_args["scale_neck"] = dist_ref_img / dist_1st_img  # align / pose = ref / 1st

            # dist_1st_img = np.linalg.norm(body_1st_img[16]-body_1st_img[17])
            # dist_ref_img = np.linalg.norm(body_ref_img[16]-body_ref_img[17])
            # align_args["scale_face"] = dist_ref_img / dist_1st_img

            dist_1st_img = np.linalg.norm(body_1st_img[16]-body_1st_img[14]) + np.linalg.norm(body_1st_img[14]-body_1st_img[0])
            dist_ref_img = np.linalg.norm(body_ref_img[16]-body_ref_img[14]) + np.linalg.norm(body_ref_img[14]-body_ref_img[0])
            align_args["scale_face_left"] = dist_ref_img / dist_1st_img

            dist_1st_img = np.linalg.norm(body_1st_img[17]-body_1st_img[15]) + np.linalg.norm(body_1st_img[15]-body_1st_img[0])
            dist_ref_img = np.linalg.norm(body_ref_img[17]-body_ref_img[15]) + np.linalg.norm(body_ref_img[15]-body_ref_img[0])
            align_args["scale_face_right"] = dist_ref_img / dist_1st_img

            dist_1st_img = np.linalg.norm(body_1st_img[2]-body_1st_img[5])  # 0.112
            dist_ref_img = np.linalg.norm(body_ref_img[2]-body_ref_img[5])  # 0.174
            align_args["scale_shoulder"] = dist_ref_img / dist_1st_img  

            dist_1st_img = np.linalg.norm(body_1st_img[2]-body_1st_img[3])  # 0.895
            dist_ref_img = np.linalg.norm(body_ref_img[2]-body_ref_img[3])  # 0.134
            s1 = dist_ref_img / dist_1st_img
            dist_1st_img = np.linalg.norm(body_1st_img[5]-body_1st_img[6])
            dist_ref_img = np.linalg.norm(body_ref_img[5]-body_ref_img[6])
            s2 = dist_ref_img / dist_1st_img
            align_args["scale_arm_upper"] = (s1+s2)/2 # 1.548

            dist_1st_img = np.linalg.norm(body_1st_img[3]-body_1st_img[4])
            dist_ref_img = np.linalg.norm(body_ref_img[3]-body_ref_img[4])
            s1 = dist_ref_img / dist_1st_img
            dist_1st_img = np.linalg.norm(body_1st_img[6]-body_1st_img[7])
            dist_ref_img = np.linalg.norm(body_ref_img[6]-body_ref_img[7])
            s2 = dist_ref_img / dist_1st_img
            align_args["scale_arm_lower"] = (s1+s2)/2

            # hand
            dist_1st_img = np.zeros(10)
            dist_ref_img = np.zeros(10)      
             
            dist_1st_img[0] = np.linalg.norm(hands_1st_img[0,0]-hands_1st_img[0,1])
            dist_1st_img[1] = np.linalg.norm(hands_1st_img[0,0]-hands_1st_img[0,5])
            dist_1st_img[2] = np.linalg.norm(hands_1st_img[0,0]-hands_1st_img[0,9])
            dist_1st_img[3] = np.linalg.norm(hands_1st_img[0,0]-hands_1st_img[0,13])
            dist_1st_img[4] = np.linalg.norm(hands_1st_img[0,0]-hands_1st_img[0,17])
            dist_1st_img[5] = np.linalg.norm(hands_1st_img[1,0]-hands_1st_img[1,1])
            dist_1st_img[6] = np.linalg.norm(hands_1st_img[1,0]-hands_1st_img[1,5])
            dist_1st_img[7] = np.linalg.norm(hands_1st_img[1,0]-hands_1st_img[1,9])
            dist_1st_img[8] = np.linalg.norm(hands_1st_img[1,0]-hands_1st_img[1,13])
            dist_1st_img[9] = np.linalg.norm(hands_1st_img[1,0]-hands_1st_img[1,17])

            dist_ref_img[0] = np.linalg.norm(hands_ref_img[0,0]-hands_ref_img[0,1])
            dist_ref_img[1] = np.linalg.norm(hands_ref_img[0,0]-hands_ref_img[0,5])
            dist_ref_img[2] = np.linalg.norm(hands_ref_img[0,0]-hands_ref_img[0,9])
            dist_ref_img[3] = np.linalg.norm(hands_ref_img[0,0]-hands_ref_img[0,13])
            dist_ref_img[4] = np.linalg.norm(hands_ref_img[0,0]-hands_ref_img[0,17])
            dist_ref_img[5] = np.linalg.norm(hands_ref_img[1,0]-hands_ref_img[1,1])
            dist_ref_img[6] = np.linalg.norm(hands_ref_img[1,0]-hands_ref_img[1,5])
            dist_ref_img[7] = np.linalg.norm(hands_ref_img[1,0]-hands_ref_img[1,9])
            dist_ref_img[8] = np.linalg.norm(hands_ref_img[1,0]-hands_ref_img[1,13])
            dist_ref_img[9] = np.linalg.norm(hands_ref_img[1,0]-hands_ref_img[1,17])

            ratio = 0   
            count = 0
            total_iters = 0 # 10
            for i in range (total_iters):
                if dist_1st_img[i] != 0:
                    ratio = ratio + dist_ref_img[i]/dist_1st_img[i]
                    count = count + 1
            if count!=0:
                align_args["scale_hand"] = (ratio/count+align_args["scale_arm_upper"]+align_args["scale_arm_lower"])/3
            else:
                align_args["scale_hand"] = (align_args["scale_arm_upper"]+align_args["scale_arm_lower"])/2

            # body 
            dist_1st_img = np.linalg.norm(body_1st_img[1] - (body_1st_img[8] + body_1st_img[11])/2 )
            dist_ref_img = np.linalg.norm(body_ref_img[1] - (body_ref_img[8] + body_ref_img[11])/2 )
            align_args["scale_body_len"]=dist_ref_img / dist_1st_img

            dist_1st_img = np.linalg.norm(body_1st_img[8]-body_1st_img[9])
            dist_ref_img = np.linalg.norm(body_ref_img[8]-body_ref_img[9])
            s1 = dist_ref_img / dist_1st_img
            dist_1st_img = np.linalg.norm(body_1st_img[11]-body_1st_img[12])
            dist_ref_img = np.linalg.norm(body_ref_img[11]-body_ref_img[12])
            s2 = dist_ref_img / dist_1st_img
            align_args["scale_leg_upper"] = (s1+s2)/2

            dist_1st_img = np.linalg.norm(body_1st_img[9]-body_1st_img[10])
            dist_ref_img = np.linalg.norm(body_ref_img[9]-body_ref_img[10])
            s1 = dist_ref_img / dist_1st_img
            dist_1st_img = np.linalg.norm(body_1st_img[12]-body_1st_img[13])
            dist_ref_img = np.linalg.norm(body_ref_img[12]-body_ref_img[13])
            s2 = dist_ref_img / dist_1st_img
            align_args["scale_leg_lower"] = (s1+s2)/2

            ####################
            ####################
            # need adjust nan
            for k,v in align_args.items():
                if np.isnan(v):
                    align_args[k]=1

            # centre offset (the offset of key point 1)
            offset = body_ref_img[1] - body_1st_img[1]
        
    
        # pose align
        pose_img, pose_ori = detector(img, args.detect_resolution, args.image_resolution, output_type='cv2', return_pose_dict=True)
        video_pose_buffer.append(pose_img)
        pose_align = align_img(img, pose_ori, align_args, args.detect_resolution, args.image_resolution)
        

        # add centre offset
        pose = pose_align
        pose['bodies']['candidate'] = pose['bodies']['candidate'] + offset
        pose['hands'] = pose['hands'] + offset
        pose['faces'] = pose['faces'] + offset


        # h不变，w从绝对坐标缩放回0-1 注意这里要回到ref的坐标系
        pose['bodies']['candidate'][:, 0] = pose['bodies']['candidate'][:, 0] / ref_ratio
        pose['hands'][:, :, 0] = pose['hands'][:, :, 0] / ref_ratio
        pose['faces'][:, :, 0] = pose['faces'][:, :, 0] / ref_ratio
        pose_list.append(pose)

    # stack
    body_list  = [pose['bodies']['candidate'][:18] for pose in pose_list]
    body_list_subset = [pose['bodies']['subset'][:1] for pose in pose_list]
    hands_list = [pose['hands'][:2] for pose in pose_list]
    faces_list = [pose['faces'][:1] for pose in pose_list]
   
    body_seq         = np.stack(body_list       , axis=0)
    body_seq_subset  = np.stack(body_list_subset, axis=0)
    hands_seq        = np.stack(hands_list      , axis=0)
    faces_seq        = np.stack(faces_list      , axis=0)


    # concatenate and paint results
    # H = 768 # paint height
    H = ref_H # paint height
    W1 = int((H/ref_H * ref_W)//2 *2)
    W2 = int((H/height * width)//2 *2)
    result_demo = [] # = Writer(args, None, H, 3*W1+2*W2, outfn, fps)
    result_pose_only = [] # Writer(args, None, H, W1, args.outfn_align_pose_video, fps)
    result_pose_aug_only = [] # Writer(args, None, H, W1, args.outfn_align_pose_video[:-4] + "_aug" + ".mp4", fps)

    offset_x=(-0.2,0.2)
    offset_y=(-0.2,0.2)
    scale=(0.7,1.3)
    aspect_ratio_range=(0.6, 1.4)
    offset = (offset_x, offset_y)

    for i in range(len(body_seq)):
        pose_t={}
        pose_t["bodies"]={}
        pose_t["bodies"]["candidate"]=body_seq[i]
        pose_t["bodies"]["subset"]=body_seq_subset[i]
        pose_t["hands"]=hands_seq[i]
        pose_t["faces"]=faces_seq[i]

        ref_img = cv2.cvtColor(refer_img, cv2.COLOR_RGB2BGR)
        ref_img = cv2.resize(ref_img, (W1, H))
        ref_pose= cv2.resize(output_refer, (W1, H))
        
        # output_transformed = draw_pose(
        #     pose_t, 
        #     int(H_in*1024/W_in), 
        #     1024, 
        #     draw_face=False,
        #     )
        # output_transformed = cv2.cvtColor(output_transformed, cv2.COLOR_BGR2RGB)
        # output_transformed = cv2.resize(output_transformed, (W1, H))
        
        output_transformed = draw_pose(     # single.mp4
            pose_t,
            ref_H*2,
            ref_W*2,
            draw_face=False,
            )
        output_transformed = cv2.cvtColor(output_transformed, cv2.COLOR_BGR2RGB)
        # output_transformed = cv2.resize(output_transformed, (W1, H), interpolation=cv2.INTER_CUBIC)
        
        output_transformed_1 = draw_pose(   # all.mp4
            pose_t,
            ref_H,
            ref_W,
            draw_face=False,
            )
        output_transformed_1 = cv2.cvtColor(output_transformed_1, cv2.COLOR_BGR2RGB)
        # output_transformed_1 = cv2.resize(output_transformed_1, (W1, H), interpolation=cv2.INTER_CUBIC)

        pose_t_aug = pose_aug_diff(pose_t.copy(), size=(ref_H, ref_W), offset=offset, scale=scale, aspect_ratio_range=aspect_ratio_range, add_aug=True)
        
        output_transformed_aug = draw_pose(     # single_aug.mp4
            pose_t_aug,
            ref_H*2,
            ref_W*2,
            draw_face=False,
            )
        output_transformed_aug = cv2.cvtColor(output_transformed_aug, cv2.COLOR_BGR2RGB)
        # output_transformed_aug = cv2.resize(output_transformed_aug, (W1, H), interpolation=cv2.INTER_CUBIC)

        output_transformed_aug_1 = draw_pose(   # all.mp4
            pose_t_aug,
            ref_H,
            ref_W,
            draw_face=False,
            )
        output_transformed_aug_1 = cv2.cvtColor(output_transformed_aug_1, cv2.COLOR_BGR2RGB)
        # output_transformed_aug_1 = cv2.resize(output_transformed_aug_1, (W1, H), interpolation=cv2.INTER_CUBIC)
        
        video_frame = cv2.resize(video_frame_buffer[i], (W2, H), interpolation=cv2.INTER_CUBIC)
        video_frame = cv2.cvtColor(video_frame, cv2.COLOR_BGR2RGB)
        video_pose  = cv2.resize(video_pose_buffer[i], (W2, H), interpolation=cv2.INTER_CUBIC)

        if (4*W1 + 2*W2) <= 16384:
            res = np.concatenate([ref_img, ref_pose, output_transformed_1, output_transformed_aug_1, video_frame, video_pose], axis=1)
        else:
            res = np.concatenate([ref_img, ref_pose, output_transformed_1, output_transformed_aug_1, video_frame], axis=1)
        result_demo.append(res) # all.mp4
        result_pose_only.append(output_transformed) # single.mp4
        result_pose_aug_only.append(output_transformed_aug) # single_aug.mp4

    print(f"pose_list len: {len(pose_list)}")
    result_demo = [frame.astype('uint8') for frame in result_demo]
    result_pose_only = [frame.astype('uint8') for frame in result_pose_only]
    result_pose_aug_only = [frame.astype('uint8') for frame in result_pose_aug_only]
    
    clip = moviepy.video.io.ImageSequenceClip.ImageSequenceClip(result_demo, fps=fps)
    clip.write_videofile(outfn, fps=fps, codec="libx264", audio=False, logger=None) # all.mp4
    
    clip = moviepy.video.io.ImageSequenceClip.ImageSequenceClip(result_pose_only, fps=fps)
    clip.write_videofile(args.outfn_align_pose_video, fps=fps, codec="libx264", audio=False, logger=None) # single.mp4
    
    clip = moviepy.video.io.ImageSequenceClip.ImageSequenceClip(result_pose_aug_only, fps=fps)
    clip.write_videofile(args.outfn_align_pose_video[:-4] + "_aug" + ".mp4", fps=fps, codec="libx264", audio=False, logger=None)    # single_aug.mp4
    
    print('pose align done')



def main():
    parser = argparse.ArgumentParser()
    # parser.add_argument('--detect_resolution', type=int, default=512, help='detect_resolution')
    # parser.add_argument('--image_resolution', type=int, default=720, help='image_resolution')
    parser.add_argument('--detect_resolution', type=int, default=1024, help='detect_resolution')
    parser.add_argument('--image_resolution', type=int, default=720, help='image_resolution')

    parser.add_argument("--yolox_config",  type=str, default=f"{this_dir}/pose/config/yolox_l_8xb8-300e_coco.py")
    parser.add_argument("--dwpose_config", type=str, default=f"{this_dir}/pose/config/dwpose-l_384x288.py")
    parser.add_argument("--yolox_ckpt",  type=str, default=f"{this_dir}/pretrained_weights/dwpose/yolox_l_8x8_300e_coco.pth")
    parser.add_argument("--dwpose_ckpt", type=str, default=f"{this_dir}/pretrained_weights/dwpose/dw-ll_ucoco_384.pth")


    parser.add_argument('--align_frame', type=int, default=0, help='the frame index of the video to align')
    parser.add_argument('--max_frame', type=int, default=300, help='maximum frame number of the video to align')
    parser.add_argument('--imgfn_refer', type=str, default="./assets/images/0.jpg", help='refer image path')
    parser.add_argument('--vidfn', type=str, default="./assets/videos/0.mp4", help='Input video path')
    parser.add_argument('--outfn_align_pose_video', type=str, default=None, help='output path of the aligned video of the refer img')
    parser.add_argument('--outfn', type=str, default=None, help='Output path of the alignment visualization')
    args = parser.parse_args()
    
    # if not os.path.exists("./assets/poses/align"):
    #     # os.makedirs("./assets/poses/")
    #     os.makedirs("./assets/poses/align")
    #     os.makedirs("./assets/poses/align_demo")
        
    img_name = os.path.basename(args.imgfn_refer).split('.')[0]
    video_name = os.path.basename(args.vidfn).split('.')[0]
    if args.outfn_align_pose_video is None:
        args.outfn_align_pose_video = "./assets/poses/align/img_{}_video_{}.mp4".format(img_name, video_name)
    if args.outfn is None:
        args.outfn = "./assets/poses/align_demo/img_{}_video_{}.mp4".format(img_name, video_name)

    os.makedirs(os.path.dirname(args.outfn), exist_ok=True)
    os.makedirs(os.path.dirname(args.outfn_align_pose_video), exist_ok=True)

    run_align_video_with_filterPose_translate_smooth(args)


    
if __name__ == '__main__':
    main()