diff --git a/src/python/pose_format/utils/mmposewholebody.py b/src/python/pose_format/utils/mmposewholebody.py new file mode 100644 index 0000000..1a883b1 --- /dev/null +++ b/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) diff --git a/src/python/pose_format/utils/mmposewholebody_test.py b/src/python/pose_format/utils/mmposewholebody_test.py new file mode 100644 index 0000000..d172a03 --- /dev/null +++ b/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 diff --git a/src/python/pyproject.toml b/src/python/pyproject.toml index 5bcd451..b308e37 100644 --- a/src/python/pyproject.toml +++ b/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",