67 lines
1.9 KiB
Python
67 lines
1.9 KiB
Python
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(),
|
|
)
|