src/models/XPose/util/keypoint_ops.py

import torch, os

def keypoint_xyxyzz_to_xyzxyz(keypoints: torch.Tensor):
    """_summary_

    Args:
        keypoints (torch.Tensor): ..., 51
    """
    res = torch.zeros_like(keypoints)
    num_points = keypoints.shape[-1] // 3
    Z = keypoints[..., :2*num_points]
    V = keypoints[..., 2*num_points:]
    res[...,0::3] = Z[..., 0::2]
    res[...,1::3] = Z[..., 1::2]
    res[...,2::3] = V[...]
    return res

def keypoint_xyzxyz_to_xyxyzz(keypoints: torch.Tensor):
    """_summary_

    Args:
        keypoints (torch.Tensor): ..., 51
    """
    res = torch.zeros_like(keypoints)
    num_points = keypoints.shape[-1] // 3
    res[...,0:2*num_points:2] = keypoints[..., 0::3]
    res[...,1:2*num_points:2] = keypoints[..., 1::3]
    res[...,2*num_points:] = keypoints[..., 2::3]
    return res