Merge pull request #15 from cleong110/pose_utils

Pose utils and some tests.

Author: AmitMY

.gitignore

@@ -1,4 +1,10 @@
 .idea/
 build/
 pose_evaluation.egg-info/
-**/__pycache__/
+**/__pycache__/
+.coverage
+.vscode/
+coverage.lcov
+**/test_data/
+*.npz
+*.code-workspace

pose_evaluation/metrics/.gitignore

@@ -1 +1 @@
-temp/
+tests

pose_evaluation/utils/conftest.py

@@ -0,0 +1,51 @@
+import json
+import copy
+from pathlib import Path
+from typing import List, Dict
+
+import pytest
+from pose_format import Pose
+from pose_format.utils.generic import fake_pose
+from pose_format.utils.openpose_135 import (
+    OpenPose_Components as openpose_135_components,
+)
+
+from pose_evaluation.utils.pose_utils import load_pose_file
+
+
+
+utils_test_data_dir = Path(__file__).parent / "test" / "test_data"
+
+
+@pytest.fixture(scope="function")
+def mediapipe_poses_test_data_paths() -> List[Path]:
+    pose_file_paths = list(utils_test_data_dir.glob("*.pose"))
+    return pose_file_paths
+
+
+@pytest.fixture(scope="function")
+def mediapipe_poses_test_data(mediapipe_poses_test_data_paths) -> List[Pose]:
+    original_poses = [
+        load_pose_file(pose_path) for pose_path in mediapipe_poses_test_data_paths
+    ]
+    # I ran into issues where if one test would modify a Pose, it would affect other tests.
+    # specifically, pose.header.components[0].name = unsupported_component_name in test_detect_format
+    # this ensures we get a fresh object each time.
+    return copy.deepcopy(original_poses)
+
+
+@pytest.fixture
+def standard_mediapipe_components_dict() -> Dict[str, List[str]]:
+    format_json = utils_test_data_dir / "mediapipe_components_and_points.json"
+    with open(format_json, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+@pytest.fixture
+def fake_openpose_poses(count: int = 3) -> List[Pose]:
+    return [fake_pose(30) for _ in range(count)]
+
+
+@pytest.fixture
+def fake_openpose_135_poses(count: int = 3) -> List[Pose]:
+    return [fake_pose(30, components=openpose_135_components) for _ in range(count)]

pose_evaluation/utils/pose_utils.py

@@ -0,0 +1,92 @@
+from pathlib import Path
+from typing import List, Tuple, Dict, Iterable
+from collections import defaultdict
+import numpy as np
+from numpy import ma
+from pose_format import Pose
+
+
+def pose_remove_world_landmarks(pose: Pose) -> Pose:
+    return pose.remove_components(["POSE_WORLD_LANDMARKS"])
+
+
+def get_component_names_and_points_dict(
+    pose: Pose,
+) -> Tuple[List[str], Dict[str, List[str]]]:
+    component_names = []
+    points_dict = defaultdict(list)
+    for component in pose.header.components:
+        component_names.append(component.name)
+
+        for point in component.points:
+            points_dict[component.name].append(point)
+
+    return component_names, points_dict
+
+
+def get_face_and_hands_from_pose(pose: Pose) -> Pose:
+    # based on MediaPipe Holistic format.
+    components_to_keep = [
+        "FACE_LANDMARKS",
+        "LEFT_HAND_LANDMARKS",
+        "RIGHT_HAND_LANDMARKS",
+    ]
+    return pose.get_components(components_to_keep)
+
+
+def load_pose_file(pose_path: Path) -> Pose:
+    pose_path = Path(pose_path).resolve()
+    with pose_path.open("rb") as f:
+        pose = Pose.read(f.read())
+    return pose
+
+
+def reduce_poses_to_intersection(
+    poses: Iterable[Pose],
+) -> List[Pose]:
+    poses = list(poses) # get a list, no need to copy
+
+    # look at the first pose
+    component_names = {c.name for c in poses[0].header.components}
+    points = {c.name: set(c.points) for c in poses[0].header.components}
+
+    # remove anything that other poses don't have
+    for pose in poses[1:]:
+        component_names.intersection_update({c.name for c in pose.header.components})
+        for component in pose.header.components:
+            points[component.name].intersection_update(set(component.points))
+
+    # change datatypes to match get_components, then update the poses
+    points_dict = {}
+    for c_name in points.keys():
+        points_dict[c_name] = list(points[c_name])
+    poses = [pose.get_components(list(component_names), points_dict) for pose in poses]
+    return poses
+
+
+def zero_pad_shorter_poses(poses: Iterable[Pose]) -> List[Pose]:
+    poses = [pose.copy() for pose in poses]
+    # arrays = [pose.body.data for pose in poses]
+
+    # first dimension is frames. Then People, joint-points, XYZ or XY
+    max_frame_count = max(len(pose.body.data) for pose in poses)
+    # Pad the shorter array with zeros
+    for pose in poses:
+        if len(pose.body.data) < max_frame_count:
+            desired_shape = list(pose.body.data.shape)
+            desired_shape[0] = max_frame_count - len(pose.body.data)
+            padding_tensor = ma.zeros(desired_shape)
+            padding_tensor_conf = ma.ones(desired_shape[:-1])
+            pose.body.data = ma.concatenate([pose.body.data, padding_tensor], axis=0)
+            pose.body.confidence = ma.concatenate(
+                [pose.body.confidence, padding_tensor_conf]
+            )
+    return poses
+
+
+def pose_hide_low_conf(pose: Pose, confidence_threshold: float = 0.2) -> None:
+    mask = pose.body.confidence <= confidence_threshold
+    pose.body.confidence[mask] = 0
+    stacked_confidence = np.stack([mask, mask, mask], axis=3)
+    masked_data = ma.masked_array(pose.body.data, mask=stacked_confidence)
+    pose.body.data = masked_data