Source code for pose_format.utils.pose_converter

from typing import List

import numpy as np

from pose_format import Pose, PoseHeader
from pose_format.numpy import NumPyPoseBody
from pose_format.pose_header import PoseHeaderComponent
from pose_format.pose_visualizer import PoseVisualizer
from pose_format.utils.holistic import HAND_POINTS as HOLISTIC_HAND_POINTS
from pose_format.utils.holistic import holistic_components
from pose_format.utils.openpose import HAND_POINTS as OPENPOSE_HAND_POINTS

LEFT_HAND_MAP = [(("hand_left_keypoints_2d", k1), ("LEFT_HAND_LANDMARKS", k2))
                 for k1, k2 in zip(OPENPOSE_HAND_POINTS, HOLISTIC_HAND_POINTS)]
RIGHT_HAND_MAP = [(("hand_right_keypoints_2d", k1), ("RIGHT_HAND_LANDMARKS", k2))
                  for k1, k2 in zip(OPENPOSE_HAND_POINTS, HOLISTIC_HAND_POINTS)]
BODY_MAP = [
    {("pose_keypoints_2d", "Nose"), ("POSE_LANDMARKS", "NOSE")},
    {("pose_keypoints_2d", "Neck"), ("POSE_LANDMARKS", ("RIGHT_SHOULDER", "LEFT_SHOULDER"))},
    {("pose_keypoints_2d", "RShoulder"), ("POSE_LANDMARKS", "RIGHT_SHOULDER")},
    {("pose_keypoints_2d", "RElbow"), ("POSE_LANDMARKS", "RIGHT_ELBOW")},
    {("pose_keypoints_2d", "RWrist"), ("POSE_LANDMARKS", "RIGHT_WRIST")},
    {("pose_keypoints_2d", "LShoulder"), ("POSE_LANDMARKS", "LEFT_SHOULDER")},
    {("pose_keypoints_2d", "LElbow"), ("POSE_LANDMARKS", "LEFT_ELBOW")},
    {("pose_keypoints_2d", "LWrist"), ("POSE_LANDMARKS", "LEFT_WRIST")},
    {("pose_keypoints_2d", "MidHip"), ("POSE_LANDMARKS", ("RIGHT_HIP", "LEFT_HIP"))},
    {("pose_keypoints_2d", "RHip"), ("POSE_LANDMARKS", "RIGHT_HIP")},
    {("pose_keypoints_2d", "RKnee"), ("POSE_LANDMARKS", "RIGHT_KNEE")},
    {("pose_keypoints_2d", "RAnkle"), ("POSE_LANDMARKS", "RIGHT_ANKLE")},
    {("pose_keypoints_2d", "LHip"), ("POSE_LANDMARKS", "LEFT_HIP")},
    {("pose_keypoints_2d", "LKnee"), ("POSE_LANDMARKS", "LEFT_KNEE")},
    {("pose_keypoints_2d", "LAnkle"), ("POSE_LANDMARKS", "LEFT_ANKLE")},
    {("pose_keypoints_2d", "REye"), ("POSE_LANDMARKS", "RIGHT_EYE")},
    {("pose_keypoints_2d", "LEye"), ("POSE_LANDMARKS", "LEFT_EYE")},
    {("pose_keypoints_2d", "REar"), ("POSE_LANDMARKS", "RIGHT_EAR")},
    {("pose_keypoints_2d", "LEar"), ("POSE_LANDMARKS", "LEFT_EAR")},
    {("pose_keypoints_2d", "LHeel"), ("POSE_LANDMARKS", "LEFT_HEEL")},
    {("pose_keypoints_2d", "RHeel"), ("POSE_LANDMARKS", "RIGHT_HEEL")},
]

POSES_MAP = BODY_MAP + LEFT_HAND_MAP + RIGHT_HAND_MAP


[docs]def convert_pose(pose: Pose, pose_components: List[PoseHeaderComponent]) -> Pose: """ converts the given pose to a new pose instance based on given pose components. Parameters ---------- pose : Pose The initial pose object to convert pose_components : List[PoseHeaderComponent] the new set of pose components to define the pose structure Returns ------- Pose Converted pose object """ pose_header = PoseHeader(version=pose.header.version, dimensions=pose.header.dimensions, components=pose_components) base_shape = (pose.body.data.shape[0], pose.body.data.shape[1], pose_header.total_points()) data = np.zeros(shape=(*base_shape, len(pose_components[0].format) - 1), dtype=np.float32) conf = np.zeros(shape=base_shape, dtype=np.float32) original_components = set([c.name for c in pose.header.components]) new_components = set([c.name for c in pose_components]) # Create a mapping mapping = {} for points in POSES_MAP: original_point = None new_point = None for component, point in points: if component in original_components: original_point = (component, point) if component in new_components and isinstance(point, str): new_point = (component, point) if original_point is not None and new_point is not None: mapping[new_point] = original_point dims = min(len(pose_header.components[0].format), len(pose.header.components[0].format)) - 1 for (c1, p1), (c2, p2) in mapping.items(): p2 = tuple([p2]) if isinstance(p2, str) else p2 p2s = [pose.header._get_point_index(c2, p) for p in list(p2)] p1_index = pose_header._get_point_index(c1, p1) data[:, :, p1_index, :dims] = pose.body.data[:, :, p2s, :dims].mean(axis=2) conf[:, :, p1_index] = pose.body.confidence[:, :, p2s].mean(axis=2) pose_body = NumPyPoseBody(fps=pose.body.fps, data=data, confidence=conf) return Pose(pose_header, pose_body)
[docs]def save_image(pose: Pose, name: str): """ Saves visualized pose as an image with a given name Parameters ---------- pose : Pose Pose to be visualized and saved name : str Name to save image to. """ visualizer = PoseVisualizer(pose, thickness=1) frame = next(iter(visualizer.draw(background_color=(255, 255, 255)))) visualizer.save_frame(name, frame)
if __name__ == "__main__": with open("sample-data/video/sample.pose", "rb") as f: original_pose = Pose.read(f.read(), NumPyPoseBody) original_pose.focus() conv_holistic = convert_pose(original_pose, holistic_components()) conv_openpose1 = convert_pose(original_pose, original_pose.header.components) conv_openpose2 = convert_pose(conv_holistic, original_pose.header.components) save_image(original_pose, "original.png") save_image(conv_holistic, "conv_holistic.png") save_image(conv_openpose1, "conv_openpose1.png") save_image(conv_openpose2, "conv_openpose2.png")