NODEDC_MISSION_CORE/tests/test_e5_instance_tracking.py

from __future__ import annotations

import importlib.util
import json
import math
import sys
from pathlib import Path
from types import ModuleType

import numpy as np
import pytest


def _worker_module() -> ModuleType:
    worker_root = Path(__file__).parents[1] / "experiments" / "perception" / "worker"
    path = worker_root / "run_e5_instance_tracking.py"
    sys.path.insert(0, str(worker_root))
    try:
        spec = importlib.util.spec_from_file_location("e5_instance_tracking_worker", path)
        assert spec is not None and spec.loader is not None
        module = importlib.util.module_from_spec(spec)
        sys.modules[spec.name] = module
        spec.loader.exec_module(module)
        return module
    finally:
        sys.path.pop(0)


def _profile_path() -> Path:
    return (
        Path(__file__).parents[1]
        / "experiments"
        / "perception"
        / "worker"
        / "e5_yolox_bytetrack_profile.json"
    )


def _detection(*, score: float, x: float = 100, class_id: int = 0) -> dict[str, object]:
    return {
        "class_id": class_id,
        "label": "person" if class_id == 0 else "car",
        "score": score,
        "bbox_xyxy": [x, 100.0, x + 60.0, 220.0],
        "valid_fov_fraction": 1.0,
    }


def test_profile_is_pinned_and_semantically_valid() -> None:
    worker = _worker_module()

    profile, digest = worker._read_profile(_profile_path())

    assert len(digest) == 64
    assert profile["model"]["id"] == "yolox_s"
    assert profile["model"]["model_sha256"] == (
        "c5c2d13e59ae883e6af3b45daea64af4833a4951c92d116ec270d9ddbe998063"
    )
    assert profile["tracking"]["algorithm"] == "bytetrack-style-two-stage-iou/v1"
    assert profile["detection"]["target_class_ids"] == [0, 1, 2, 3, 5, 7]


def test_profile_rejects_duplicate_target_classes(tmp_path: Path) -> None:
    worker = _worker_module()
    profile = json.loads(_profile_path().read_text(encoding="utf-8"))
    profile["detection"]["target_class_ids"] = [0, 0]
    path = tmp_path / "profile.json"
    path.write_text(json.dumps(profile), encoding="utf-8")

    with pytest.raises(RuntimeError, match="target class"):
        worker._read_profile(path)


def test_valid_fraction_uses_box_area_and_center_gate() -> None:
    worker = _worker_module()
    mask = np.zeros((10, 10), dtype=bool)
    mask[2:8, 2:8] = True
    integral = np.pad(mask.astype(np.int64), ((1, 0), (1, 0))).cumsum(0).cumsum(1)

    fraction, center_inside, area = worker._valid_fraction(
        np.asarray([1.0, 1.0, 5.0, 5.0]), integral
    )

    assert area == 16
    assert fraction == pytest.approx(9 / 16)
    assert center_inside is True


def test_yolox_decode_and_fov_filter_admit_one_person() -> None:
    worker = _worker_module()
    profile = json.loads(_profile_path().read_text(encoding="utf-8"))
    output = np.zeros((1, 8400, 85), dtype=np.float32)
    output[0, 0, :4] = [40.0, 30.0, math.log(10.0), math.log(10.0)]
    output[0, 0, 4] = 0.9
    output[0, 0, 5] = 0.9

    detections, rejected = worker._detections(
        output,
        profile,
        np.ones((600, 800), dtype=bool),
    )

    assert rejected == {}
    assert len(detections) == 1
    assert detections[0]["label"] == "person"
    assert detections[0]["score"] == pytest.approx(0.81)
    assert detections[0]["bbox_xyxy"] == pytest.approx([350.0, 250.0, 450.0, 350.0])


def test_nms_suppresses_lower_score_box_contained_inside_object() -> None:
    worker = _worker_module()
    boxes = np.asarray(
        [
            [0.0, 0.0, 100.0, 100.0],
            [10.0, 10.0, 60.0, 60.0],
        ]
    )
    scores = np.asarray([0.9, 0.5])

    assert worker._nms(boxes, scores, 0.45, 0.8) == [0]


def test_two_stage_tracker_keeps_id_through_low_score_detection() -> None:
    worker = _worker_module()
    profile = json.loads(_profile_path().read_text(encoding="utf-8"))
    tracker = worker.TwoStageTracker(profile["tracking"])

    assert tracker.update([_detection(score=0.8)], 0) == []
    second = tracker.update([_detection(score=0.7, x=103)], 1)
    third = tracker.update([_detection(score=0.15, x=106)], 2)

    assert [track.track_id for track in second] == [1]
    assert [track.track_id for track in third] == [1]
    assert third[0].hits == 3
    assert tracker.created == 1


def test_tracker_does_not_cross_match_classes() -> None:
    worker = _worker_module()
    profile = json.loads(_profile_path().read_text(encoding="utf-8"))
    tracker = worker.TwoStageTracker(profile["tracking"])

    tracker.update([_detection(score=0.8)], 0)
    tracks = tracker.update([_detection(score=0.8, class_id=2)], 1)

    assert tracks == []
    assert tracker.created == 2
    assert {track.class_id for track in tracker.tracks} == {0, 2}


def test_clip_timeline_may_start_at_nonzero_source_index(tmp_path: Path) -> None:
    worker = _worker_module()
    path = tmp_path / "timeline.jsonl"
    rows = [
        {"frame_index": 0, "source_frame_index": 1000, "session_seconds": 135.1},
        {"frame_index": 1, "source_frame_index": 1001, "session_seconds": 135.2},
    ]
    path.write_text("\n".join(json.dumps(row) for row in rows) + "\n", encoding="utf-8")

    assert worker._read_timeline(path, 2) == rows