dlc_processor: add lightweight validation of shape and dtype (+tests)

Author: deruyter92

dlclivegui/services/dlc_processor.py

@@ -15,7 +15,7 @@
 import numpy as np
 from PySide6.QtCore import QObject, Signal
 
-from dlclivegui.config import DLCProcessorSettings
+from dlclivegui.config import DLCProcessorSettings, ModelType
 from dlclivegui.processors.processor_utils import instantiate_from_scan
 from dlclivegui.temp import Engine  # type: ignore # TODO use main package enum when released
 
@@ -37,6 +37,63 @@
 class PoseResult:
     pose: np.ndarray | None
     timestamp: float
+    packet: "PosePacketV0 | None" = None
+
+
+@dataclass(slots=True, frozen=True)
+class PoseSource:
+    backend: str    # e.g. "DLCLive"
+    model_type: ModelType | None = None
+
+
+@dataclass(slots=True, frozen=True)
+class PosePacketV0:
+    schema_version: int = 0
+    keypoints: np.ndarray | None = None
+    keypoint_names: list[str] | None = None
+    individual_ids: list[str] | None = None
+    source: PoseSource = PoseSource(backend="DLCLive")
+    raw: Any | None = None
+
+
+def validate_pose_array(pose: Any, *, source_backend: str = "DLCLive") -> np.ndarray:
+    """
+    Validate pose output shape and dtype.
+
+    Accepted runner output shapes:
+    - (K, 3): single-animal
+    - (N, K, 3): multi-animal
+    """
+    try:
+        arr = np.asarray(pose)
+    except Exception as exc:
+        raise ValueError(
+            f"{source_backend} returned an invalid pose output format: could not convert to array ({exc})"
+        ) from exc
+
+    if arr.ndim not in (2, 3):
+        raise ValueError(
+            f"{source_backend} returned an invalid pose output format: expected a 2D or 3D array, got ndim={arr.ndim}, shape={arr.shape!r}"
+        )
+
+    if arr.shape[-1] != 3:
+        raise ValueError(
+            f"{source_backend} returned an invalid pose output format: expected last dimension size 3 (x, y, likelihood), got shape={arr.shape!r}"
+        )
+
+    if arr.ndim == 2 and arr.shape[0] <= 0:
+        raise ValueError(f"{source_backend} returned an invalid pose output format: expected at least one keypoint")
+    if arr.ndim == 3 and (arr.shape[0] <= 0 or arr.shape[1] <= 0):
+        raise ValueError(
+            f"{source_backend} returned an invalid pose output format: expected at least one individual and one keypoint, got shape={arr.shape!r}"
+        )
+
+    if not np.issubdtype(arr.dtype, np.number):
+        raise ValueError(
+            f"{source_backend} returned an invalid pose output format: expected numeric values, got dtype={arr.dtype}"
+        )
+
+    return arr
 
 
 @dataclass
@@ -269,8 +326,17 @@ def _process_frame(
         # Time GPU inference (and processor overhead when present)
         with self._timed_processor() as proc_holder:
             inference_start = time.perf_counter()
-            pose = self._dlc.get_pose(frame, frame_time=timestamp)
+            raw_pose: Any = self._dlc.get_pose(frame, frame_time=timestamp)
             inference_time = time.perf_counter() - inference_start
+        pose_arr: np.ndarray = validate_pose_array(raw_pose, source_backend="DLCLive")
+        pose_packet = PosePacketV0(
+            schema_version=0,
+            keypoints=pose_arr,
+            keypoint_names=None,
+            individual_ids=None,
+            source=PoseSource(backend="DLCLive", model_type=self._settings.model_type),
+            raw=raw_pose,
+        )
 
         processor_overhead = 0.0
         gpu_inference_time = inference_time
@@ -280,7 +346,7 @@ def _process_frame(
 
         # Emit pose (measure signal overhead)
         signal_start = time.perf_counter()
-        self.pose_ready.emit(PoseResult(pose=pose, timestamp=timestamp))
+        self.pose_ready.emit(PoseResult(pose=pose_packet.keypoints, timestamp=timestamp, packet=pose_packet))
         signal_time = time.perf_counter() - signal_start
 
         end_ts = time.perf_counter()

tests/services/test_pose_contract.py

@@ -0,0 +1,40 @@
+import numpy as np
+import pytest
+
+from dlclivegui.services.dlc_processor import validate_pose_array
+
+
+@pytest.mark.unit
+def test_validate_pose_array_keeps_single_animal_shape():
+    pose = np.ones((5, 3), dtype=np.float64)
+    out = validate_pose_array(pose)
+    assert out.shape == (5, 3)
+    assert out.dtype == np.float64
+
+
+@pytest.mark.unit
+def test_validate_pose_array_accepts_multi_animal():
+    pose = np.ones((2, 5, 3), dtype=np.float32)
+    out = validate_pose_array(pose)
+    assert out.shape == (2, 5, 3)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "bad_pose,expected",
+    [
+        (np.ones((5, 2), dtype=np.float32), "last dimension size 3"),
+        (np.ones((2, 5, 4), dtype=np.float32), "last dimension size 3"),
+        (np.ones((3,), dtype=np.float32), "expected a 2D or 3D array"),
+    ],
+)
+def test_validate_pose_array_rejects_invalid_shapes(bad_pose, expected):
+    with pytest.raises(ValueError, match=expected):
+        validate_pose_array(bad_pose)
+
+
+@pytest.mark.unit
+def test_validate_pose_array_rejects_non_numeric():
+    pose = np.array([[["x", "y", "p"]]], dtype=object)
+    with pytest.raises(ValueError, match="expected numeric values"):
+        validate_pose_array(pose)