Add MMPose wholebody estimator

src/python/pose_format/utils/mmposewholebody.py

@@ -0,0 +1,113 @@
+import numpy as np
+import numpy.ma as ma
+
+try:
+    from mmpose.apis import MMPoseInferencer
+except ImportError:
+    raise ImportError(
+        "Please install MMPose and its dependencies. For GPU support, mmcv must be installed\n"
+        "from the OpenMMLab CUDA-specific index (see https://mmcv.readthedocs.io/en/latest/get_started/installation.html).\n"
+        "The remaining packages: pip install 'mmpose>=1.3.2' 'mmengine>=0.10.7' 'mmdet>=3.3.0'"
+    )
+
+from ..numpy.pose_body import NumPyPoseBody
+from ..pose import Pose
+from ..pose_header import PoseHeader, PoseHeaderDimensions
+from .cocowholebody133_header import cocowholebody_components
+
+NUM_KEYPOINTS = 133
+
+
+def estimate_mmpose_wholebody(input_path: str,
+                         version: float = 0.2,
+                         fps: float = 24,
+                         width: int = 1000,
+                         height: int = 1000,
+                         depth: int = 0) -> Pose:
+    """
+    Run MMPose wholebody inference on a video and return a Pose object.
+
+    Parameters
+    ----------
+    input_path : str
+        Path to the input video file.
+    version : float
+        Pose format version written to the header.
+    fps : float
+        Frames per second stored in the pose body.
+    width : int
+        Frame width in pixels stored in the header dimensions.
+    height : int
+        Frame height in pixels stored in the header dimensions.
+    depth : int
+        Depth dimension size (0 for 2D poses).
+
+    Returns
+    -------
+    Pose
+        Loaded pose with header and body.
+    """
+    header = PoseHeader(
+        version=version,
+        dimensions=PoseHeaderDimensions(width=width, height=height, depth=depth),
+        components=cocowholebody_components(),
+    )
+    body = _process_video(input_path, fps)
+    return Pose(header, body)
+
+
+def _process_video(input_path: str, fps: float, use_cpu: bool = False) -> NumPyPoseBody:
+    """
+    Run MMPose wholebody inference and convert frame results to NumPyPoseBody.
+
+    Parameters
+    ----------
+    input_path : str
+        Path to the input video file.
+    fps : float
+        Frames per second to store in the pose body.
+    use_cpu : bool
+        If True, run inference on CPU (slow; useful when no GPU is available).
+
+    Returns
+    -------
+    NumPyPoseBody
+    """
+    device = 'cpu' if use_cpu else None
+    inferencer_kwargs = {'wholebody': True}
+    if device is not None:
+        inferencer_kwargs['device'] = device
+
+    inferencer = MMPoseInferencer('wholebody', **({'device': device} if device else {}))
+    result_generator = inferencer(input_path, show=False, return_vis=False)
+
+    frames_xy = []
+    frames_conf = []
+    frames_mask = []  # True = valid, False = masked out (no detection)
+
+    for result in result_generator:
+        predictions_by_frame = result['predictions']  # list of per-person dicts for this frame
+
+        if len(predictions_by_frame) == 0 or len(predictions_by_frame[0]) == 0:
+            # No person detected in this frame. Insert a zeroed, fully-masked row so
+            # the frame count stays aligned with the video. Callers can distinguish
+            # "no detection" from a real zero-coordinate keypoint via the mask.
+            frames_xy.append(np.zeros((1, NUM_KEYPOINTS, 2), dtype=np.float32))
+            frames_conf.append(np.zeros((1, NUM_KEYPOINTS), dtype=np.float32))
+            frames_mask.append(True)   # True = mask this frame entirely
+        else:
+            person = predictions_by_frame[0][0]
+            frames_xy.append(np.array(person['keypoints'], dtype=np.float32)[None])    # (1, 133, 2)
+            frames_conf.append(np.array(person['keypoint_scores'], dtype=np.float32)[None])  # (1, 133)
+            frames_mask.append(False)  # False = keep (not masked)
+
+    xy_data = np.concatenate(frames_xy, axis=0)[:, None, :, :]    # (T, 1, 133, 2)
+    conf_data = np.concatenate(frames_conf, axis=0)[:, None, :]   # (T, 1, 133)
+
+    # Build the masked array: mask=True on empty frames so downstream code
+    # can treat them as missing rather than as detected-at-origin.
+    mask = np.array(frames_mask)                                    # (T,)
+    xy_mask = mask[:, None, None, None] * np.ones_like(xy_data, dtype=bool)
+    masked_xy = ma.array(xy_data, mask=xy_mask)
+
+    return NumPyPoseBody(fps=fps, data=masked_xy, confidence=conf_data)

src/python/pose_format/utils/mmposewholebody_test.py

@@ -0,0 +1,118 @@
+import sys
+from unittest.mock import MagicMock
+
+import numpy as np
+import numpy.ma as ma
+import pytest
+
+# Stub the MMPose package and its dependencies before our module is imported.
+# mmposewholebody.py does `from mmpose.apis import MMPoseInferencer` at module level,
+# so sys.modules must be populated before the first import of that module.
+for _mod in ["mmpose", "mmpose.apis", "mmcv", "mmengine", "mmdet"]:
+    sys.modules.setdefault(_mod, MagicMock())
+
+from pose_format.utils import mmposewholebody  # noqa: E402 — must come after the stubs above
+from pose_format.utils.mmposewholebody import estimate_mmpose_wholebody  # noqa: E402
+from pose_format.utils.cocowholebody133_header import cocowholebody_components  # noqa: E402
+
+NUM_KEYPOINTS = 133
+
+
+# ---------------------------------------------------------------------------
+# Header / components tests — no MMPose installation required
+# ---------------------------------------------------------------------------
+
+def test_components_total_keypoints():
+    assert sum(len(c.points) for c in cocowholebody_components()) == NUM_KEYPOINTS
+
+
+def test_components_names():
+    names = [c.name for c in cocowholebody_components()]
+    assert names == ["BODY", "FACE", "LEFT_HAND", "RIGHT_HAND"]
+
+
+def test_components_point_format():
+    for c in cocowholebody_components():
+        assert c.format == "XYC"
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _fake_result(num_keypoints: int = NUM_KEYPOINTS):
+    """Single-frame MMPose result with one detected person."""
+    return {
+        "predictions": [[{
+            "keypoints": np.random.rand(num_keypoints, 2).tolist(),
+            "keypoint_scores": np.random.rand(num_keypoints).tolist(),
+        }]]
+    }
+
+
+def _empty_result():
+    """Single-frame MMPose result with no detected person."""
+    return {"predictions": []}
+
+
+def _make_inferencer(results):
+    """Return a patched MMPoseInferencer class whose instance yields `results`."""
+    fake_instance = MagicMock()
+    fake_instance.return_value = iter(results)
+    return MagicMock(return_value=fake_instance)
+
+
+# ---------------------------------------------------------------------------
+# Loader tests (MMPoseInferencer is mocked)
+# ---------------------------------------------------------------------------
+
+def test_load_shape(monkeypatch, tmp_path):
+    """Output Pose has the right frame/keypoint shape."""
+    monkeypatch.setattr(mmposewholebody, "MMPoseInferencer",
+                        _make_inferencer([_fake_result(), _fake_result(), _fake_result()]))
+    pose = estimate_mmpose_wholebody(str(tmp_path / "video.mp4"), fps=25.0, width=1280, height=720)
+
+    assert pose.body.data.shape == (3, 1, NUM_KEYPOINTS, 2)
+    assert pose.body.fps == 25.0
+    assert pose.header.dimensions.width == 1280
+    assert pose.header.dimensions.height == 720
+
+
+def test_load_component_names(monkeypatch, tmp_path):
+    monkeypatch.setattr(mmposewholebody, "MMPoseInferencer",
+                        _make_inferencer([_fake_result()]))
+    pose = estimate_mmpose_wholebody(str(tmp_path / "video.mp4"))
+
+    assert [c.name for c in pose.header.components] == ["BODY", "FACE", "LEFT_HAND", "RIGHT_HAND"]
+
+
+def test_empty_frame_is_masked(monkeypatch, tmp_path):
+    """Frames with no detection are present in the output but fully masked."""
+    results = [_fake_result(), _empty_result(), _fake_result()]
+    monkeypatch.setattr(mmposewholebody, "MMPoseInferencer", _make_inferencer(results))
+    pose = estimate_mmpose_wholebody(str(tmp_path / "video.mp4"))
+
+    # All three frames must be present so frame count matches the video.
+    assert pose.body.data.shape[0] == 3
+
+    # Frame 1 (index 1) must be fully masked; frames 0 and 2 must not be.
+    assert pose.body.data[1].mask.all(), "empty frame should be fully masked"
+    assert not pose.body.data[0].mask.all(), "detected frame should not be fully masked"
+    assert not pose.body.data[2].mask.all(), "detected frame should not be fully masked"
+
+
+def test_all_empty_frames(monkeypatch, tmp_path):
+    """A video where no person is ever detected produces a fully masked Pose."""
+    results = [_empty_result(), _empty_result()]
+    monkeypatch.setattr(mmposewholebody, "MMPoseInferencer", _make_inferencer(results))
+    pose = estimate_mmpose_wholebody(str(tmp_path / "video.mp4"))
+
+    assert pose.body.data.shape[0] == 2
+    assert pose.body.data.mask.all()
+
+
+def test_version_default(monkeypatch, tmp_path):
+    monkeypatch.setattr(mmposewholebody, "MMPoseInferencer",
+                        _make_inferencer([_fake_result()]))
+    pose = estimate_mmpose_wholebody(str(tmp_path / "video.mp4"))
+    assert pose.header.version == 0.2

src/python/pyproject.toml

@@ -33,6 +33,13 @@ mediapipe = [
     "mediapipe<0.10.30",
 ]
 
+mmpose = [
+    "mmcv>=2.1.0",
+    "mmengine>=0.10.7",
+    "mmdet>=3.3.0",
+    "mmpose>=1.3.2",
+]
+
 [tool.setuptools]
 packages = [
     "pose_format",