tests/test_smoke.py

"""Run flow/inference.py on the sample clip and compare with examples/expected_flow_output.json.

Run from the repository root: pytest tests/test_smoke.py
Needs the packages in requirements.txt and pytest. The first run downloads RAFT-small from torchvision.
The test runs on CPU. Probabilities may differ by 0.02 and yaw by 0.5 degrees across machines.
"""

import json
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent


def test_flow_output_matches_expected() -> None:
    cmd = [
        sys.executable,
        str(ROOT / "flow" / "inference.py"),
        str(ROOT / "examples" / "bristol-walk-10s.mp4"),
        "--config", str(ROOT / "flow" / "config_walking.json"),
        "--stride", "49",
        "--device", "cpu",
    ]
    out = json.loads(subprocess.run(cmd, check=True, capture_output=True, text=True).stdout)
    expected = json.loads((ROOT / "examples" / "expected_flow_output.json").read_text())
    assert len(out["windows"]) == len(expected["windows"])
    for got, want in zip(out["windows"], expected["windows"]):
        assert got["start_frame"] == want["start_frame"]
        assert abs(got["yaw_deg"] - want["yaw_deg"]) < 0.5
        assert abs(got["pitch_deg"] - want["pitch_deg"]) < 0.5
        for key, value in want["key_probabilities"].items():
            assert abs(got["key_probabilities"][key] - value) < 0.02