feat(v1): add verified pose adapter

2026-07-21 09:53:03 +08:00
parent 2805905486
commit f80fb51a85
8 changed files with 267 additions and 15 deletions
@@ -0,0 +1,172 @@
+"""Verified Ultralytics Pose adapter for the V1 event pipeline."""
+
+import hashlib
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, List, Optional, Sequence, Tuple
+
+import numpy as np
+from ultralytics import YOLO
+
+
+class ModelValidationError(ValueError):
+    """Raised when a model cannot satisfy the locked V1 Pose contract."""
+
+
+@dataclass(frozen=True)
+class Keypoint:
+    x: float
+    y: float
+    confidence: float
+
+
+@dataclass(frozen=True)
+class PersonPose:
+    box_xyxy: Tuple[float, float, float, float]
+    box_confidence: float
+    keypoints: Sequence[Keypoint]
+
+
+ModelFactory = Callable[[str], Any]
+_SHA256 = re.compile(r"^[0-9a-fA-F]{64}$")
+
+
+def sha256_file(path: Path) -> str:
+    digest = hashlib.sha256()
+    with Path(path).open("rb") as stream:
+        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+class PoseAdapter:
+    """Load only the locked 17-keypoint person Pose model."""
+
+    def __init__(
+        self,
+        model_path: Path,
+        expected_sha256: str,
+        confidence_threshold: float = 0.25,
+        model_factory: Optional[ModelFactory] = None,
+    ) -> None:
+        self._model_path = Path(model_path)
+        self._confidence = float(confidence_threshold)
+        if not 0.0 <= self._confidence <= 1.0:
+            raise ModelValidationError("confidence threshold must be between 0 and 1")
+        self._validate_hash(expected_sha256)
+        factory = model_factory or YOLO
+        try:
+            self._model = factory(str(self._model_path))
+        except Exception as exc:
+            raise ModelValidationError("unable to load Pose model: {0}".format(exc))
+        self._person_class_ids = self._validate_model_contract(self._model)
+
+    @property
+    def model_path(self) -> Path:
+        return self._model_path
+
+    def infer(self, image: np.ndarray) -> Sequence[PersonPose]:
+        results = self._model(image, conf=self._confidence, verbose=False)
+        return self.from_results(results, self._person_class_ids)
+
+    @classmethod
+    def from_results(
+        cls, results: Any, person_class_ids: Optional[Sequence[int]] = None
+    ) -> Sequence[PersonPose]:
+        """Convert Ultralytics result objects into framework-independent poses."""
+
+        if not isinstance(results, (list, tuple)):
+            result_items = [results]
+        else:
+            result_items = results
+        poses: List[PersonPose] = []
+        for result in result_items:
+            boxes = getattr(result, "boxes", None)
+            keypoints = getattr(result, "keypoints", None)
+            if boxes is None or keypoints is None:
+                raise ModelValidationError("Pose result is missing boxes or keypoints")
+            names = getattr(result, "names", {})
+            allowed_ids = (
+                set(person_class_ids)
+                if person_class_ids is not None
+                else cls._person_ids_from_names(names)
+            )
+            box_rows = cls._rows(getattr(boxes, "xyxy", None), "boxes.xyxy")
+            confidence_rows = cls._rows(getattr(boxes, "conf", None), "boxes.conf")
+            class_rows = cls._rows(getattr(boxes, "cls", None), "boxes.cls")
+            keypoint_rows = cls._rows(getattr(keypoints, "data", None), "keypoints.data")
+            if not (
+                len(box_rows)
+                == len(confidence_rows)
+                == len(class_rows)
+                == len(keypoint_rows)
+            ):
+                raise ModelValidationError("Pose result arrays have inconsistent lengths")
+            for index, box in enumerate(box_rows):
+                class_id = int(class_rows[index])
+                if class_id not in allowed_ids:
+                    continue
+                if len(box) != 4:
+                    raise ModelValidationError("person box must contain four coordinates")
+                points = keypoint_rows[index]
+                if len(points) != 17:
+                    raise ModelValidationError("person Pose result must contain exactly 17 keypoints")
+                parsed_points = []
+                for point in points:
+                    if len(point) != 3:
+                        raise ModelValidationError("keypoint must contain x, y and confidence")
+                    parsed_points.append(
+                        Keypoint(float(point[0]), float(point[1]), float(point[2]))
+                    )
+                poses.append(
+                    PersonPose(
+                        box_xyxy=tuple(float(value) for value in box),
+                        box_confidence=float(confidence_rows[index]),
+                        keypoints=tuple(parsed_points),
+                    )
+                )
+        return tuple(poses)
+
+    def _validate_hash(self, expected_sha256: str) -> None:
+        if not _SHA256.match(expected_sha256 or ""):
+            raise ModelValidationError("expected SHA-256 must be a 64-character hex value")
+        if not self._model_path.is_file():
+            raise ModelValidationError("Pose model file does not exist: {0}".format(self._model_path))
+        actual_sha256 = sha256_file(self._model_path)
+        if actual_sha256.lower() != expected_sha256.lower():
+            raise ModelValidationError("Pose model SHA-256 does not match expected value")
+
+    @classmethod
+    def _validate_model_contract(cls, model: Any) -> Sequence[int]:
+        if getattr(model, "task", None) != "pose":
+            raise ModelValidationError("model task must be pose")
+        person_ids = cls._person_ids_from_names(getattr(model, "names", {}))
+        raw_shape = getattr(getattr(model, "model", None), "kpt_shape", None)
+        if raw_shape is None:
+            raw_shape = getattr(model, "kpt_shape", None)
+        if tuple(raw_shape or ()) != (17, 3):
+            raise ModelValidationError("model keypoint shape must be exactly (17, 3)")
+        return tuple(person_ids)
+
+    @staticmethod
+    def _person_ids_from_names(names: Any) -> Sequence[int]:
+        items = names.items() if isinstance(names, dict) else enumerate(names or ())
+        person_ids = [int(index) for index, name in items if str(name).lower() == "person"]
+        if not person_ids:
+            raise ModelValidationError("model must expose a person class")
+        return tuple(person_ids)
+
+    @staticmethod
+    def _rows(value: Any, field_name: str) -> List[Any]:
+        if value is None:
+            raise ModelValidationError("Pose result is missing {0}".format(field_name))
+        if hasattr(value, "cpu"):
+            value = value.cpu()
+        if hasattr(value, "numpy"):
+            value = value.numpy()
+        if isinstance(value, np.ndarray):
+            return value.tolist()
+        if hasattr(value, "tolist"):
+            return value.tolist()
+        return list(value)
@@ -0,0 +1,66 @@
+import hashlib
+
+import numpy as np
+import pytest
+
+from v1.pose import ModelValidationError, PoseAdapter
+
+
+class _FakeKeypoints:
+    def __init__(self, data):
+        self.data = data
+
+
+class _FakeBoxes:
+    def __init__(self):
+        self.xyxy = np.array([[10.0, 20.0, 50.0, 90.0], [1.0, 2.0, 3.0, 4.0]])
+        self.conf = np.array([0.9, 0.8])
+        self.cls = np.array([0, 1])
+
+
+class _FakeResult:
+    names = {0: "person", 1: "chair"}
+
+    def __init__(self):
+        person = [[float(index), float(index + 1), 0.9] for index in range(17)]
+        chair = [[float(index), float(index + 1), 0.8] for index in range(17)]
+        self.boxes = _FakeBoxes()
+        self.keypoints = _FakeKeypoints(np.array([person, chair]))
+
+
+def test_pose_adapter_rejects_hash_mismatch_before_loading_model(tmp_path):
+    model_path = tmp_path / "pose.pt"
+    model_path.write_bytes(b"not a real model")
+
+    with pytest.raises(ModelValidationError, match="SHA-256"):
+        PoseAdapter(model_path, expected_sha256="0" * 64)
+
+
+def test_from_results_extracts_person_box_confidence_and_seventeen_keypoints():
+    poses = PoseAdapter.from_results(_FakeResult())
+
+    assert len(poses) == 1
+    assert poses[0].box_xyxy == (10.0, 20.0, 50.0, 90.0)
+    assert poses[0].box_confidence == 0.9
+    assert len(poses[0].keypoints) == 17
+    assert poses[0].keypoints[5].confidence == 0.9
+
+
+def test_pose_adapter_rejects_non_pose_model_after_hash_validation(tmp_path):
+    model_path = tmp_path / "model.pt"
+    model_path.write_bytes(b"model bytes")
+    expected_sha256 = hashlib.sha256(model_path.read_bytes()).hexdigest()
+
+    class _WrongTaskModel:
+        task = "detect"
+        names = {0: "person"}
+
+        class model:
+            kpt_shape = (17, 3)
+
+    with pytest.raises(ModelValidationError, match="task"):
+        PoseAdapter(
+            model_path,
+            expected_sha256=expected_sha256,
+            model_factory=lambda _path: _WrongTaskModel(),
+        )