# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import os
import cv2
import time
import math
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from typing import Dict, List
import random
from pose2d_utils import AAPoseMeta


def draw_handpose(canvas, keypoints, hand_score_th=0.6):
    """
    Draw keypoints and connections representing hand pose on a given canvas.

    Args:
        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the hand pose.
        keypoints (List[Keypoint]| None): A list of Keypoint objects representing the hand keypoints to be drawn
                                          or None if no keypoints are present.

    Returns:
        np.ndarray: A 3D numpy array representing the modified canvas with the drawn hand pose.

    Note:
        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
    """
    eps = 0.01

    H, W, C = canvas.shape
    stickwidth = max(int(min(H, W) / 200), 1)

    edges = [
        [0, 1],
        [1, 2],
        [2, 3],
        [3, 4],
        [0, 5],
        [5, 6],
        [6, 7],
        [7, 8],
        [0, 9],
        [9, 10],
        [10, 11],
        [11, 12],
        [0, 13],
        [13, 14],
        [14, 15],
        [15, 16],
        [0, 17],
        [17, 18],
        [18, 19],
        [19, 20],
    ]

    for ie, (e1, e2) in enumerate(edges):
        k1 = keypoints[e1]
        k2 = keypoints[e2]
        if k1 is None or k2 is None:
            continue
        if k1[2] < hand_score_th or k2[2] < hand_score_th:
            continue

        x1 = int(k1[0])
        y1 = int(k1[1])
        x2 = int(k2[0])
        y2 = int(k2[1])
        if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
            cv2.line(
                canvas,
                (x1, y1),
                (x2, y2),
                matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * 255,
                thickness=stickwidth,
            )

    for keypoint in keypoints:

        if keypoint is None:
            continue
        if keypoint[2] < hand_score_th:
            continue

        x, y = keypoint[0], keypoint[1]
        x = int(x)
        y = int(y)
        if x > eps and y > eps:
            cv2.circle(canvas, (x, y), stickwidth, (0, 0, 255), thickness=-1)
    return canvas


def draw_handpose_new(canvas, keypoints, stickwidth_type='v2', hand_score_th=0.6):
    """
    Draw keypoints and connections representing hand pose on a given canvas.

    Args:
        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the hand pose.
        keypoints (List[Keypoint]| None): A list of Keypoint objects representing the hand keypoints to be drawn
                                          or None if no keypoints are present.

    Returns:
        np.ndarray: A 3D numpy array representing the modified canvas with the drawn hand pose.

    Note:
        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
    """
    eps = 0.01

    H, W, C = canvas.shape
    if stickwidth_type == 'v1':
        stickwidth = max(int(min(H, W) / 200), 1)
    elif stickwidth_type == 'v2':
        stickwidth = max(max(int(min(H, W) / 200) - 1, 1) // 2, 1)

    edges = [
        [0, 1],
        [1, 2],
        [2, 3],
        [3, 4],
        [0, 5],
        [5, 6],
        [6, 7],
        [7, 8],
        [0, 9],
        [9, 10],
        [10, 11],
        [11, 12],
        [0, 13],
        [13, 14],
        [14, 15],
        [15, 16],
        [0, 17],
        [17, 18],
        [18, 19],
        [19, 20],
    ]

    for ie, (e1, e2) in enumerate(edges):
        k1 = keypoints[e1]
        k2 = keypoints[e2]
        if k1 is None or k2 is None:
            continue
        if k1[2] < hand_score_th or k2[2] < hand_score_th:
            continue

        x1 = int(k1[0])
        y1 = int(k1[1])
        x2 = int(k2[0])
        y2 = int(k2[1])
        if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
            cv2.line(
                canvas,
                (x1, y1),
                (x2, y2),
                matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * 255,
                thickness=stickwidth,
            )

    for keypoint in keypoints:

        if keypoint is None:
            continue
        if keypoint[2] < hand_score_th:
            continue

        x, y = keypoint[0], keypoint[1]
        x = int(x)
        y = int(y)
        if x > eps and y > eps:
            cv2.circle(canvas, (x, y), stickwidth, (0, 0, 255), thickness=-1)
    return canvas


def draw_ellipse_by_2kp(img, keypoint1, keypoint2, color, threshold=0.6):
    H, W, C = img.shape
    stickwidth = max(int(min(H, W) / 200), 1)

    if keypoint1[-1] < threshold or keypoint2[-1] < threshold:
        return img

    Y = np.array([keypoint1[0], keypoint2[0]])
    X = np.array([keypoint1[1], keypoint2[1]])
    mX = np.mean(X)
    mY = np.mean(Y)
    length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
    angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
    polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
    cv2.fillConvexPoly(img, polygon, [int(float(c) * 0.6) for c in color])
    return img


def split_pose2d_kps_to_aa(kp2ds: np.ndarray) -> List[np.ndarray]:
    """Convert the 133 keypoints from pose2d to body and hands keypoints.

    Args:
        kp2ds (np.ndarray): [133, 2]

    Returns:
        List[np.ndarray]: _description_
    """
    kp2ds_body = (
        kp2ds[[0, 6, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 17, 20]]
        + kp2ds[[0, 5, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 18, 21]]
    ) / 2
    kp2ds_lhand = kp2ds[91:112]
    kp2ds_rhand = kp2ds[112:133]
    return kp2ds_body.copy(), kp2ds_lhand.copy(), kp2ds_rhand.copy()


def draw_aapose_by_meta(img, meta: AAPoseMeta, threshold=0.5, stick_width_norm=200, draw_hand=True, draw_head=True):
    kp2ds = np.concatenate([meta.kps_body, meta.kps_body_p[:, None]], axis=1)
    kp2ds_lhand = np.concatenate([meta.kps_lhand, meta.kps_lhand_p[:, None]], axis=1)
    kp2ds_rhand = np.concatenate([meta.kps_rhand, meta.kps_rhand_p[:, None]], axis=1)
    pose_img = draw_aapose(img, kp2ds, threshold, kp2ds_lhand=kp2ds_lhand, kp2ds_rhand=kp2ds_rhand, stick_width_norm=stick_width_norm, draw_hand=draw_hand, draw_head=draw_head)
    return pose_img

def draw_aapose_by_meta_new(img, meta: AAPoseMeta, threshold=0.5, stickwidth_type='v2', draw_hand=True, draw_head=True):
    kp2ds = np.concatenate([meta.kps_body, meta.kps_body_p[:, None]], axis=1)
    kp2ds_lhand = np.concatenate([meta.kps_lhand, meta.kps_lhand_p[:, None]], axis=1)
    kp2ds_rhand = np.concatenate([meta.kps_rhand, meta.kps_rhand_p[:, None]], axis=1)
    pose_img = draw_aapose_new(img, kp2ds, threshold, kp2ds_lhand=kp2ds_lhand, kp2ds_rhand=kp2ds_rhand,
                               stickwidth_type=stickwidth_type, draw_hand=draw_hand, draw_head=draw_head)
    return pose_img

def draw_hand_by_meta(img, meta: AAPoseMeta, threshold=0.5, stick_width_norm=200):
    kp2ds = np.concatenate([meta.kps_body, meta.kps_body_p[:, None] * 0], axis=1)
    kp2ds_lhand = np.concatenate([meta.kps_lhand, meta.kps_lhand_p[:, None]], axis=1)
    kp2ds_rhand = np.concatenate([meta.kps_rhand, meta.kps_rhand_p[:, None]], axis=1)
    pose_img = draw_aapose(img, kp2ds, threshold, kp2ds_lhand=kp2ds_lhand, kp2ds_rhand=kp2ds_rhand, stick_width_norm=stick_width_norm, draw_hand=True, draw_head=False)
    return pose_img


def draw_aaface_by_meta(img, meta: AAPoseMeta, threshold=0.5, stick_width_norm=200, draw_hand=False, draw_head=True):
    kp2ds = np.concatenate([meta.kps_body, meta.kps_body_p[:, None]], axis=1)
    # kp2ds_lhand = np.concatenate([meta.kps_lhand, meta.kps_lhand_p[:, None]], axis=1)
    # kp2ds_rhand = np.concatenate([meta.kps_rhand, meta.kps_rhand_p[:, None]], axis=1)
    pose_img = draw_M(img, kp2ds, threshold, kp2ds_lhand=None, kp2ds_rhand=None, stick_width_norm=stick_width_norm, draw_hand=draw_hand, draw_head=draw_head)
    return pose_img


def draw_aanose_by_meta(img, meta: AAPoseMeta, threshold=0.5, stick_width_norm=100, draw_hand=False):
    kp2ds = np.concatenate([meta.kps_body, meta.kps_body_p[:, None]], axis=1)
    # kp2ds_lhand = np.concatenate([meta.kps_lhand, meta.kps_lhand_p[:, None]], axis=1)
    # kp2ds_rhand = np.concatenate([meta.kps_rhand, meta.kps_rhand_p[:, None]], axis=1)
    pose_img = draw_nose(img, kp2ds, threshold, kp2ds_lhand=None, kp2ds_rhand=None, stick_width_norm=stick_width_norm, draw_hand=draw_hand)
    return pose_img


def gen_face_motion_seq(img, metas: List[AAPoseMeta], threshold=0.5, stick_width_norm=200):

    return 


def draw_M(
    img,
    kp2ds,
    threshold=0.6,
    data_to_json=None,
    idx=-1,
    kp2ds_lhand=None,
    kp2ds_rhand=None,
    draw_hand=False,
    stick_width_norm=200,
    draw_head=True
):
    """
    Draw keypoints and connections representing hand pose on a given canvas.

    Args:
        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the hand pose.
        keypoints (List[Keypoint]| None): A list of Keypoint objects representing the hand keypoints to be drawn
                                          or None if no keypoints are present.

    Returns:
        np.ndarray: A 3D numpy array representing the modified canvas with the drawn hand pose.

    Note:
        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
    """

    new_kep_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",
    ]
    # kp2ds_body = (kp2ds.copy()[[0, 6, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 17, 20]] + \
    #              kp2ds.copy()[[0, 5, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 18, 21]]) / 2
    kp2ds = kp2ds.copy()
    # import ipdb; ipdb.set_trace()
    kp2ds[[1,2,3,4,5,6,7,8,9,10,11,12,13,18,19], 2] = 0
    if not draw_head:
        kp2ds[[0,14,15,16,17], 2] = 0
    kp2ds_body = kp2ds
    # kp2ds_body = kp2ds_body[:18]

    # kp2ds_lhand = kp2ds.copy()[91:112]
    # kp2ds_rhand = kp2ds.copy()[112:133]

    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],
        # [11, 20],  # foot
    ]

    colors = [
        # [255, 0, 0],
        # [255, 85, 0],
        # [255, 170, 0],
        # [255, 255, 0],
        # [170, 255, 0],
        # [85, 255, 0],
        # [0, 255, 0],
        # [0, 255, 85],
        # [0, 255, 170],
        # [0, 255, 255],
        # [0, 170, 255],
        # [0, 85, 255],
        # [0, 0, 255],
        # [85, 0, 255],
        [170, 0, 255],
        [255, 0, 255],
        [255, 0, 170],
        [255, 0, 85],
        # foot
        # [200, 200, 0],
        # [100, 100, 0],
    ]

    H, W, C = img.shape
    stickwidth = max(int(min(H, W) / stick_width_norm), 1)

    for _idx, ((k1_index, k2_index), color) in enumerate(zip(limbSeq, colors)):
        keypoint1 = kp2ds_body[k1_index - 1]
        keypoint2 = kp2ds_body[k2_index - 1]

        if keypoint1[-1] < threshold or keypoint2[-1] < threshold:
            continue

        Y = np.array([keypoint1[0], keypoint2[0]])
        X = np.array([keypoint1[1], keypoint2[1]])
        mX = np.mean(X)
        mY = np.mean(Y)
        length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
        angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
        polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
        cv2.fillConvexPoly(img, polygon, [int(float(c) * 0.6) for c in color])

    for _idx, (keypoint, color) in enumerate(zip(kp2ds_body, colors)):
        if keypoint[-1] < threshold:
            continue
        x, y = keypoint[0], keypoint[1]
        # cv2.circle(canvas, (int(x), int(y)), 4, color, thickness=-1)
        cv2.circle(img, (int(x), int(y)), stickwidth, color, thickness=-1)

    if draw_hand:
        img = draw_handpose(img, kp2ds_lhand, hand_score_th=threshold)
        img = draw_handpose(img, kp2ds_rhand, hand_score_th=threshold)

    kp2ds_body[:, 0] /= W
    kp2ds_body[:, 1] /= H

    if data_to_json is not None:
        if idx == -1:
            data_to_json.append(
                {
                    "image_id": "frame_{:05d}.jpg".format(len(data_to_json) + 1),
                    "height": H,
                    "width": W,
                    "category_id": 1,
                    "keypoints_body": kp2ds_body.tolist(),
                    "keypoints_left_hand": kp2ds_lhand.tolist(),
                    "keypoints_right_hand": kp2ds_rhand.tolist(),
                }
            )
        else:
            data_to_json[idx] = {
                "image_id": "frame_{:05d}.jpg".format(idx + 1),
                "height": H,
                "width": W,
                "category_id": 1,
                "keypoints_body": kp2ds_body.tolist(),
                "keypoints_left_hand": kp2ds_lhand.tolist(),
                "keypoints_right_hand": kp2ds_rhand.tolist(),
            }
    return img


def draw_nose(
    img,
    kp2ds,
    threshold=0.6,
    data_to_json=None,
    idx=-1,
    kp2ds_lhand=None,
    kp2ds_rhand=None,
    draw_hand=False,
    stick_width_norm=200,
):
    """
    Draw keypoints and connections representing hand pose on a given canvas.

    Args:
        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the hand pose.
        keypoints (List[Keypoint]| None): A list of Keypoint objects representing the hand keypoints to be drawn
                                          or None if no keypoints are present.

    Returns:
        np.ndarray: A 3D numpy array representing the modified canvas with the drawn hand pose.

    Note:
        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
    """

    new_kep_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",
    ]
    # kp2ds_body = (kp2ds.copy()[[0, 6, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 17, 20]] + \
    #              kp2ds.copy()[[0, 5, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 18, 21]]) / 2
    kp2ds = kp2ds.copy()
    kp2ds[1:, 2] = 0
    # kp2ds[0, 2] = 1
    kp2ds_body = kp2ds
    # kp2ds_body = kp2ds_body[:18]

    # kp2ds_lhand = kp2ds.copy()[91:112]
    # kp2ds_rhand = kp2ds.copy()[112:133]

    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],
        # [11, 20],  # foot
    ]

    colors = [
        # [255, 0, 0],
        # [255, 85, 0],
        # [255, 170, 0],
        # [255, 255, 0],
        # [170, 255, 0],
        # [85, 255, 0],
        # [0, 255, 0],
        # [0, 255, 85],
        # [0, 255, 170],
        # [0, 255, 255],
        # [0, 170, 255],
        # [0, 85, 255],
        # [0, 0, 255],
        # [85, 0, 255],
        [170, 0, 255],
        # [255, 0, 255],
        # [255, 0, 170],
        # [255, 0, 85],
        # foot
        # [200, 200, 0],
        # [100, 100, 0],
    ]

    H, W, C = img.shape
    stickwidth = max(int(min(H, W) / stick_width_norm), 1)

    # for _idx, ((k1_index, k2_index), color) in enumerate(zip(limbSeq, colors)):
    #     keypoint1 = kp2ds_body[k1_index - 1]
    #     keypoint2 = kp2ds_body[k2_index - 1]

    #     if keypoint1[-1] < threshold or keypoint2[-1] < threshold:
    #         continue

    #     Y = np.array([keypoint1[0], keypoint2[0]])
    #     X = np.array([keypoint1[1], keypoint2[1]])
    #     mX = np.mean(X)
    #     mY = np.mean(Y)
    #     length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
    #     angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
    #     polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
    #     cv2.fillConvexPoly(img, polygon, [int(float(c) * 0.6) for c in color])

    for _idx, (keypoint, color) in enumerate(zip(kp2ds_body, colors)):
        if keypoint[-1] < threshold:
            continue
        x, y = keypoint[0], keypoint[1]
        # cv2.circle(canvas, (int(x), int(y)), 4, color, thickness=-1)
        cv2.circle(img, (int(x), int(y)), stickwidth, color, thickness=-1)

    if draw_hand:
        img = draw_handpose(img, kp2ds_lhand, hand_score_th=threshold)
        img = draw_handpose(img, kp2ds_rhand, hand_score_th=threshold)

    kp2ds_body[:, 0] /= W
    kp2ds_body[:, 1] /= H

    if data_to_json is not None:
        if idx == -1:
            data_to_json.append(
                {
                    "image_id": "frame_{:05d}.jpg".format(len(data_to_json) + 1),
                    "height": H,
                    "width": W,
                    "category_id": 1,
                    "keypoints_body": kp2ds_body.tolist(),
                    "keypoints_left_hand": kp2ds_lhand.tolist(),
                    "keypoints_right_hand": kp2ds_rhand.tolist(),
                }
            )
        else:
            data_to_json[idx] = {
                "image_id": "frame_{:05d}.jpg".format(idx + 1),
                "height": H,
                "width": W,
                "category_id": 1,
                "keypoints_body": kp2ds_body.tolist(),
                "keypoints_left_hand": kp2ds_lhand.tolist(),
                "keypoints_right_hand": kp2ds_rhand.tolist(),
            }
    return img


def draw_aapose(
    img,
    kp2ds,
    threshold=0.6,
    data_to_json=None,
    idx=-1,
    kp2ds_lhand=None,
    kp2ds_rhand=None,
    draw_hand=False,
    stick_width_norm=200,
    draw_head=True
):
    """
    Draw keypoints and connections representing hand pose on a given canvas.

    Args:
        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the hand pose.
        keypoints (List[Keypoint]| None): A list of Keypoint objects representing the hand keypoints to be drawn
                                          or None if no keypoints are present.

    Returns:
        np.ndarray: A 3D numpy array representing the modified canvas with the drawn hand pose.

    Note:
        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
    """

    new_kep_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",
    ]
    # kp2ds_body = (kp2ds.copy()[[0, 6, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 17, 20]] + \
    #              kp2ds.copy()[[0, 5, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 18, 21]]) / 2
    kp2ds = kp2ds.copy()
    if not draw_head:
        kp2ds[[0,14,15,16,17], 2] = 0
    kp2ds_body = kp2ds

    # kp2ds_lhand = kp2ds.copy()[91:112]
    # kp2ds_rhand = kp2ds.copy()[112:133]

    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],
        [11, 20],  # foot
    ]

    colors = [
        [255, 0, 0],
        [255, 85, 0],
        [255, 170, 0],
        [255, 255, 0],
        [170, 255, 0],
        [85, 255, 0],
        [0, 255, 0],
        [0, 255, 85],
        [0, 255, 170],
        [0, 255, 255],
        [0, 170, 255],
        [0, 85, 255],
        [0, 0, 255],
        [85, 0, 255],
        [170, 0, 255],
        [255, 0, 255],
        [255, 0, 170],
        [255, 0, 85],
        # foot
        [200, 200, 0],
        [100, 100, 0],
    ]

    H, W, C = img.shape
    stickwidth = max(int(min(H, W) / stick_width_norm), 1)

    for _idx, ((k1_index, k2_index), color) in enumerate(zip(limbSeq, colors)):
        keypoint1 = kp2ds_body[k1_index - 1]
        keypoint2 = kp2ds_body[k2_index - 1]

        if keypoint1[-1] < threshold or keypoint2[-1] < threshold:
            continue

        Y = np.array([keypoint1[0], keypoint2[0]])
        X = np.array([keypoint1[1], keypoint2[1]])
        mX = np.mean(X)
        mY = np.mean(Y)
        length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
        angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
        polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
        cv2.fillConvexPoly(img, polygon, [int(float(c) * 0.6) for c in color])

    for _idx, (keypoint, color) in enumerate(zip(kp2ds_body, colors)):
        if keypoint[-1] < threshold:
            continue
        x, y = keypoint[0], keypoint[1]
        # cv2.circle(canvas, (int(x), int(y)), 4, color, thickness=-1)
        cv2.circle(img, (int(x), int(y)), stickwidth, color, thickness=-1)

    if draw_hand:
        img = draw_handpose(img, kp2ds_lhand, hand_score_th=threshold)
        img = draw_handpose(img, kp2ds_rhand, hand_score_th=threshold)

    kp2ds_body[:, 0] /= W
    kp2ds_body[:, 1] /= H

    if data_to_json is not None:
        if idx == -1:
            data_to_json.append(
                {
                    "image_id": "frame_{:05d}.jpg".format(len(data_to_json) + 1),
                    "height": H,
                    "width": W,
                    "category_id": 1,
                    "keypoints_body": kp2ds_body.tolist(),
                    "keypoints_left_hand": kp2ds_lhand.tolist(),
                    "keypoints_right_hand": kp2ds_rhand.tolist(),
                }
            )
        else:
            data_to_json[idx] = {
                "image_id": "frame_{:05d}.jpg".format(idx + 1),
                "height": H,
                "width": W,
                "category_id": 1,
                "keypoints_body": kp2ds_body.tolist(),
                "keypoints_left_hand": kp2ds_lhand.tolist(),
                "keypoints_right_hand": kp2ds_rhand.tolist(),
            }
    return img


def draw_aapose_new(
    img,
    kp2ds,
    threshold=0.6,
    data_to_json=None,
    idx=-1,
    kp2ds_lhand=None,
    kp2ds_rhand=None,
    draw_hand=False,
    stickwidth_type='v2',
    draw_head=True
):
    """
    Draw keypoints and connections representing hand pose on a given canvas.

    Args:
        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the hand pose.
        keypoints (List[Keypoint]| None): A list of Keypoint objects representing the hand keypoints to be drawn
                                          or None if no keypoints are present.

    Returns:
        np.ndarray: A 3D numpy array representing the modified canvas with the drawn hand pose.

    Note:
        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
    """

    new_kep_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",
    ]
    # kp2ds_body = (kp2ds.copy()[[0, 6, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 17, 20]] + \
    #              kp2ds.copy()[[0, 5, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 18, 21]]) / 2
    kp2ds = kp2ds.copy()
    if not draw_head:
        kp2ds[[0,14,15,16,17], 2] = 0
    kp2ds_body = kp2ds

    # kp2ds_lhand = kp2ds.copy()[91:112]
    # kp2ds_rhand = kp2ds.copy()[112:133]

    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],
        [11, 20],  # foot
    ]

    colors = [
        [255, 0, 0],
        [255, 85, 0],
        [255, 170, 0],
        [255, 255, 0],
        [170, 255, 0],
        [85, 255, 0],
        [0, 255, 0],
        [0, 255, 85],
        [0, 255, 170],
        [0, 255, 255],
        [0, 170, 255],
        [0, 85, 255],
        [0, 0, 255],
        [85, 0, 255],
        [170, 0, 255],
        [255, 0, 255],
        [255, 0, 170],
        [255, 0, 85],
        # foot
        [200, 200, 0],
        [100, 100, 0],
    ]

    H, W, C = img.shape
    H, W, C = img.shape

    if stickwidth_type == 'v1':
        stickwidth = max(int(min(H, W) / 200), 1)
    elif stickwidth_type == 'v2':
        stickwidth = max(int(min(H, W) / 200) - 1, 1)
    else:
        raise

    for _idx, ((k1_index, k2_index), color) in enumerate(zip(limbSeq, colors)):
        keypoint1 = kp2ds_body[k1_index - 1]
        keypoint2 = kp2ds_body[k2_index - 1]

        if keypoint1[-1] < threshold or keypoint2[-1] < threshold:
            continue

        Y = np.array([keypoint1[0], keypoint2[0]])
        X = np.array([keypoint1[1], keypoint2[1]])
        mX = np.mean(X)
        mY = np.mean(Y)
        length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
        angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
        polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
        cv2.fillConvexPoly(img, polygon, [int(float(c) * 0.6) for c in color])

    for _idx, (keypoint, color) in enumerate(zip(kp2ds_body, colors)):
        if keypoint[-1] < threshold:
            continue
        x, y = keypoint[0], keypoint[1]
        # cv2.circle(canvas, (int(x), int(y)), 4, color, thickness=-1)
        cv2.circle(img, (int(x), int(y)), stickwidth, color, thickness=-1)

    if draw_hand:
        img = draw_handpose_new(img, kp2ds_lhand, stickwidth_type=stickwidth_type, hand_score_th=threshold)
        img = draw_handpose_new(img, kp2ds_rhand, stickwidth_type=stickwidth_type, hand_score_th=threshold)

    kp2ds_body[:, 0] /= W
    kp2ds_body[:, 1] /= H

    if data_to_json is not None:
        if idx == -1:
            data_to_json.append(
                {
                    "image_id": "frame_{:05d}.jpg".format(len(data_to_json) + 1),
                    "height": H,
                    "width": W,
                    "category_id": 1,
                    "keypoints_body": kp2ds_body.tolist(),
                    "keypoints_left_hand": kp2ds_lhand.tolist(),
                    "keypoints_right_hand": kp2ds_rhand.tolist(),
                }
            )
        else:
            data_to_json[idx] = {
                "image_id": "frame_{:05d}.jpg".format(idx + 1),
                "height": H,
                "width": W,
                "category_id": 1,
                "keypoints_body": kp2ds_body.tolist(),
                "keypoints_left_hand": kp2ds_lhand.tolist(),
                "keypoints_right_hand": kp2ds_rhand.tolist(),
            }
    return img


def draw_bbox(img, bbox, color=(255, 0, 0)):
    img = load_image(img)
    bbox = [int(bbox_tmp) for bbox_tmp in bbox]
    cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
    return img


def draw_kp2ds(img, kp2ds, threshold=0, color=(255, 0, 0), skeleton=None, reverse=False):
    img = load_image(img, reverse)

    if skeleton is not None:
        if skeleton == "coco17":
            skeleton_list = [
                [6, 8],
                [8, 10],
                [5, 7],
                [7, 9],
                [11, 13],
                [13, 15],
                [12, 14],
                [14, 16],
                [5, 6],
                [6, 12],
                [12, 11],
                [11, 5],
            ]
            color_list = [
                (255, 0, 0),
                (0, 255, 0),
                (0, 0, 255),
                (255, 255, 0),
                (255, 0, 255),
                (0, 255, 255),
            ]
        elif skeleton == "cocowholebody":
            skeleton_list = [
                [6, 8],
                [8, 10],
                [5, 7],
                [7, 9],
                [11, 13],
                [13, 15],
                [12, 14],
                [14, 16],
                [5, 6],
                [6, 12],
                [12, 11],
                [11, 5],
                [15, 17],
                [15, 18],
                [15, 19],
                [16, 20],
                [16, 21],
                [16, 22],
                [91, 92, 93, 94, 95],
                [91, 96, 97, 98, 99],
                [91, 100, 101, 102, 103],
                [91, 104, 105, 106, 107],
                [91, 108, 109, 110, 111],
                [112, 113, 114, 115, 116],
                [112, 117, 118, 119, 120],
                [112, 121, 122, 123, 124],
                [112, 125, 126, 127, 128],
                [112, 129, 130, 131, 132],
            ]
            color_list = [
                (255, 0, 0),
                (0, 255, 0),
                (0, 0, 255),
                (255, 255, 0),
                (255, 0, 255),
                (0, 255, 255),
            ]
        else:
            color_list = [color]
        for _idx, _skeleton in enumerate(skeleton_list):
            for i in range(len(_skeleton) - 1):
                cv2.line(
                    img,
                    (int(kp2ds[_skeleton[i], 0]), int(kp2ds[_skeleton[i], 1])),
                    (int(kp2ds[_skeleton[i + 1], 0]), int(kp2ds[_skeleton[i + 1], 1])),
                    color_list[_idx % len(color_list)],
                    3,
                )

    for _idx, kp2d in enumerate(kp2ds):
        if kp2d[2] > threshold:
            cv2.circle(img, (int(kp2d[0]), int(kp2d[1])), 3, color, -1)
            # cv2.putText(img,
            #         str(_idx),
            #         (int(kp2d[0, i, 0])*1,
            #             int(kp2d[0, i, 1])*1),
            #         cv2.FONT_HERSHEY_SIMPLEX,
            #         0.75,
            #         color,
            #         2
            #         )

    return img


def draw_mask(img, mask, background=0, return_rgba=False):
    img = load_image(img)
    h, w, _ = img.shape
    if type(background) == int:
        background = np.ones((h, w, 3)).astype(np.uint8) * 255 * background
    backgournd = cv2.resize(background, (w, h))
    img_rgba = np.concatenate([img, mask], -1)
    return alphaMerge(img_rgba, background, 0, 0, return_rgba=True)


def draw_pcd(pcd_list, save_path=None):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")

    color_list = ["r", "g", "b", "y", "p"]

    for _idx, _pcd in enumerate(pcd_list):
        ax.scatter(_pcd[:, 0], _pcd[:, 1], _pcd[:, 2], c=color_list[_idx], marker="o")

    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")

    if save_path is not None:
        plt.savefig(save_path)
    else:
        plt.savefig("tmp.png")


def load_image(img, reverse=False):
    if type(img) == str:
        img = cv2.imread(img)
    if reverse:
        img = img.astype(np.float32)
        img = img[:, :, ::-1]
        img = img.astype(np.uint8)
    return img


def draw_skeleten(meta):
    kps = []
    for i, kp in enumerate(meta["keypoints_body"]):
        if kp is None:
            # if kp is None:
            kps.append([0, 0, 0])
        else:
            kps.append([*kp, 1])
    kps = np.array(kps)

    kps[:, 0] *= meta["width"]
    kps[:, 1] *= meta["height"]
    pose_img = np.zeros([meta["height"], meta["width"], 3], dtype=np.uint8)

    pose_img = draw_aapose(
        pose_img,
        kps,
        draw_hand=True,
        kp2ds_lhand=meta["keypoints_left_hand"],
        kp2ds_rhand=meta["keypoints_right_hand"],
    )
    return pose_img


def draw_skeleten_with_pncc(pncc: np.ndarray, meta: Dict) -> np.ndarray:
    """
    Args:
        pncc: [H,W,3]
        meta: required keys: keypoints_body: [N, 3] keypoints_left_hand, keypoints_right_hand
    Return:
        np.ndarray [H, W, 3]
    """
    # preprocess keypoints
    kps = []
    for i, kp in enumerate(meta["keypoints_body"]):
        if kp is None:
            # if kp is None:
            kps.append([0, 0, 0])
        elif i in [14, 15, 16, 17]:
            kps.append([0, 0, 0])
        else:
            kps.append([*kp])
    kps = np.stack(kps)

    kps[:, 0] *= pncc.shape[1]
    kps[:, 1] *= pncc.shape[0]

    # draw neck
    canvas = np.zeros_like(pncc)
    if kps[0][2] > 0.6 and kps[1][2] > 0.6:
        canvas = draw_ellipse_by_2kp(canvas, kps[0], kps[1], [0, 0, 255])

    # draw pncc
    mask = (pncc > 0).max(axis=2)
    canvas[mask] = pncc[mask]
    pncc = canvas

    # draw other skeleten
    kps[0] = 0

    meta["keypoints_left_hand"][:, 0] *= meta["width"]
    meta["keypoints_left_hand"][:, 1] *= meta["height"]

    meta["keypoints_right_hand"][:, 0] *= meta["width"]
    meta["keypoints_right_hand"][:, 1] *= meta["height"]
    pose_img = draw_aapose(
        pncc,
        kps,
        draw_hand=True,
        kp2ds_lhand=meta["keypoints_left_hand"],
        kp2ds_rhand=meta["keypoints_right_hand"],
    )
    return pose_img


FACE_CUSTOM_STYLE = {
    "eyeball": {"indexs": [68, 69], "color": [255, 255, 255], "connect": False},
    "left_eyebrow": {"indexs": [17, 18, 19, 20, 21], "color": [0, 255, 0]},
    "right_eyebrow": {"indexs": [22, 23, 24, 25, 26], "color": [0, 0, 255]},
    "left_eye": {"indexs": [36, 37, 38, 39, 40, 41], "color": [255, 255, 0], "close": True},
    "right_eye": {"indexs": [42, 43, 44, 45, 46, 47], "color": [255, 0, 255], "close": True},
    "mouth_outside": {"indexs": list(range(48, 60)), "color": [100, 255, 50], "close": True},
    "mouth_inside": {"indexs": [60, 61, 62, 63, 64, 65, 66, 67], "color": [255, 100, 50], "close": True},
}


def draw_face_kp(img, kps, thickness=2, style=FACE_CUSTOM_STYLE):
    """
    Args:
        img: [H, W, 3]
        kps: [70, 2]
    """
    img = img.copy()
    for key, item in style.items():
        pts = np.array(kps[item["indexs"]]).astype(np.int32)
        connect = item.get("connect", True)
        color = item["color"]
        close = item.get("close", False)
        if connect:
            cv2.polylines(img, [pts], close, color, thickness=thickness)
        else:
            for kp in pts:
                kp = np.array(kp).astype(np.int32)
                cv2.circle(img, kp, thickness * 2, color=color, thickness=-1)
    return img


def draw_traj(metas: List[AAPoseMeta], threshold=0.6):

    colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
                [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
                [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85], [100, 255, 50], [255, 100, 50], 
                # foot
                [200, 200, 0],
                [100, 100, 0]
                ]
    limbSeq = [
                    [1, 2], [1, 5],     # shoulders
                    [2, 3], [3, 4],     # left arm
                    [5, 6], [6, 7],     # right arm
                    [1, 8], [8, 9], [9, 10],    # right leg 
                    [1, 11], [11, 12], [12, 13],  # left leg
                     # face (nose, eyes, ears)
                    [13, 18], [10, 19] # foot
                ]
    
    face_seq = [[1, 0], [0, 14], [14, 16], [0, 15], [15, 17]] 
    kp_body = np.array([meta.kps_body for meta in metas])
    kp_body_p = np.array([meta.kps_body_p for meta in metas])
    
    
    face_seq = random.sample(face_seq, 2)

    kp_lh = np.array([meta.kps_lhand for meta in metas])
    kp_rh = np.array([meta.kps_rhand for meta in metas])

    kp_lh_p = np.array([meta.kps_lhand_p for meta in metas])
    kp_rh_p = np.array([meta.kps_rhand_p for meta in metas])

    # kp_lh = np.concatenate([kp_lh, kp_lh_p], axis=-1)
    # kp_rh = np.concatenate([kp_rh, kp_rh_p], axis=-1)
    
    new_limbSeq = []
    key_point_list = []
    for _idx, ((k1_index, k2_index)) in enumerate(limbSeq):
        
        vis = (kp_body_p[:, k1_index] > threshold) * (kp_body_p[:, k2_index] > threshold) * 1
        if vis.sum() * 1.0 / vis.shape[0] > 0.4:
            new_limbSeq.append([k1_index, k2_index])

    for _idx, ((k1_index, k2_index)) in enumerate(limbSeq):

        keypoint1 = kp_body[:, k1_index - 1]
        keypoint2 = kp_body[:, k2_index - 1]
        interleave = random.randint(4, 7)
        randind = random.randint(0, interleave - 1)
        # randind = random.rand(range(interleave), sampling_num)

        Y = np.array([keypoint1[:, 0], keypoint2[:, 0]]) 
        X = np.array([keypoint1[:, 1], keypoint2[:, 1]])

        vis = (keypoint1[:, -1] > threshold) * (keypoint2[:, -1] > threshold) * 1

        # for randidx in randind:
        t = randind / interleave
        x = (1-t)*Y[0, :] + t*Y[1, :]
        y = (1-t)*X[0, :] + t*X[1, :]

        # np.array([1])
        x = x.astype(int)
        y = y.astype(int)

        new_array = np.array([x, y, vis]).T
        
        key_point_list.append(new_array)
    
    indx_lh = random.randint(0, kp_lh.shape[1] - 1)
    lh = kp_lh[:, indx_lh, :]
    lh_p = kp_lh_p[:, indx_lh:indx_lh+1]
    lh = np.concatenate([lh, lh_p], axis=-1)
    
    indx_rh = random.randint(0, kp_rh.shape[1] - 1)
    rh = kp_rh[:, random.randint(0, kp_rh.shape[1] - 1), :]
    rh_p = kp_rh_p[:, indx_rh:indx_rh+1]
    rh = np.concatenate([rh, rh_p], axis=-1)



    lh[-1, :] = (lh[-1, :] > threshold) * 1
    rh[-1, :] = (rh[-1, :] > threshold) * 1

    # print(rh.shape, new_array.shape)
    # exit()
    key_point_list.append(lh.astype(int))
    key_point_list.append(rh.astype(int))

    
    key_points_list = np.stack(key_point_list)
    num_points = len(key_points_list)
    sample_colors = random.sample(colors, num_points)

    stickwidth = max(int(min(metas[0].width, metas[0].height) / 150), 2)

    image_list_ori = []
    for i in range(key_points_list.shape[-2]):
        _image_vis = np.zeros((metas[0].width, metas[0].height, 3))
        points = key_points_list[:, i, :]
        for idx, point in enumerate(points):
            x, y, vis = point
            if vis == 1:
                cv2.circle(_image_vis, (x, y), stickwidth, sample_colors[idx], thickness=-1)
        
        image_list_ori.append(_image_vis)
    
    return image_list_ori

    return [np.zeros([meta.width, meta.height, 3], dtype=np.uint8) for meta in metas]


if __name__ == "__main__":
    meta = {
        "image_id": "00472.jpg",
        "height": 540,
        "width": 414,
        "category_id": 1,
        "keypoints_body": [
            [0.5084776947463768, 0.11350188078703703],
            [0.504467655495169, 0.20419560185185184],
            [0.3982016153381642, 0.198046875],
            [0.3841664779589372, 0.34869068287037036],
            [0.3901815368357488, 0.4670536747685185],
            [0.610733695652174, 0.2103443287037037],
            [0.6167487545289855, 0.3517650462962963],
            [0.6448190292874396, 0.4762767650462963],
            [0.4523371452294686, 0.47320240162037036],
            [0.4503321256038647, 0.6776475694444445],
            [0.47639738073671495, 0.8544234664351852],
            [0.5766483620169082, 0.47320240162037036],
            [0.5666232638888888, 0.6761103877314815],
            [0.534542949879227, 0.863646556712963],
            [0.4864224788647343, 0.09505570023148148],
            [0.5285278910024155, 0.09351851851851851],
            [0.46236224335748793, 0.10581597222222222],
            [0.5586031853864735, 0.10274160879629629],
            [0.4994551064311594, 0.9405056423611111],
            [0.4152442821557971, 0.9312825520833333],
        ],
        "keypoints_left_hand": [
            [267.78515625, 263.830078125, 1.2840936183929443],
            [265.294921875, 269.640625, 1.2546794414520264],
            [263.634765625, 277.111328125, 1.2863062620162964],
            [262.8046875, 285.412109375, 1.267038345336914],
            [261.14453125, 292.8828125, 1.280144453048706],
            [273.595703125, 281.26171875, 1.2592815160751343],
            [271.10546875, 291.22265625, 1.3256099224090576],
            [265.294921875, 294.54296875, 1.2368024587631226],
            [261.14453125, 294.54296875, 0.9771889448165894],
            [274.42578125, 282.091796875, 1.250044584274292],
            [269.4453125, 291.22265625, 1.2571144104003906],
            [264.46484375, 292.8828125, 1.177802324295044],
            [260.314453125, 292.052734375, 0.9283463358879089],
            [273.595703125, 282.091796875, 1.1834490299224854],
            [269.4453125, 290.392578125, 1.188171625137329],
            [265.294921875, 290.392578125, 1.192609429359436],
            [261.974609375, 289.5625, 0.9366656541824341],
            [271.935546875, 281.26171875, 1.0946396589279175],
            [268.615234375, 287.072265625, 0.9906131029129028],
            [265.294921875, 287.90234375, 1.0219476222991943],
            [262.8046875, 287.072265625, 0.9240120053291321],
        ],
        "keypoints_right_hand": [
            [161.53515625, 258.849609375, 1.2069408893585205],
            [168.17578125, 263.0, 1.1846840381622314],
            [173.986328125, 269.640625, 1.1435924768447876],
            [173.986328125, 277.94140625, 1.1802611351013184],
            [173.986328125, 286.2421875, 1.2599592208862305],
            [165.685546875, 275.451171875, 1.0633569955825806],
            [167.345703125, 286.2421875, 1.1693341732025146],
            [169.8359375, 291.22265625, 1.2698509693145752],
            [170.666015625, 294.54296875, 1.0619274377822876],
            [160.705078125, 276.28125, 1.0995020866394043],
            [163.1953125, 287.90234375, 1.2735884189605713],
            [166.515625, 291.22265625, 1.339503526687622],
            [169.005859375, 294.54296875, 1.0835273265838623],
            [157.384765625, 277.111328125, 1.0866981744766235],
            [161.53515625, 287.072265625, 1.2468621730804443],
            [164.025390625, 289.5625, 1.2817761898040771],
            [166.515625, 292.052734375, 1.099466323852539],
            [155.724609375, 277.111328125, 1.1065717935562134],
            [159.044921875, 285.412109375, 1.1924479007720947],
            [160.705078125, 287.072265625, 1.1304771900177002],
            [162.365234375, 287.90234375, 1.0040509700775146],
        ],
    }
    demo_meta = AAPoseMeta(meta)
    res = draw_traj([demo_meta]*5)
    cv2.imwrite("traj.png", res[0][..., ::-1])