update tests and readme: results-validation with valid_frames mask, ResultStatus

Author: deruyter92

fmpose3d/inference_api/README.md

@@ -100,6 +100,10 @@ Convenience constructor for the **animal** pipeline. Sets `model_type="fmpose3d_
 #### `predict(source, *, camera_rotation, seed, progress)` → `Pose3DResult`
 
 End-to-end prediction: 2D estimation followed by 3D lifting in a single call.
+Raises `ValueError` when 2D estimation is unusable for lifting
+(`Pose2DResult.status` is `empty` or `invalid`).
+For partial 2D detections, invalid frames are masked to `NaN` in
+`Pose3DResult.poses_3d` and `Pose3DResult.poses_3d_world`.
 
 | Parameter | Type | Description |
 |---|---|---|
@@ -121,7 +125,9 @@ Runs only the 2D pose estimation step.
 | `source` | `Source` | Same flexible input as `predict()`. |
 | `progress` | `ProgressCallback \| None` | Optional progress callback. |
 
-**Returns:** `Pose2DResult` containing `keypoints`, `scores`, and `image_size`.
+**Returns:** `Pose2DResult` containing `keypoints`, `scores`, `image_size`,
+and `valid_frames_mask`. The object also exposes derived properties
+`status` and `status_message`.
 
 ---
 
@@ -168,13 +174,35 @@ Source = Union[str, Path, np.ndarray, Sequence[Union[str, Path, np.ndarray]]]
 | `keypoints` | `ndarray` | 2D keypoints, shape `(num_persons, num_frames, J, 2)`. |
 | `scores` | `ndarray` | Per-joint confidence, shape `(num_persons, num_frames, J)`. |
 | `image_size` | `tuple[int, int]` | `(height, width)` of source frames. |
+| `valid_frames_mask` | `ndarray \| None` | Boolean mask, shape `(num_frames,)`, indicating frames with valid detections. |
+
+Computed properties:
+
+- `status` → `ResultStatus`
+- `status_message` → `str`
+
+#### `ResultStatus`
+
+String enum values:
+
+- `success` — valid detections in all frames
+- `partial` — valid detections in a subset of frames
+- `empty` — no valid detections in any frame
+- `invalid` — output predictions are unusable/malformed
+- `unknown` — validity metadata missing or malformed
 
 #### `Pose3DResult`
 
 | Field | Type | Description |
 |---|---|---|
 | `poses_3d` | `ndarray` | Root-relative 3D poses, shape `(num_frames, J, 3)`. |
 | `poses_3d_world` | `ndarray` | Post-processed 3D poses, shape `(num_frames, J, 3)`. For humans: world-coordinate poses. For animals: limb-regularized poses. |
+| `valid_frames_mask` | `ndarray \| None` | Boolean mask, shape `(num_frames,)`, indicating frames with valid 3D output. |
+
+Computed properties:
+
+- `status` → `ResultStatus`
+- `status_message` → `str`
 
 
 
@@ -187,14 +215,14 @@ Source = Union[str, Path, np.ndarray, Sequence[Union[str, Path, np.ndarray]]]
 Default 2D estimator for the human pipeline. Wraps HRNet + YOLO with a COCO → H36M keypoint conversion.
 
 - `setup_runtime()` — Loads YOLO + HRNet models.
-- `predict(frames: ndarray)` → `(keypoints, scores)` — Returns H36M-format 2D keypoints from BGR frames `(N, H, W, C)`.
+- `predict(frames: ndarray)` → `(keypoints, scores, valid_frames_mask)` — Returns H36M-format 2D keypoints from BGR frames `(N, H, W, C)` plus a frame-level validity mask.
 
 #### `SuperAnimalEstimator(cfg: SuperAnimalConfig | None)`
 
 2D estimator for the animal pipeline. Uses DeepLabCut SuperAnimal and maps quadruped80K keypoints to the 26-joint Animal3D layout.
 
 - `setup_runtime()` — No-op (DLC loads lazily).
-- `predict(frames: ndarray)` → `(keypoints, scores)` — Returns Animal3D-format 2D keypoints from BGR frames.
+- `predict(frames: ndarray)` → `(keypoints, scores, valid_frames_mask)` — Returns Animal3D-format 2D keypoints plus a frame-level validity mask.
 
 ---
 

tests/fmpose3d_api/test_fmpose3d.py

@@ -25,6 +25,7 @@
     HRNetEstimator,
     HumanPostProcessor,
     Pose2DResult,
+    ResultStatus,
     Pose3DResult,
     SuperAnimalEstimator,
     _default_components,
@@ -615,8 +616,9 @@ def test_predict_end_to_end_with_mock_estimator(self):
 
         mock_kpts = np.random.randn(1, 1, 17, 2).astype("float32")
         mock_scores = np.ones((1, 1, 17), dtype="float32")
+        mock_mask = np.array([True], dtype=bool)
         api._estimator_2d = MagicMock()
-        api._estimator_2d.predict.return_value = (mock_kpts, mock_scores)
+        api._estimator_2d.predict.return_value = (mock_kpts, mock_scores, mock_mask)
         api._estimator_2d.setup_runtime = MagicMock()
 
         frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
@@ -626,6 +628,24 @@ def test_predict_end_to_end_with_mock_estimator(self):
         assert result.poses_3d.shape == (1, 17, 3)
         api._estimator_2d.predict.assert_called_once()
 
+    def test_predict_applies_partial_2d_mask_to_3d(self):
+        """predict() masks invalid 2D frames to NaN in 3D outputs."""
+        api = _make_ready_api("fmpose3d_humans", test_augmentation=False)
+        mock_kpts = np.random.randn(1, 3, 17, 2).astype("float32")
+        mock_scores = np.ones((1, 3, 17), dtype="float32")
+        mask = np.array([True, False, True], dtype=bool)
+        api._estimator_2d = MagicMock()
+        api._estimator_2d.predict.return_value = (mock_kpts, mock_scores, mask)
+        api._estimator_2d.setup_runtime = MagicMock()
+
+        frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
+        result = api.predict([frame, frame, frame], seed=42)
+
+        np.testing.assert_array_equal(result.valid_frames_mask, mask)
+        assert result.status == ResultStatus.PARTIAL
+        assert np.all(np.isnan(result.poses_3d[1]))
+        assert np.all(np.isnan(result.poses_3d_world[1]))
+
 
 # =========================================================================
 # Unit tests — dataclasses
@@ -648,12 +668,41 @@ def test_pose2d_result_default_image_size(self):
         )
         assert result.image_size == (0, 0)
 
+    def test_pose2d_status_success(self):
+        result = Pose2DResult(
+            keypoints=np.zeros((1, 2, 17, 2)),
+            scores=np.zeros((1, 2, 17)),
+            valid_frames_mask=np.array([True, True], dtype=bool),
+        )
+        assert result.status == ResultStatus.SUCCESS
+        assert "all frames" in result.status_message
+
+    def test_pose2d_status_partial(self):
+        result = Pose2DResult(
+            keypoints=np.zeros((1, 2, 17, 2)),
+            scores=np.zeros((1, 2, 17)),
+            valid_frames_mask=np.array([True, False], dtype=bool),
+        )
+        assert result.status == ResultStatus.PARTIAL
+        assert "subset" in result.status_message
+
+    def test_pose2d_status_invalid_mask_length(self):
+        result = Pose2DResult(
+            keypoints=np.zeros((1, 2, 17, 2)),
+            scores=np.zeros((1, 2, 17)),
+            valid_frames_mask=np.array([True], dtype=bool),
+        )
+        assert result.status == ResultStatus.INVALID
+        assert "mismatches" in result.status_message
+
     def test_pose3d_result(self):
         p3d = np.random.randn(10, 17, 3)
         pw = np.random.randn(10, 17, 3)
-        result = Pose3DResult(poses_3d=p3d, poses_3d_world=pw)
+        mask = np.ones((10,), dtype=bool)
+        result = Pose3DResult(poses_3d=p3d, poses_3d_world=pw, valid_frames_mask=mask)
         assert result.poses_3d is p3d
         assert result.poses_3d_world is pw
+        assert result.status == ResultStatus.SUCCESS
 
 
 # =========================================================================
@@ -672,12 +721,13 @@ def test_predict_returns_zeros_when_no_bodyparts(self):
             "deeplabcut.pose_estimation_pytorch.apis.superanimal_analyze_images",
         ) as mock_fn:
             mock_fn.return_value = {"frame.png": {"bodyparts": None}}
-            kpts, scores = estimator.predict(frames)
+            kpts, scores, mask = estimator.predict(frames)
 
         assert kpts.shape == (1, 2, 26, 2)
         np.testing.assert_allclose(kpts, 0.0)
         assert scores.shape == (1, 2, 26)
-        np.testing.assert_allclose(scores, 1.0)
+        np.testing.assert_allclose(scores, 0.0)
+        np.testing.assert_array_equal(mask, np.array([False, False]))
 
     def test_predict_maps_valid_bodyparts(self):
         """Valid DLC bodyparts are mapped to Animal3D layout."""
@@ -692,9 +742,10 @@ def test_predict_maps_valid_bodyparts(self):
             "deeplabcut.pose_estimation_pytorch.apis.superanimal_analyze_images",
         ) as mock_fn:
             mock_fn.return_value = {"frame.png": {"bodyparts": fake_bp}}
-            kpts, scores = estimator.predict(frames)
+            kpts, scores, mask = estimator.predict(frames)
 
         assert kpts.shape == (1, 1, 26, 2)
         assert scores.shape == (1, 1, 26)
+        np.testing.assert_array_equal(mask, np.array([True]))
         # target[24] ← source[0] → (0*3, 0*3+1) = (0.0, 1.0)
         np.testing.assert_allclose(kpts[0, 0, 24], fake_bp[0, 0, :2])