parking_solution/apps/trainer/tests/test_train.py

"""The training job end to end on the synthetic volume — needs the `train` extra (torch);
skipped where it is not installed (CI syncs without it, like the vision service)."""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pytest

torch = pytest.importorskip("torch")

from trainer.cli import main  # noqa: E402
from trainer.infer import OnnxClassifier  # noqa: E402
from trainer.preprocess import Sidecar  # noqa: E402

COMMON = ["--no-pretrained", "--input-size", "64", "--no-cache", "--seed", "3"]


def test_features_run_writes_model_sidecar_report_and_evaluates(
    collector_dir: Path, tmp_path: Path, capsys
) -> None:  # type: ignore[no-untyped-def]
    out = tmp_path / "out"
    rc = main(
        [
            "train",
            "--data",
            str(collector_dir),
            "--out",
            str(out),
            "--version",
            "vtest",
            "--mode",
            "features",
            "--epochs",
            "150",
            "--min-accuracy",
            "0.0",
            *COMMON,
        ]
    )
    assert rc == 0
    d = out / "vtest"
    assert {p.name for p in d.iterdir()} == {"bodytype.onnx", "bodytype.json", "report.md", "metrics.json"}
    side = Sidecar.read(d / "bodytype.json")
    assert side.classes == ["sedan", "suv", "van"] and side.input_size == 64 and side.mode == "features"
    assert side.labels == {"train": 96, "val": 24} and side.metrics["floor"] == 0.0
    metrics = json.loads((d / "metrics.json").read_text())
    assert metrics["n"] == 24 and metrics["onnx_agreement"] == 1.0
    # Colour-coded classes: even a random backbone's pooled features separate them.
    assert metrics["accuracy"] >= 0.9
    report = (d / "report.md").read_text()
    assert (
        "MODEL WRITTEN" in report
        and "truck (5)" in report
        and "crop is missing on disk (skipped): 1" in report
    )

    # The exported graph takes raw 0–255 RGB and answers by itself.
    clf = OnnxClassifier(d / "bodytype.onnx")
    probs, kept = clf.predict_files(
        [s for s in sorted((collector_dir / "crops" / "booth-2").glob("*.jpg"))][:6]
    )
    assert probs.shape == (6, 3) and kept == [0, 1, 2, 3, 4, 5]
    assert np.allclose(probs.sum(axis=1), 1.0, atol=1e-4)

    # evaluate: labels reviewed after training (none — the fixture's reviews predate it) and the
    # unlabelled pile (6 entry samples).
    capsys.readouterr()
    assert main(["evaluate", "--data", str(collector_dir), "--model", str(d / "bodytype.onnx")]) == 0
    res = json.loads(capsys.readouterr().out)
    assert res["model"] == "vtest" and res["reviewedSince"] is None
    assert res["unlabelled"]["n"] == 6 and sum(res["unlabelled"]["predicted"].values()) == 6
    assert (
        main(
            [
                "evaluate",
                "--data",
                str(collector_dir),
                "--model",
                str(d / "bodytype.onnx"),
                "--since",
                "2026-09-01T00:00:00Z",
            ]
        )
        == 0
    )
    res2 = json.loads(capsys.readouterr().out)
    assert res2["reviewedSince"]["n"] == 125 - 5  # trucks are not a class the model knows


def test_below_the_floor_writes_the_report_but_no_model(collector_dir: Path, tmp_path: Path) -> None:
    out = tmp_path / "out"
    rc = main(
        [
            "train",
            "--data",
            str(collector_dir),
            "--out",
            str(out),
            "--version",
            "vlow",
            "--mode",
            "features",
            "--epochs",
            "5",
            "--min-accuracy",
            "1.01",
            *COMMON,
        ]
    )
    assert rc == 3
    d = out / "vlow"
    assert {p.name for p in d.iterdir()} == {"report.md", "metrics.json"}
    assert "MODEL NOT WRITTEN" in (d / "report.md").read_text()


def test_not_enough_labels_is_exit_2(collector_dir: Path, tmp_path: Path) -> None:
    out = tmp_path / "out"
    rc = main(["train", "--data", str(collector_dir), "--out", str(out), "--min-per-class", "100", *COMMON])
    assert rc == 2
    assert not out.exists()


def test_finetune_runs_and_uses_the_feature_cache(collector_dir: Path, tmp_path: Path) -> None:
    out = tmp_path / "out"
    args = [
        "train",
        "--data",
        str(collector_dir),
        "--out",
        str(out),
        "--mode",
        "finetune",
        "--backbone",
        "mobilenet_v3_small",
        "--epochs",
        "1",
        "--batch",
        "16",
        "--min-accuracy",
        "0.0",
        "--no-pretrained",
        "--input-size",
        "64",
        "--seed",
        "3",
    ]
    assert main([*args, "--version", "vft"]) == 0
    cache = out / "cache" / "features-mobilenet_v3_small-64.npz"
    assert cache.exists()
    z = np.load(cache)
    assert len(z["ids"]) == 96 and z["feats"].shape == (96, 576)
    assert Sidecar.read(out / "vft" / "bodytype.json").mode == "finetune"