parking_solution/apps/trainer/trainer/cli.py

"""parking-trainer — inspect / train / evaluate / publish.

    parking-trainer inspect  --data /data
    parking-trainer train    --data /data --out /out [--mode features|finetune] [--min-accuracy 0.85]
    parking-trainer evaluate --model /out/<version>/bodytype.onnx --data /data
    parking-trainer publish  /out/<version> --url https://<gitea>/api/packages/<owner>/generic/parking-bodytype

Exit codes: 0 ok · 2 not enough labels · 3 trained but below the floor (report written, model
NOT written) · 1 anything else.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
import urllib.request
from datetime import datetime, timezone
from pathlib import Path

from .data import (
    VEHICLE_CLASSES,
    class_weights,
    load_labelled,
    load_reviewed_since,
    load_unlabelled,
    make_split,
    suggested_epochs,
    summarise,
)
from .preprocess import CROP_MARGIN, Sidecar
from .report import compute_metrics, render_report

MODEL_FILE = "bodytype.onnx"
SIDECAR_FILE = "bodytype.json"
REPORT_FILE = "report.md"
METRICS_FILE = "metrics.json"


def _log(msg: str) -> None:
    print(f"[trainer] {msg}", file=sys.stderr, flush=True)


def _now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


# ----------------------------------------------------------------------------------
# inspect
# ----------------------------------------------------------------------------------


def cmd_inspect(a: argparse.Namespace) -> int:
    samples, missing = load_labelled(a.data)
    split = make_split(samples, a.val_fraction, a.min_per_class)
    out = {
        "labelled": summarise(samples),
        "missingCrops": missing,
        "run": {
            "classes": list(split.classes),
            "train": split.counts("train"),
            "val": split.counts("val"),
            "dropped": split.dropped,
            "minPerClass": a.min_per_class,
            "valFraction": a.val_fraction,
        },
        "ready": len(split.classes) >= 2,
    }
    print(json.dumps(out, indent=2))
    return 0 if out["ready"] else 2


# ----------------------------------------------------------------------------------
# train
# ----------------------------------------------------------------------------------


def cmd_train(a: argparse.Namespace) -> int:
    try:
        from . import model as M
    except ImportError as exc:  # torch missing
        _log(f"the training stack is not installed ({exc}); install with: uv sync --extra train")
        return 1
    import numpy as np

    samples, missing = load_labelled(a.data)
    split = make_split(samples, a.val_fraction, a.min_per_class)
    if len(split.classes) < 2:
        _log(
            f"not enough labels: {len(samples)} usable, classes with >= {a.min_per_class}: "
            f"{list(split.classes)} (dropped {split.dropped}); nothing to train"
        )
        return 2
    version = a.version or datetime.now(timezone.utc).strftime("v%Y%m%d-%H%M")
    out_dir = a.out / version
    epochs = a.epochs or suggested_epochs(len(split.train), a.mode)
    weights = class_weights(split)
    idx = split.class_index
    _log(
        f"{version}: {a.mode} on {a.backbone}, classes {list(split.classes)}, "
        f"{len(split.train)} train / {len(split.val)} val, {epochs} epochs"
    )

    t0 = time.perf_counter()
    x_train, train = M.load_images(split.train, a.input_size)
    x_val, val = M.load_images(split.val, a.input_size)
    y_train = [idx[s.label] for s in train]
    y_val = [idx[s.label] for s in val]
    _log(f"decoded {len(train)} + {len(val)} crops in {time.perf_counter() - t0:.1f}s")
    if len(val) == 0 or len(set(y_train)) < 2:
        _log("not enough decodable crops on both sides of the split")
        return 2

    backbone = M.build_backbone(a.backbone, pretrained=not a.no_pretrained)
    cache = (
        None if a.no_cache else M.FeatureCache(a.out / "cache" / f"features-{a.backbone}-{a.input_size}.npz")
    )
    t0 = time.perf_counter()
    f_train = M.features_for(backbone, train, x_train, cache)
    _log(
        f"features for {len(train)} train crops in {time.perf_counter() - t0:.1f}s "
        f"(cache: {cache.path if cache else 'off'})"
    )
    head_epochs = epochs if a.mode == "features" else max(30, epochs * 5)
    head = M.train_head(f_train, y_train, len(split.classes), weights, head_epochs, seed=a.seed)
    net = M.Classifier.make(backbone, head)

    if a.mode == "finetune":
        t0 = time.perf_counter()
        net = M.finetune(
            net, x_train, y_train, weights, epochs, batch=a.batch, lr=a.lr, seed=a.seed, log=_log
        )
        _log(f"fine-tuned in {(time.perf_counter() - t0) / 60:.1f} min")

    logits = M.predict_logits(net, x_val)
    y_pred = logits.argmax(axis=1).tolist()
    metrics = compute_metrics(split.classes, y_val, y_pred, camera=[s.vision_class for s in val])

    # Export and check the graph gives the same answers as the torch model.
    out_dir.mkdir(parents=True, exist_ok=True)
    tmp_model = out_dir / (MODEL_FILE + ".tmp")
    M.export_onnx(net, a.input_size, tmp_model)
    from .infer import OnnxClassifier

    sidecar = Sidecar(
        version=version,
        classes=list(split.classes),
        input_size=a.input_size,
        backbone=a.backbone,
        mode=a.mode,
        trained_at=_now(),
        labels={"train": len(train), "val": len(val)},
    )
    tmp_side = out_dir / (SIDECAR_FILE + ".tmp")
    sidecar.write(tmp_side)
    onnx_pred = OnnxClassifier(tmp_model, tmp_side).predict_inputs(x_val.astype(np.float32)).argmax(axis=1)
    metrics.onnx_agreement = float((onnx_pred == np.array(y_pred)).mean()) if len(y_pred) else None
    sidecar.metrics = {
        "accuracy": metrics.accuracy,
        "macroRecall": metrics.macro_recall,
        "perClass": {c: m.__dict__ for c, m in metrics.per_class.items()},
        "floor": a.min_accuracy,
    }

    written = metrics.accuracy >= a.min_accuracy and (metrics.onnx_agreement or 0.0) >= 0.99
    notes = []
    if metrics.onnx_agreement is not None and metrics.onnx_agreement < 0.99:
        notes.append(
            f"ONNX export disagrees with the torch model ({metrics.onnx_agreement:.3f}); model withheld"
        )
    report = render_report(
        version=version,
        trained_at=sidecar.trained_at,
        mode=a.mode,
        backbone=a.backbone,
        epochs=epochs,
        classes=split.classes,
        train_counts=split.counts("train"),
        val_counts=split.counts("val"),
        dropped=split.dropped,
        missing_files=missing,
        weights=weights,
        metrics=metrics,
        min_accuracy=a.min_accuracy,
        written=written,
        notes=notes,
    )
    (out_dir / REPORT_FILE).write_text(report)
    (out_dir / METRICS_FILE).write_text(json.dumps(metrics.to_dict(), indent=2) + "\n")
    if written:
        sidecar.write(out_dir / SIDECAR_FILE)
        tmp_model.replace(out_dir / MODEL_FILE)
        tmp_side.unlink(missing_ok=True)
        _log(
            f"MODEL WRITTEN: {out_dir / MODEL_FILE} "
            f"(accuracy {metrics.accuracy:.3f} >= floor {a.min_accuracy})"
        )
    else:
        tmp_model.unlink(missing_ok=True)
        tmp_side.unlink(missing_ok=True)
        _log(
            f"MODEL NOT WRITTEN: accuracy {metrics.accuracy:.3f} < floor {a.min_accuracy}; "
            f"see {out_dir / REPORT_FILE}"
        )
    print(report)
    return 0 if written else 3


# ----------------------------------------------------------------------------------
# evaluate — an existing model against labels that arrived AFTER it was trained, and its
# view of the unlabelled pile (the ongoing accuracy check without labelling everything)
# ----------------------------------------------------------------------------------


def cmd_evaluate(a: argparse.Namespace) -> int:
    from .infer import OnnxClassifier

    clf = OnnxClassifier(a.model)
    since = a.since or clf.sidecar.trained_at
    classes = clf.classes
    result: dict[str, object] = {"model": clf.sidecar.version, "classes": classes, "since": since}

    reviewed = [s for s in load_reviewed_since(a.data, since) if s.label in classes]
    if reviewed:
        probs, kept = clf.predict_files([s.path for s in reviewed])
        rows = [reviewed[i] for i in kept]
        y_true = [classes.index(s.label) for s in rows]
        y_pred = probs.argmax(axis=1).tolist()
        m = compute_metrics(classes, y_true, y_pred, camera=[s.vision_class for s in rows])
        result["reviewedSince"] = m.to_dict()
    else:
        result["reviewedSince"] = None

    pending = load_unlabelled(a.data, a.limit)
    if pending:
        probs, kept = clf.predict_files([s.path for s in pending])
        rows = [pending[i] for i in kept]
        pred = probs.argmax(axis=1)
        conf = probs.max(axis=1)
        hist = {c: int((pred == i).sum()) for i, c in enumerate(classes)}
        result["unlabelled"] = {
            "n": len(rows),
            "predicted": hist,
            "meanConfidence": float(conf.mean()) if len(rows) else None,
            "belowHalf": int((conf < 0.5).sum()),
            "agreesWithDetector": float(
                sum(1 for p, s in zip(pred, rows, strict=True) if classes[int(p)] == s.vision_class)
                / len(rows)
            )
            if rows
            else None,
        }
    else:
        result["unlabelled"] = None
    print(json.dumps(result, indent=2))
    return 0


# ----------------------------------------------------------------------------------
# publish — the versioned files to a Gitea generic package (weights are not code, they
# do not live in git; the vision image fetches them by URL at build)
# ----------------------------------------------------------------------------------


def cmd_publish(a: argparse.Namespace) -> int:
    d: Path = a.dir
    files = [d / MODEL_FILE, d / SIDECAR_FILE, d / REPORT_FILE, d / METRICS_FILE]
    for f in files[:2]:
        if not f.exists():
            _log(f"{f} missing — nothing to publish (a run below the floor writes no model)")
            return 1
    version = Sidecar.read(d / SIDECAR_FILE).version
    if not a.url:
        _log("no publish url: pass --url or set TRAINER_PUBLISH_URL")
        return 1
    token = a.token or os.environ.get("TRAINER_PUBLISH_TOKEN", "")
    if not token:
        _log("no token: pass --token or set TRAINER_PUBLISH_TOKEN")
        return 1
    base = a.url.rstrip("/") + "/" + version
    for f in files:
        if not f.exists():
            continue
        req = urllib.request.Request(f"{base}/{f.name}", data=f.read_bytes(), method="PUT")
        req.add_header("Authorization", f"token {token}")
        req.add_header("Content-Type", "application/octet-stream")
        with urllib.request.urlopen(req, timeout=120) as r:
            _log(f"PUT {base}/{f.name} → {r.status}")
    _log(f"published {version}; pin it in apps/vision/models/bodytype.version and rebuild the vision image")
    return 0


def cmd_serve(a: argparse.Namespace) -> int:
    from .server import serve

    serve(
        a.data,
        a.out,
        a.host,
        a.port,
        os.environ.get("TRAINER_PUBLISH_URL", ""),
        os.environ.get("TRAINER_PUBLISH_TOKEN", ""),
    )
    return 0


# ----------------------------------------------------------------------------------


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="parking-trainer", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    sub = p.add_subparsers(dest="cmd", required=True)

    def data_args(sp: argparse.ArgumentParser) -> None:
        sp.add_argument(
            "--data",
            type=Path,
            default=Path(os.environ.get("TRAINER_DATA_DIR", "/data")),
            help="collector volume (collector.sqlite + crops/)",
        )

    def split_args(sp: argparse.ArgumentParser) -> None:
        sp.add_argument(
            "--min-per-class",
            type=int,
            default=20,
            help="classes with fewer reviewed crops are dropped from the run",
        )
        sp.add_argument(
            "--val-fraction", type=float, default=0.2, help="newest fraction held out for validation"
        )

    i = sub.add_parser("inspect", help="what a run would train on")
    data_args(i)
    split_args(i)
    i.set_defaults(fn=cmd_inspect)

    t = sub.add_parser("train", help="train, evaluate, export (or refuse)")
    data_args(t)
    split_args(t)
    t.add_argument(
        "--out",
        type=Path,
        default=Path(os.environ.get("TRAINER_OUT_DIR", "/out")),
        help="output root; a <version>/ folder is created under it",
    )
    t.add_argument("--mode", choices=["features", "finetune"], default="features")
    t.add_argument(
        "--backbone", choices=["resnet18", "mobilenet_v3_small", "efficientnet_b0"], default="resnet18"
    )
    t.add_argument("--epochs", type=int, default=0, help="0 = pick from the data size")
    t.add_argument("--batch", type=int, default=32)
    t.add_argument("--lr", type=float, default=1e-4, help="fine-tune learning rate")
    t.add_argument("--input-size", type=int, default=224)
    t.add_argument(
        "--min-accuracy", type=float, default=0.85, help="validation floor below which NO model is written"
    )
    t.add_argument("--version", default="", help="model version (default v<date>-<time>)")
    t.add_argument("--seed", type=int, default=7)
    t.add_argument(
        "--no-pretrained", action="store_true", help="random init (tests only — never for a real run)"
    )
    t.add_argument("--no-cache", action="store_true", help="do not read/write the feature cache")
    t.set_defaults(fn=cmd_train)

    e = sub.add_parser(
        "evaluate", help="an existing model vs labels reviewed after it was trained + the unlabelled pile"
    )
    data_args(e)
    e.add_argument(
        "--model", type=Path, required=True, help="path to bodytype.onnx (sidecar .json beside it)"
    )
    e.add_argument("--since", default="", help="ISO time; default = the model's trained_at")
    e.add_argument(
        "--limit", type=int, default=2000, help="how many unlabelled crops to score (newest first)"
    )
    e.set_defaults(fn=cmd_evaluate)

    u = sub.add_parser("publish", help="PUT a version folder to a Gitea generic package")
    u.add_argument("dir", type=Path, help="the <version>/ folder a passing run wrote")
    u.add_argument(
        "--url",
        default=os.environ.get("TRAINER_PUBLISH_URL", ""),
        help="https://<gitea>/api/packages/<owner>/generic/parking-bodytype (or TRAINER_PUBLISH_URL)",
    )
    u.add_argument("--token", default="", help="Gitea token with package:write (or TRAINER_PUBLISH_TOKEN)")
    u.set_defaults(fn=cmd_publish)

    s = sub.add_parser("serve", help="the job API the collector's Training section talks to")
    data_args(s)
    s.add_argument("--out", type=Path, default=Path(os.environ.get("TRAINER_OUT_DIR", "/out")))
    s.add_argument("--host", default=os.environ.get("TRAINER_HOST", "0.0.0.0"))
    s.add_argument("--port", type=int, default=int(os.environ.get("TRAINER_PORT", "8091")))
    s.set_defaults(fn=cmd_serve)
    return p


def main(argv: list[str] | None = None) -> int:
    a = build_parser().parse_args(argv)
    try:
        return int(a.fn(a))
    except FileNotFoundError as exc:
        _log(str(exc))
        return 1


__all__ = ["main", "build_parser", "VEHICLE_CLASSES", "CROP_MARGIN"]

if __name__ == "__main__":
    sys.exit(main())