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