diff --git a/README.md b/README.md index 3032454..e27bebe 100644 --- a/README.md +++ b/README.md @@ -252,6 +252,22 @@ directory = "/path/to/openpose_135/directory" pose = load_openpose_135_directory(directory, fps=24, width=1000, height=1000) ``` +##### Loading AlphaPose WholeBody JSON + +```python +from pose_format.utils.alphapose import load_alphapose_wholebody_from_json + +pose = load_alphapose_wholebody_from_json("alphapose.json") # 136-keypoint (default) +``` + +For the 133-keypoint variant: + +```python +from pose_format.utils.alphapose_133 import load_alphapose_wholebody_from_json as load_alphapose_133 + +pose = load_alphapose_133("alphapose.json") +``` + #### 7. Generating Fake Pose Data for Testing Purposes: ```python diff --git a/src/python/pose_format/utils/alphapose.py b/src/python/pose_format/utils/alphapose.py new file mode 100644 index 0000000..17d731e --- /dev/null +++ b/src/python/pose_format/utils/alphapose.py @@ -0,0 +1,258 @@ +import re +import json +import numpy as np +from ..numpy.pose_body import NumPyPoseBody +from ..pose import Pose +from ..pose_header import PoseHeader, PoseHeaderComponent, PoseHeaderDimensions + +# --- Shared constants (same across all AlphaPose WholeBody variants) --- + +FACE_POINTS = [f"face_{i}" for i in range(68)] +GENERAL_HAND_POINTS = [f"hand_{i}" for i in range(21)] + +# Left/right-specific point name lists — exported for reference / downstream use +LEFT_HAND_POINTS = [f"left_hand_{i}" for i in range(21)] +RIGHT_HAND_POINTS = [f"right_hand_{i}" for i in range(21)] + +# Single hand limb table using generic point names, shared by both hand components +HAND_LIMBS_NAMES = [ + ("hand_0", "hand_1"), ("hand_1", "hand_2"), ("hand_2", "hand_3"), ("hand_3", "hand_4"), + ("hand_0", "hand_5"), ("hand_5", "hand_6"), ("hand_6", "hand_7"), ("hand_7", "hand_8"), + ("hand_0", "hand_9"), ("hand_9", "hand_10"), ("hand_10", "hand_11"), ("hand_11", "hand_12"), + ("hand_0", "hand_13"), ("hand_13", "hand_14"), ("hand_14", "hand_15"), ("hand_15", "hand_16"), + ("hand_0", "hand_17"), ("hand_17", "hand_18"), ("hand_18", "hand_19"), ("hand_19", "hand_20"), +] + +FACE_LIMBS_NAMES = [ + ("face_0", "face_1"), ("face_1", "face_2"), ("face_2", "face_3"), ("face_3", "face_4"), + ("face_4", "face_5"), ("face_5", "face_6"), ("face_6", "face_7"), ("face_7", "face_8"), + ("face_8", "face_9"), ("face_9", "face_10"), ("face_10", "face_11"), ("face_11", "face_12"), + ("face_12", "face_13"), ("face_13", "face_14"), ("face_14", "face_15"), ("face_15", "face_16"), + ("face_17", "face_18"), ("face_18", "face_19"), ("face_19", "face_20"), ("face_20", "face_21"), + ("face_22", "face_23"), ("face_23", "face_24"), ("face_24", "face_25"), ("face_25", "face_26"), + ("face_27", "face_28"), ("face_28", "face_29"), ("face_29", "face_30"), + ("face_31", "face_32"), ("face_32", "face_33"), ("face_33", "face_34"), ("face_34", "face_35"), + ("face_36", "face_37"), ("face_37", "face_38"), ("face_38", "face_39"), + ("face_39", "face_40"), ("face_40", "face_41"), + ("face_42", "face_43"), ("face_43", "face_44"), ("face_44", "face_45"), + ("face_45", "face_46"), ("face_46", "face_47"), + ("face_48", "face_49"), ("face_49", "face_50"), ("face_50", "face_51"), ("face_51", "face_52"), + ("face_52", "face_53"), ("face_53", "face_54"), ("face_54", "face_55"), ("face_55", "face_56"), + ("face_56", "face_57"), ("face_57", "face_58"), ("face_58", "face_59"), ("face_59", "face_60"), + ("face_60", "face_61"), ("face_61", "face_62"), ("face_62", "face_63"), ("face_63", "face_64"), + ("face_64", "face_65"), ("face_65", "face_66"), ("face_66", "face_67"), +] + +# --- 136-keypoint body (WholeBody-136 has neck, head_top, pelvis) --- + +BODY_POINTS = [ + 'nose', 'left_eye', 'right_eye', 'left_ear', 'right_ear', + 'left_shoulder', 'right_shoulder', 'left_elbow', 'right_elbow', + 'left_wrist', 'right_wrist', 'left_hip', 'right_hip', + 'left_knee', 'right_knee', 'left_ankle', 'right_ankle', + 'head_top', 'neck', 'pelvis', + 'left_big_toe', 'right_big_toe', 'left_small_toe', 'right_small_toe', + 'left_heel', 'right_heel', +] + +BODY_LIMBS_NAMES = [ + ("nose", "left_eye"), ("nose", "right_eye"), + ("left_eye", "left_ear"), ("right_eye", "right_ear"), + ("head_top", "neck"), + ("left_shoulder", "neck"), ("right_shoulder", "neck"), + ("left_shoulder", "left_elbow"), ("left_elbow", "left_wrist"), + ("right_shoulder", "right_elbow"), ("right_elbow", "right_wrist"), + ("neck", "pelvis"), ("pelvis", "left_hip"), ("pelvis", "right_hip"), + ("left_hip", "left_knee"), ("left_knee", "left_ankle"), + ("right_hip", "right_knee"), ("right_knee", "right_ankle"), + ("left_ankle", "left_heel"), ("left_heel", "left_big_toe"), ("left_heel", "left_small_toe"), + ("right_ankle", "right_heel"), ("right_heel", "right_big_toe"), ("right_heel", "right_small_toe"), +] + + +def _map_limbs(points, limbs): + index_map = {name: idx for idx, name in enumerate(points)} + return [(index_map[a], index_map[b]) for (a, b) in limbs] + + +def get_alphapose_components(): + """ + Returns AlphaPose WholeBody-136 component definitions. + + Returns + ------- + list of PoseHeaderComponent + Components for body, face, left hand, and right hand. + """ + hand_limbs = _map_limbs(GENERAL_HAND_POINTS, HAND_LIMBS_NAMES) + return [ + PoseHeaderComponent( + name="BODY_136", + points=BODY_POINTS, + limbs=_map_limbs(BODY_POINTS, BODY_LIMBS_NAMES), + colors=[(0, 255, 0)], + point_format="XYC" + ), + PoseHeaderComponent( + name="FACE_136", + points=FACE_POINTS, + limbs=_map_limbs(FACE_POINTS, FACE_LIMBS_NAMES), + colors=[(255, 255, 255)], + point_format="XYC" + ), + PoseHeaderComponent( + name="LEFT_HAND_136", + points=GENERAL_HAND_POINTS, + limbs=hand_limbs, + colors=[(0, 255, 0)], + point_format="XYC" + ), + PoseHeaderComponent( + name="RIGHT_HAND_136", + points=GENERAL_HAND_POINTS, + limbs=hand_limbs, + colors=[(255, 128, 0)], + point_format="XYC" + ), + ] + + +AlphaPose_Components = get_alphapose_components() + + +def load_alphapose_json(json_path): + """ + Load AlphaPose results in either of two formats. + + FORMAT A (original list): + [{"image_id": "0.jpg", "keypoints": [x0, y0, c0, ...]}, ...] + + FORMAT B (extended dict with metadata): + {"frames": [...], "metadata": {"fps": float, "width": int, "height": int}} + + Returns + ------- + frames : list + Sorted list of frame detections. + meta : dict or None + Metadata dict if present, else None. + """ + with open(json_path, "r") as f: + raw = json.load(f) + + if isinstance(raw, dict) and "frames" in raw: + frames = raw["frames"] + meta = { + "fps": raw.get("metadata", {}).get("fps", None), + "width": raw.get("metadata", {}).get("width", None), + "height": raw.get("metadata", {}).get("height", None), + } + else: + frames = raw + meta = None + + def extract_frame_number(item): + matches = re.findall(r"\d+", item["image_id"]) + return int(matches[-1]) if matches else -1 + + return sorted(frames, key=extract_frame_number), meta + + +def parse_keypoints_and_confidence(flat): + """ + Parse AlphaPose flat keypoint list [x0, y0, c0, x1, y1, c1, ...]. + + Parameters + ---------- + flat : list or array-like + Flat list of keypoint values. + + Returns + ------- + xy : ndarray, shape (N, 2) + conf : ndarray, shape (N,) + n_keypoints : int + """ + n_values = len(flat) + if n_values == 136 * 3: + n_keypoints = 136 + elif n_values == 133 * 3: + n_keypoints = 133 + else: + raise ValueError( + f"Expected 136 (408 values) or 133 (399 values) keypoints, got {n_values} values." + ) + arr = np.array(flat).reshape(-1, 3) + return arr[:, :2], arr[:, 2], n_keypoints + + +def _apply_metadata(metadata, fps, width, height): + if metadata is not None: + fps = metadata["fps"] if metadata.get("fps") is not None else fps + width = metadata["width"] if metadata.get("width") is not None else width + height = metadata["height"] if metadata.get("height") is not None else height + return fps, width, height + + +def load_alphapose_wholebody_from_json(input_path: str, + version: float = 0.2, + fps: float = 24, + width: int = 1000, + height: int = 1000, + depth: int = 0) -> Pose: + """ + Load an AlphaPose WholeBody JSON file into a Pose object. + + Automatically detects whether the file contains 133 or 136 keypoints and + builds the appropriate components. Use alphapose_133.py directly to + enforce strict 133-only loading. + + Parameters + ---------- + input_path : str + Path to the AlphaPose JSON file. + version : float + Pose format version written to the header. + fps : float + Frames per second. Overridden by JSON metadata if present. + width : int + Frame width in pixels. Overridden by JSON metadata if present. + height : int + Frame height in pixels. Overridden by JSON metadata if present. + depth : int + Depth dimension size (0 for 2D poses). + + Returns + ------- + Pose + Loaded pose with header and body. + """ + frames, metadata = load_alphapose_json(input_path) + fps, width, height = _apply_metadata(metadata, fps, width, height) + + frames_xy = [] + frames_conf = [] + n_keypoints_detected = None + + for item in frames: + xy, conf, n_keypoints = parse_keypoints_and_confidence(item["keypoints"]) + n_keypoints_detected = n_keypoints + frames_xy.append(xy) + frames_conf.append(conf) + + xy_data = np.stack(frames_xy, axis=0)[:, None, :, :] + conf_data = np.stack(frames_conf, axis=0)[:, None, :] + + if n_keypoints_detected == 136: + components = AlphaPose_Components + else: + from .alphapose_133 import AlphaPose133_Components + components = AlphaPose133_Components + + header = PoseHeader( + version=version, + dimensions=PoseHeaderDimensions(width=width, height=height, depth=depth), + components=components, + ) + body = NumPyPoseBody(fps=fps, data=xy_data, confidence=conf_data) + return Pose(header, body) diff --git a/src/python/pose_format/utils/alphapose_133.py b/src/python/pose_format/utils/alphapose_133.py new file mode 100644 index 0000000..8ceec93 --- /dev/null +++ b/src/python/pose_format/utils/alphapose_133.py @@ -0,0 +1,145 @@ +import numpy as np +from ..numpy.pose_body import NumPyPoseBody +from ..pose import Pose +from ..pose_header import PoseHeader, PoseHeaderComponent, PoseHeaderDimensions +from .alphapose import ( + FACE_POINTS, FACE_LIMBS_NAMES, GENERAL_HAND_POINTS, + LEFT_HAND_POINTS, RIGHT_HAND_POINTS, HAND_LIMBS_NAMES, + _map_limbs, load_alphapose_json, parse_keypoints_and_confidence, _apply_metadata, +) + +# 133-keypoint body (no neck, head_top, or pelvis compared to 136) +BODY_POINTS = [ + "nose", "left_eye", "right_eye", "left_ear", "right_ear", + "left_shoulder", "right_shoulder", "left_elbow", "right_elbow", + "left_wrist", "right_wrist", "left_hip", "right_hip", + "left_knee", "right_knee", "left_ankle", "right_ankle", + "left_big_toe", "left_small_toe", "left_heel", + "right_big_toe", "right_small_toe", "right_heel", +] + +BODY_LIMBS_NAMES = [ + ("left_ankle", "left_knee"), ("left_knee", "left_hip"), + ("right_ankle", "right_knee"), ("right_knee", "right_hip"), + ("left_hip", "right_hip"), + ("left_shoulder", "left_hip"), ("right_shoulder", "right_hip"), + ("left_shoulder", "right_shoulder"), + ("left_shoulder", "left_elbow"), ("right_shoulder", "right_elbow"), + ("left_elbow", "left_wrist"), ("right_elbow", "right_wrist"), + ("left_eye", "right_eye"), ("nose", "left_eye"), ("nose", "right_eye"), + ("left_eye", "left_ear"), ("right_eye", "right_ear"), + ("left_ear", "left_shoulder"), ("right_ear", "right_shoulder"), + ("left_ankle", "left_big_toe"), ("left_ankle", "left_small_toe"), ("left_ankle", "left_heel"), + ("right_ankle", "right_big_toe"), ("right_ankle", "right_small_toe"), ("right_ankle", "right_heel"), +] + + +def get_alphapose_133_components(): + """ + Returns AlphaPose WholeBody-133 component definitions. + + Returns + ------- + list of PoseHeaderComponent + Components for body, face, left hand, and right hand. + """ + hand_limbs = _map_limbs(GENERAL_HAND_POINTS, HAND_LIMBS_NAMES) + return [ + PoseHeaderComponent( + name="BODY_133", + points=BODY_POINTS, + limbs=_map_limbs(BODY_POINTS, BODY_LIMBS_NAMES), + colors=[(0, 255, 0)], + point_format="XYC" + ), + PoseHeaderComponent( + name="FACE_133", + points=FACE_POINTS, + limbs=_map_limbs(FACE_POINTS, FACE_LIMBS_NAMES), + colors=[(255, 255, 255)], + point_format="XYC" + ), + PoseHeaderComponent( + name="LEFT_HAND_133", + points=GENERAL_HAND_POINTS, + limbs=hand_limbs, + colors=[(0, 255, 0)], + point_format="XYC" + ), + PoseHeaderComponent( + name="RIGHT_HAND_133", + points=GENERAL_HAND_POINTS, + limbs=hand_limbs, + colors=[(255, 128, 0)], + point_format="XYC" + ), + ] + + +AlphaPose133_Components = get_alphapose_133_components() + + +def load_alphapose_wholebody_from_json(input_path: str, + version: float = 0.2, + fps: float = 24, + width: int = 1000, + height: int = 1000, + depth: int = 0) -> Pose: + """ + Load an AlphaPose WholeBody-133 JSON file into a Pose object. + + Raises ValueError if the file contains 136-keypoint data; use + pose_format.utils.alphapose.load_alphapose_wholebody_from_json for + auto-detection. + + Parameters + ---------- + input_path : str + Path to the AlphaPose JSON file. + version : float + Pose format version written to the header. + fps : float + Frames per second. Overridden by JSON metadata if present. + width : int + Frame width in pixels. Overridden by JSON metadata if present. + height : int + Frame height in pixels. Overridden by JSON metadata if present. + depth : int + Depth dimension size (0 for 2D poses). + + Returns + ------- + Pose + Loaded pose with header and body. + + Raises + ------ + ValueError + If the JSON contains 136-keypoint data. + """ + frames, metadata = load_alphapose_json(input_path) + fps, width, height = _apply_metadata(metadata, fps, width, height) + + frames_xy = [] + frames_conf = [] + + for item in frames: + xy, conf, n_keypoints = parse_keypoints_and_confidence(item["keypoints"]) + if n_keypoints == 136: + raise ValueError( + "This file contains 136-keypoint AlphaPose data. " + "Use pose_format.utils.alphapose.load_alphapose_wholebody_from_json instead." + ) + frames_xy.append(xy) + frames_conf.append(conf) + + xy_data = np.stack(frames_xy, axis=0)[:, None, :, :] + conf_data = np.stack(frames_conf, axis=0)[:, None, :] + + header = PoseHeader( + version=version, + dimensions=PoseHeaderDimensions(width=width, height=height, depth=depth), + components=AlphaPose133_Components, + ) + body = NumPyPoseBody(fps=fps, data=xy_data, confidence=conf_data) + return Pose(header, body) diff --git a/src/python/pose_format/utils/alphapose_test.py b/src/python/pose_format/utils/alphapose_test.py new file mode 100644 index 0000000..64f208e --- /dev/null +++ b/src/python/pose_format/utils/alphapose_test.py @@ -0,0 +1,153 @@ +import json +import tempfile +import pytest +import numpy as np + +from pose_format.utils.alphapose import load_alphapose_wholebody_from_json +from pose_format.utils.alphapose_133 import ( + load_alphapose_wholebody_from_json as load_alphapose_133_from_json, +) + + +def _make_flat_keypoints(n: int, value: float = 1.0): + return [value] * (n * 3) + + +def _make_json_format_a(n_keypoints: int, n_frames: int = 3): + return [ + {"image_id": f"frame_{i:04d}.jpg", "keypoints": _make_flat_keypoints(n_keypoints)} + for i in range(n_frames) + ] + + +def _make_json_format_b(n_keypoints: int, n_frames: int = 3, fps=30.0, width=640, height=480): + return { + "frames": _make_json_format_a(n_keypoints, n_frames), + "metadata": {"fps": fps, "width": width, "height": height}, + } + + +def _write_json(data) -> str: + f = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) + json.dump(data, f) + f.flush() + return f.name + + +# --------------------------------------------------------------------------- +# Auto-detecting loader (alphapose.py) +# --------------------------------------------------------------------------- + +def test_auto_detect_136_shape(): + path = _write_json(_make_json_format_a(136, n_frames=3)) + pose = load_alphapose_wholebody_from_json(path) + assert pose.body.data.shape == (3, 1, 136, 2) + assert pose.body.confidence.shape == (3, 1, 136) + + +def test_auto_detect_133_shape(): + path = _write_json(_make_json_format_a(133, n_frames=3)) + pose = load_alphapose_wholebody_from_json(path) + assert pose.body.data.shape == (3, 1, 133, 2) + assert pose.body.confidence.shape == (3, 1, 133) + + +def test_auto_detect_136_component_names(): + path = _write_json(_make_json_format_a(136)) + pose = load_alphapose_wholebody_from_json(path) + assert [c.name for c in pose.header.components] == ["BODY_136", "FACE_136", "LEFT_HAND_136", "RIGHT_HAND_136"] + + +def test_auto_detect_133_component_names(): + path = _write_json(_make_json_format_a(133)) + pose = load_alphapose_wholebody_from_json(path) + assert [c.name for c in pose.header.components] == ["BODY_133", "FACE_133", "LEFT_HAND_133", "RIGHT_HAND_133"] + + +def test_auto_detect_rejects_unknown_count(): + path = _write_json(_make_json_format_a(100)) + with pytest.raises(ValueError): + load_alphapose_wholebody_from_json(path) + + +# --------------------------------------------------------------------------- +# 136-keypoint loader: defaults and metadata +# --------------------------------------------------------------------------- + +def test_load_136_defaults(): + path = _write_json(_make_json_format_a(136)) + pose = load_alphapose_wholebody_from_json(path) + assert pose.body.fps == 24 + assert pose.header.dimensions.width == 1000 + assert pose.header.dimensions.height == 1000 + + +def test_load_136_metadata_override(): + path = _write_json(_make_json_format_b(136, fps=60.0, width=1920, height=1080)) + pose = load_alphapose_wholebody_from_json(path) + assert pose.body.fps == 60.0 + assert pose.header.dimensions.width == 1920 + assert pose.header.dimensions.height == 1080 + + +# --------------------------------------------------------------------------- +# 133-keypoint strict loader (alphapose_133.py) +# --------------------------------------------------------------------------- + +def test_load_133_shape(): + path = _write_json(_make_json_format_a(133, n_frames=3)) + pose = load_alphapose_133_from_json(path) + assert pose.body.data.shape == (3, 1, 133, 2) + assert pose.body.confidence.shape == (3, 1, 133) + + +def test_load_133_defaults(): + path = _write_json(_make_json_format_a(133)) + pose = load_alphapose_133_from_json(path) + assert pose.body.fps == 24 + assert pose.header.dimensions.width == 1000 + assert pose.header.dimensions.height == 1000 + + +def test_load_133_metadata_override(): + path = _write_json(_make_json_format_b(133, fps=25.0, width=720, height=576)) + pose = load_alphapose_133_from_json(path) + assert pose.body.fps == 25.0 + assert pose.header.dimensions.width == 720 + assert pose.header.dimensions.height == 576 + + +def test_load_133_rejects_136(): + path = _write_json(_make_json_format_a(136)) + with pytest.raises(ValueError, match="136-keypoint"): + load_alphapose_133_from_json(path) + + +def test_load_133_component_names(): + path = _write_json(_make_json_format_a(133)) + pose = load_alphapose_133_from_json(path) + assert [c.name for c in pose.header.components] == ["BODY_133", "FACE_133", "LEFT_HAND_133", "RIGHT_HAND_133"] + + +# --------------------------------------------------------------------------- +# Frame sorting +# --------------------------------------------------------------------------- + +def test_frame_sort_order(): + # Image IDs where the frame number is NOT the first integer ("v2/frame_NNN.jpg"). + # A correct sort uses the last integer; an incorrect one would sort by "2" instead. + def kpts(sentinel): + # x=sentinel, y=sentinel, conf=1.0 for every keypoint (non-zero conf avoids masking) + return [v for _ in range(136) for v in (sentinel, sentinel, 1.0)] + + data = [ + {"image_id": "v2/frame_0003.jpg", "keypoints": kpts(3.0)}, + {"image_id": "v2/frame_0001.jpg", "keypoints": kpts(1.0)}, + {"image_id": "v2/frame_0002.jpg", "keypoints": kpts(2.0)}, + ] + path = _write_json(data) + pose = load_alphapose_wholebody_from_json(path) + # Correctly sorted: frame_0001 → [0], frame_0002 → [1], frame_0003 → [2] + assert pose.body.data[0, 0, 0, 0] == pytest.approx(1.0) + assert pose.body.data[1, 0, 0, 0] == pytest.approx(2.0) + assert pose.body.data[2, 0, 0, 0] == pytest.approx(3.0) diff --git a/src/python/pose_format/utils/generic.py b/src/python/pose_format/utils/generic.py index bff3e63..357aedb 100644 --- a/src/python/pose_format/utils/generic.py +++ b/src/python/pose_format/utils/generic.py @@ -9,12 +9,13 @@ from pose_format.utils.openpose import OpenPose_Components from pose_format.utils.openpose import BODY_POINTS as OPENPOSE_BODY_POINTS from pose_format.utils.openpose_135 import OpenPose_Components as OpenPose135_Components +from pose_format.utils.alphapose import get_alphapose_components +from pose_format.utils.alphapose_133 import get_alphapose_133_components # from pose_format.utils.holistic import holistic_components # The import above creates an error: ImportError: Please install mediapipe with: pip install mediapipe -KnownPoseFormat = Literal["holistic", "openpose", "openpose_135"] - +KnownPoseFormat = Literal["holistic", "openpose", "openpose_135", "alphapose_133", "alphapose_136"] def get_component_names( pose_or_header_or_components: Union[Pose,PoseHeader]) -> List[str]: @@ -41,6 +42,10 @@ def detect_known_pose_format(pose_or_header: Union[Pose,PoseHeader]) -> KnownPos openpose_135_components = [c.name for c in OpenPose135_Components] + alphapose_133_components = [c.name for c in get_alphapose_133_components()] + + alphapose_136_components = [c.name for c in get_alphapose_components()] + for component_name in component_names: if component_name in mediapipe_components: return "holistic" @@ -48,6 +53,10 @@ def detect_known_pose_format(pose_or_header: Union[Pose,PoseHeader]) -> KnownPos return "openpose" if component_name in openpose_135_components: return "openpose_135" + if component_name in alphapose_133_components: + return "alphapose_133" + if component_name in alphapose_136_components: + return "alphapose_136" raise ValueError( f"Could not detect pose format, unknown pose header schema with component names: {component_names}" @@ -87,6 +96,17 @@ def pose_hide_legs(pose: Pose, remove: bool = False) -> Pose: # if any of the items in point_ points_to_remove_dict = {"pose_keypoints_2d": point_names_to_remove} + elif known_pose_format == "alphapose_133" or known_pose_format == "alphapose_136": + variant = known_pose_format[len("alphapose_"):] + point_names_to_remove = [ + "left_hip", "right_hip", + "left_knee", "right_knee", + "left_ankle", "right_ankle", + "left_big_toe", "left_small_toe", "left_heel", + "right_big_toe", "right_small_toe", "right_heel", + ] + points_to_remove_dict = {f"BODY_{variant}": point_names_to_remove} + else: raise NotImplementedError( f"Unsupported pose header schema {known_pose_format} for {pose_hide_legs.__name__}: {pose.header}" @@ -124,6 +144,10 @@ def pose_shoulders(pose_header: PoseHeader) -> Tuple[Tuple[str, str], Tuple[str, if known_pose_format == "openpose": return ("pose_keypoints_2d", "RShoulder"), ("pose_keypoints_2d", "LShoulder") + if known_pose_format == "alphapose_133" or known_pose_format == "alphapose_136": + variant = known_pose_format[len("alphapose_"):] + return (f"BODY_{variant}", "right_shoulder"), (f"BODY_{variant}", "left_shoulder") + raise NotImplementedError( f"Unsupported pose header schema {known_pose_format} for {pose_shoulders.__name__}: {pose_header}" ) @@ -142,6 +166,13 @@ def hands_indexes(pose_header: PoseHeader)-> List[int]: pose_header.get_point_index("hand_left_keypoints_2d", "M_CMC"), pose_header.get_point_index("hand_right_keypoints_2d", "M_CMC"), ] + + if known_pose_format == "alphapose_133" or known_pose_format == "alphapose_136": + variant = known_pose_format[len("alphapose_"):] + return [ + pose_header.get_point_index(f"LEFT_HAND_{variant}", "hand_9"), + pose_header.get_point_index(f"RIGHT_HAND_{variant}", "hand_9"), + ] raise NotImplementedError( f"Unsupported pose header schema {known_pose_format} for {hands_indexes.__name__}: {pose_header}" ) @@ -163,7 +194,10 @@ def hands_components(pose_header: PoseHeader)-> Tuple[Tuple[str, str], Tuple[str if known_pose_format == "openpose": return ("hand_left_keypoints_2d", "hand_right_keypoints_2d"), ("BASE", "P_CMC", "I_CMC"), ("BASE", "M_CMC") - + + if known_pose_format == "alphapose_133" or known_pose_format == "alphapose_136": + variant = known_pose_format[len("alphapose_"):] + return (f"LEFT_HAND_{variant}", f"RIGHT_HAND_{variant}"), ("hand_0", "hand_17", "hand_5"), ("hand_0", "hand_9") raise NotImplementedError( f"Unsupported pose header schema '{known_pose_format}' for {hands_components.__name__}: {pose_header}" ) @@ -209,6 +243,10 @@ def get_standard_components_for_known_format(known_pose_format: KnownPoseFormat) return OpenPose_Components if known_pose_format == "openpose_135": return OpenPose135_Components + if known_pose_format == "alphapose_133": + return get_alphapose_133_components() + if known_pose_format == "alphapose_136": + return get_alphapose_components() raise NotImplementedError(f"Unsupported pose header schema {known_pose_format}") @@ -274,6 +312,9 @@ def get_hand_wrist_index(pose: Pose, hand: str)-> int: return pose.header.get_point_index(f"{hand.upper()}_HAND_LANDMARKS", "WRIST") if known_pose_format == "openpose": return pose.header.get_point_index(f"hand_{hand.lower()}_keypoints_2d", "BASE") + if known_pose_format == "alphapose_133" or known_pose_format == "alphapose_136": + variant = known_pose_format[len("alphapose_"):] + return pose.header.get_point_index(f"{hand.upper()}_HAND_{variant}", "hand_0") raise NotImplementedError( f"Unsupported pose header schema {known_pose_format} for {get_hand_wrist_index.__name__}: {pose.header}" ) @@ -285,6 +326,9 @@ def get_body_hand_wrist_index(pose: Pose, hand: str)-> int: return pose.header.get_point_index("POSE_LANDMARKS", f"{hand.upper()}_WRIST") if known_pose_format == "openpose": return pose.header.get_point_index("pose_keypoints_2d", f"{hand.upper()[0]}Wrist") + if known_pose_format == "alphapose_133" or known_pose_format == "alphapose_136": + variant = known_pose_format[len("alphapose_"):] + return pose.header.get_point_index(f"BODY_{variant}", f"{hand.lower()}_wrist") raise NotImplementedError( f"Unsupported pose header schema {known_pose_format} for {get_body_hand_wrist_index.__name__}: {pose.header}" ) diff --git a/src/python/pose_format/utils/generic_test.py b/src/python/pose_format/utils/generic_test.py index e1b5912..c072bc8 100644 --- a/src/python/pose_format/utils/generic_test.py +++ b/src/python/pose_format/utils/generic_test.py @@ -85,8 +85,8 @@ def test_pose_shoulders(fake_poses: List[Pose]): for pose in fake_poses: shoulders = pose_shoulders(pose.header) assert len(shoulders) == 2 - assert "RIGHT" in shoulders[0][1] or shoulders[0][1][0] == "R" - assert "LEFT" in shoulders[1][1] or shoulders[1][1][0] == "L" + assert "RIGHT" in shoulders[0][1] or shoulders[0][1][0] in ("R", "r") + assert "LEFT" in shoulders[1][1] or shoulders[1][1][0] in ("L", "l") @pytest.mark.parametrize("fake_poses", TEST_POSE_FORMATS, indirect=["fake_poses"]) @@ -224,6 +224,22 @@ def test_pose_remove_legs(fake_poses: List[Pose]): component_index = c_names.index("pose_keypoints_2d") pose_with_legs_removed = pose_hide_legs(pose, remove=True) + for point_name in points_that_should_be_removed: + assert point_name not in pose_with_legs_removed.header.components[component_index].points, f"{pose_with_legs_removed.header.components[component_index].name},{pose_with_legs_removed.header.components[component_index].points}" + assert point_name in pose.header.components[component_index].points + + elif known_pose_format == "alphapose_133" or known_pose_format == "alphapose_136": + c_names = [c.name for c in pose.header.components] + points_that_should_be_removed = [ + "left_hip", "right_hip", + "left_knee", "right_knee", + "left_ankle", "right_ankle", + "left_big_toe", "left_small_toe", "left_heel", + "right_big_toe", "right_small_toe", "right_heel", + ] + component_index = c_names.index(f"BODY_{known_pose_format[-3:]}") + pose_with_legs_removed = pose_hide_legs(pose, remove=True) + for point_name in points_that_should_be_removed: assert point_name not in pose_with_legs_removed.header.components[component_index].points, f"{pose_with_legs_removed.header.components[component_index].name},{pose_with_legs_removed.header.components[component_index].points}" assert point_name in pose.header.components[component_index].points @@ -272,6 +288,9 @@ def test_fake_pose(known_pose_format: KnownPoseFormat): assert point_formats[0] == "XYC" elif detected_format == 'openpose_135': assert point_formats[0] == "XYC" + elif detected_format == 'alphapose_133' or detected_format == 'alphapose_136': + assert point_formats[0] == "XYC" + assert detected_format == known_pose_format assert pose.body.fps == fps