"""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 @property def confidence_threshold(self) -> float: return self._confidence def set_confidence_threshold(self, value: float) -> None: """Update the inference confidence so settings changes take effect.""" confidence = float(value) if not 0.0 <= confidence <= 1.0: raise ModelValidationError("confidence threshold must be between 0 and 1") self._confidence = confidence 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)