silver_pose/v1/pose.py

"""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)