pose_tracking_exp/tests/test_tracker_single_person.py

from pathlib import Path

import numpy as np
import pytest

pytest.importorskip("rpt")

from pose_tracking_exp.common.joints import BODY20_INDEX_BY_NAME
from pose_tracking_exp.schema import CameraCalibration, CameraFrame, FrameBundle, ProposalCluster, SceneConfig, TrackerConfig
from pose_tracking_exp.tracking import PoseTracker


def _make_scene() -> SceneConfig:
    cameras = (
        CameraCalibration(
            name="cam0",
            width=640,
            height=480,
            K=np.asarray([[500.0, 0.0, 320.0], [0.0, 500.0, 240.0], [0.0, 0.0, 1.0]], dtype=np.float64),
            DC=np.zeros(5, dtype=np.float64),
            R=np.eye(3, dtype=np.float64),
            T=np.zeros(3, dtype=np.float64),
        ),
        CameraCalibration(
            name="cam1",
            width=640,
            height=480,
            K=np.asarray([[500.0, 0.0, 320.0], [0.0, 500.0, 240.0], [0.0, 0.0, 1.0]], dtype=np.float64),
            DC=np.zeros(5, dtype=np.float64),
            R=np.eye(3, dtype=np.float64),
            T=np.asarray([1.0, 0.0, 0.0], dtype=np.float64),
        ),
    )
    return SceneConfig(
        room_size=np.asarray([6.0, 4.0, 3.0], dtype=np.float64),
        room_center=np.asarray([0.0, 0.0, 1.0], dtype=np.float64),
        cameras=cameras,
    )


def _make_bundle(bundle_index: int) -> FrameBundle:
    views = tuple(
        CameraFrame(
            camera_name=camera_name,
            frame_index=bundle_index,
            timestamp_unix_ns=1_000_000_000 + bundle_index * 33_000_000,
            detections=(),
            source_size=(640, 480),
        )
        for camera_name in ("cam0", "cam1")
    )
    return FrameBundle(
        bundle_index=bundle_index,
        timestamp_unix_ns=views[0].timestamp_unix_ns,
        views=views,
    )


def _make_proposal(root_x: float, *, score: float = 1.0) -> ProposalCluster:
    pose = np.zeros((20, 4), dtype=np.float64)
    joint_positions = {
        "hip_middle": [root_x, 1.0, 3.0],
        "hip_left": [root_x + 0.12, 1.0, 3.0],
        "hip_right": [root_x - 0.12, 1.0, 3.0],
        "shoulder_middle": [root_x, 1.52, 3.0],
        "shoulder_left": [root_x + 0.18, 1.52, 3.0],
        "shoulder_right": [root_x - 0.18, 1.52, 3.0],
        "elbow_left": [root_x + 0.42, 1.48, 3.02],
        "elbow_right": [root_x - 0.42, 1.48, 3.02],
        "wrist_left": [root_x + 0.64, 1.45, 3.04],
        "wrist_right": [root_x - 0.64, 1.45, 3.04],
        "knee_left": [root_x + 0.1, 0.58, 3.0],
        "knee_right": [root_x - 0.1, 0.58, 3.0],
        "ankle_left": [root_x + 0.1, 0.15, 3.02],
        "ankle_right": [root_x - 0.1, 0.15, 3.02],
        "head": [root_x, 1.82, 3.02],
        "nose": [root_x, 1.8, 3.06],
        "eye_left": [root_x + 0.03, 1.81, 3.05],
        "eye_right": [root_x - 0.03, 1.81, 3.05],
        "ear_left": [root_x + 0.06, 1.81, 3.02],
        "ear_right": [root_x - 0.06, 1.81, 3.02],
    }
    for name, position in joint_positions.items():
        pose[BODY20_INDEX_BY_NAME[name], :3] = position
        pose[BODY20_INDEX_BY_NAME[name], 3] = score
    return ProposalCluster(
        pose3d=pose,
        root=np.asarray([root_x, 1.0, 3.0], dtype=np.float64),
        source_views=frozenset({"cam0", "cam1"}),
        support_size=2,
        mean_score=score,
    )


def test_single_person_mode_caps_active_tracks(monkeypatch) -> None:
    tracker = PoseTracker(
        _make_scene(),
        TrackerConfig(
            max_active_tracks=1,
            tentative_min_age=1,
            tentative_hits_required=1,
            tentative_promote_score=0.0,
            active_miss_to_lost=5,
            proposal_min_score=0.5,
        ),
    )
    proposals_by_bundle = {
        0: (_make_proposal(0.0, score=0.95), _make_proposal(0.9, score=0.7)),
        1: (_make_proposal(0.05, score=0.96), _make_proposal(0.85, score=0.75)),
    }

    monkeypatch.setattr(
        tracker,
        "_build_proposals",
        lambda bundle, unmatched: proposals_by_bundle[bundle.bundle_index],
    )

    results = tracker.run([_make_bundle(0), _make_bundle(1)])

    assert len(results[0].active_tracks) == 1
    assert len(results[1].active_tracks) == 1
    assert not results[1].tentative_tracks
    assert [track.track_id for track in results[1].active_tracks] == [1]


def test_single_person_mode_reuses_lost_track_id(monkeypatch) -> None:
    tracker = PoseTracker(
        _make_scene(),
        TrackerConfig(
            max_active_tracks=1,
            tentative_min_age=1,
            tentative_hits_required=1,
            tentative_promote_score=0.0,
            active_miss_to_lost=1,
            lost_delete_age=10,
            proposal_min_score=0.5,
        ),
    )
    proposals_by_bundle = {
        0: (_make_proposal(0.0, score=0.95),),
        1: (),
        2: (_make_proposal(0.05, score=0.96),),
    }

    monkeypatch.setattr(
        tracker,
        "_build_proposals",
        lambda bundle, unmatched: proposals_by_bundle[bundle.bundle_index],
    )

    results = tracker.run([_make_bundle(0), _make_bundle(1), _make_bundle(2)])

    assert [track.track_id for track in results[0].active_tracks] == [1]
    assert [track.track_id for track in results[1].lost_tracks] == [1]
    assert [track.track_id for track in results[2].active_tracks] == [1]
    assert tracker.diagnostics_snapshot().reacquisitions >= 1