173 lines
6.7 KiB
Python
173 lines
6.7 KiB
Python
"""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)
|