silver_pose/v1/tests/test_pose.py

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


class _RecordingPoseModel:
    task = "pose"
    names = {0: "person"}

    class model:
        kpt_shape = (17, 3)

    def __init__(self):
        self.calls = []

    def __call__(self, image, conf, verbose):
        self.calls.append(conf)

        class _Empty:
            names = {0: "person"}

            class boxes:
                xyxy = []
                conf = []
                cls = []

            class keypoints:
                data = []

        return _Empty()


def test_set_confidence_threshold_is_forwarded_to_inference(tmp_path):
    weights = tmp_path / "pose.pt"
    weights.write_bytes(b"fake-weights")
    expected_sha256 = hashlib.sha256(weights.read_bytes()).hexdigest()
    model = _RecordingPoseModel()
    adapter = PoseAdapter(
        weights, expected_sha256, confidence_threshold=0.25, model_factory=lambda _p: model
    )

    adapter.set_confidence_threshold(0.6)
    adapter.infer(np.zeros((10, 10, 3), dtype=np.uint8))

    assert adapter.confidence_threshold == 0.6
    assert model.calls == [0.6]


def test_set_confidence_threshold_rejects_out_of_range(tmp_path):
    weights = tmp_path / "pose.pt"
    weights.write_bytes(b"fake-weights")
    expected_sha256 = hashlib.sha256(weights.read_bytes()).hexdigest()
    adapter = PoseAdapter(
        weights, expected_sha256, model_factory=lambda _p: _RecordingPoseModel()
    )

    with pytest.raises(ModelValidationError):
        adapter.set_confidence_threshold(1.5)


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(),
        )