code/scripts/ref_golden.py

#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Goldens of the fp32 CPU reference (``tt_sceneseg.reference``) for the TT port's tests. CPU only (no device).

    OMP_NUM_THREADS=4 tools/research-venv/bin/python -I bundles/sceneseg-p150/code/scripts/ref_golden.py

Writes, for both channel orders (``bgr`` deployed, ``rgb`` training order):

- SMALL, shipped with the bundle (``code/tt_sceneseg/tests/goldens/``): per shipped sample (the 9 tutorial images)
  ``<stem>.npz`` = the 320x640 class map, the published mask at the source size and the 3 logits at 4096 fixed pixels
  (``sample_yx``), and ``<stem>.json`` = statistics (class fractions, decision margins; every tap's shape / mean / std
  / absmax / sha256 for ``highway_normal_1``).
- the container smoke's stored references, ``code/tt_sceneseg/samples/highway_normal_1.<order>.reference.json``
  (``ReferencePipeline.reference_body``: the ``/predict`` body without timings).
- LARGE, never shipped (``research/sceneseg/goldens/``): ``taps/<id>.<order>.npz`` = the module taps of
  ``config.MODULE_TAPS`` (fp32; what the TT module gates compare and teacher-force) for the gate samples, and
  ``public/<id>.npz`` = logits (fp32) and class maps of both orders for every converted public frame of
  research/sceneseg/public_data (PandaSet, nuScenes, comma10k, tutorial), plus ``index.json`` (sizes, sha256, the
  agreement with the research ORT goldens, versions).

The inputs are built with the node's exact pre-processing (``host.preprocess``); public frames start from the
converter's lossless decode (``inputs/<id>/source_bgr.png``). Nothing downloaded is executed (``-I``; weights are
read as data by ``reference.weights``).
"""
from __future__ import annotations

import argparse
import datetime as dt
import hashlib
import json
import sys
import time
from pathlib import Path

BUNDLE = Path(__file__).resolve().parents[2]
WS = BUNDLE.parents[1]
sys.path.insert(0, str(BUNDLE / "code"))

import numpy as np  # noqa: E402

from tt_sceneseg.reference import assets, config  # noqa: E402

PKG = BUNDLE / "code" / "tt_sceneseg"
SMALL = PKG / "tests" / "goldens"
SAMPLES = PKG / "samples"
LARGE = WS / "research" / "sceneseg" / "goldens"
PUBLIC = WS / "research" / "sceneseg" / "public_data"
TUTORIAL = ("highway_normal_1", "highway_normal_2", "highway_normal_3", "highway_normal_4", "highway_normal_5",
            "highway_rain_1", "highway_rain_2", "highway_rain_3", "highway_snow")
SMOKE_SAMPLE = "highway_normal_1"
TAP_SAMPLES = (("sample_highway_normal_1", None), ("ps019_f40_front_camera", ("bgr",)))   # (id, orders or None=all)
N_SAMPLES = 4096


def sha256_bytes(b: bytes) -> str:
    return hashlib.sha256(b).hexdigest()


def sha256_file(p: Path) -> str:
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for block in iter(lambda: f.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()


def sha256_array(a: np.ndarray) -> str:
    a = np.ascontiguousarray(a)
    return sha256_bytes(f"{a.dtype.str}{a.shape}".encode() + a.tobytes())


def stats(a: np.ndarray) -> dict:
    a64 = np.asarray(a, np.float64)
    return {"shape": list(a.shape), "mean": float(a64.mean()), "std": float(a64.std()),
            "absmax": float(np.abs(a64).max()), "sha256": sha256_array(np.asarray(a, np.float32))}


def margin_stats(logits: np.ndarray) -> dict:
    """How many pixels a small logit change can flip (the pixels a bf16 port is likely to change)."""
    lg = np.sort(logits[0].astype(np.float64), axis=0)
    gap = lg[-1] - lg[-2]
    cls = np.argmax(logits[0], axis=0)
    out = {f"frac_within_{t}": float((gap < t).mean()) for t in (0.05, 0.25, 1.0)}
    fg = cls == config.FOREGROUND_CLASS
    out["fg_frac_within_0.25"] = float((gap[fg] < 0.25).mean()) if fg.any() else 0.0
    return out


def class_fractions(cls: np.ndarray) -> dict:
    return {config.CLASS_NAMES[c]: float((cls == c).mean()) for c in range(config.NUM_CLASSES)}


def load_bgr_png(path: Path) -> np.ndarray:
    from PIL import Image

    with Image.open(path) as im:
        return np.ascontiguousarray(np.asarray(im.convert("RGB"))[:, :, ::-1])


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--threads", type=int, default=4)
    ap.add_argument("--skip-public", action="store_true", help="no research/sceneseg/goldens/public")
    ap.add_argument("--skip-taps", action="store_true", help="no research/sceneseg/goldens/taps")
    ap.add_argument("--only", choices=("comma10k",), help="only this part (the index is updated in place)")
    a = ap.parse_args()
    import torch

    torch.set_num_threads(a.threads)
    from tt_sceneseg.host.postprocess import postprocess
    from tt_sceneseg.host.preprocess import prepare_bgr, prepare_image
    from tt_sceneseg.reference.model import SceneSegReference
    from tt_sceneseg.reference.pipeline import ReferencePipeline
    from tt_sceneseg.ttaw.golden import TapRegistry
    from tt_sceneseg.ttaw.metrics import pcc

    onnx = assets.onnx_path()
    t0 = time.perf_counter()
    ref = SceneSegReference.from_onnx(onnx)
    versions = {"torch": torch.__version__, "numpy": np.__version__, "python": sys.version.split()[0]}
    common = {"onnx": config.ONNX_FILE, "onnx_sha256": config.ONNX_SHA256, "versions": versions,
              "written": dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"),
              "writer": "code/scripts/ref_golden.py (fp32 CPU reference, tt_sceneseg.reference)"}
    print(f"reference loaded in {time.perf_counter() - t0:.1f} s; torch {torch.__version__}, {a.threads} threads",
          flush=True)
    rng = np.random.default_rng(0)
    sample_yx = np.stack([rng.integers(0, config.INPUT_HW[0], N_SAMPLES),
                          rng.integers(0, config.INPUT_HW[1], N_SAMPLES)])
    SMALL.mkdir(parents=True, exist_ok=True)
    LARGE.mkdir(parents=True, exist_ok=True)
    index = {"description": "Large goldens of the fp32 CPU reference of sceneseg-p150 (never shipped)", **common,
             "files": {}, "public": {}}
    if a.only:
        if (LARGE / "index.json").is_file():
            index = json.loads((LARGE / "index.json").read_text())
        comma10k_goldens(ref, index, common)
        (LARGE / "index.json").write_text(json.dumps(index, indent=1) + "\n")
        print(f"done in {time.perf_counter() - t0:.0f} s", flush=True)
        return 0
    logits_cache = {}

    # ---- the shipped samples: small goldens, smoke references, module taps of the smoke sample ----------------
    for stem in TUTORIAL:
        path = SAMPLES / f"{stem}.png"
        arrays = {"sample_yx": sample_yx.astype(np.int32)}
        meta = {**common, "sample": f"samples/{stem}.png", "sample_sha256": sha256_file(path), "variants": {}}
        for order in config.CHANNEL_ORDERS:
            t1 = time.perf_counter()
            prep = prepare_image(str(path), channel_order=order)
            want_taps = stem == SMOKE_SAMPLE and not a.skip_taps
            taps = TapRegistry(exclude=config.LARGE_TAPS) if stem == SMOKE_SAMPLE else None
            logits = ref.run_numpy(prep.network_input(), taps=taps)["logits"]
            res = postprocess(prep.src_hw, logits=logits)
            logits_cache[(stem, order)] = logits
            arrays[f"{order}.class_map"] = res.class_map
            arrays[f"{order}.mask"] = res.mask
            arrays[f"{order}.logits_samples"] = logits[0][:, sample_yx[0], sample_yx[1]].astype(np.float32)
            v = {"src_hw": list(prep.src_hw), "input_u8_sha256": sha256_array(prep.u8_bgr),
                 "network_input_sha256": sha256_array(prep.network_input()), "logits": stats(logits),
                 "class_fraction_320x640": class_fractions(res.class_map),
                 "foreground_fraction_src": res.foreground_fraction, "margins": margin_stats(logits),
                 "seconds": round(time.perf_counter() - t1, 1)}
            if taps is not None:
                v["taps"] = {name: stats(val) for name, val in taps.to_dict().items()}
                if want_taps:
                    out = LARGE / "taps" / f"sample_{stem}.{order}.npz"
                    out.parent.mkdir(parents=True, exist_ok=True)
                    np.savez(out, **{k: taps[k] for k in config.MODULE_TAPS})
                    index["files"][str(out.relative_to(LARGE))] = {"bytes": out.stat().st_size,
                                                                   "sha256": sha256_file(out),
                                                                   "taps": list(config.MODULE_TAPS)}
            meta["variants"][order] = v
            print(f"{stem} {order}: fg {res.foreground_fraction:.4f} {v['seconds']} s", flush=True)
        np.savez_compressed(SMALL / f"{stem}.npz", **arrays)
        (SMALL / f"{stem}.json").write_text(json.dumps(meta, indent=1) + "\n")

    # the container smoke's stored references (the /predict body of the CPU reference)
    for order in config.CHANNEL_ORDERS:
        pipe = ReferencePipeline(onnx, order, threads=a.threads, model=ref)
        body = pipe.reference_body(str(SAMPLES / f"{SMOKE_SAMPLE}.png"))
        (SAMPLES / f"{SMOKE_SAMPLE}.{order}.reference.json").write_text(json.dumps(body, indent=1) + "\n")
        print(f"reference body {order}: {body['class_counts']}", flush=True)

    # ---- the public frames: logits + class maps of both orders (large), module taps of the gate frames --------
    if not a.skip_public and (PUBLIC / "inputs").is_dir():
        pub = LARGE / "public"
        pub.mkdir(parents=True, exist_ok=True)
        tap_orders = dict(TAP_SAMPLES)
        for d in sorted(p for p in (PUBLIC / "inputs").iterdir() if (p / "source_bgr.png").is_file()):
            fid = d.name
            bgr = load_bgr_png(d / "source_bgr.png")
            arrays, rec = {}, {"orders": {}}
            for order in config.CHANNEL_ORDERS:
                t1 = time.perf_counter()
                prep = prepare_bgr(bgr, channel_order=order)
                stem = fid[4:] if fid.startswith("tut_") else None
                want_taps = (not a.skip_taps and fid in tap_orders and (tap_orders[fid] is None
                                                                        or order in tap_orders[fid]))
                taps = TapRegistry(include=config.MODULE_TAPS) if want_taps else None
                if stem is not None and (stem, order) in logits_cache and taps is None:
                    logits = logits_cache[(stem, order)]
                    same = sha256_array(prepare_image(str(SAMPLES / f"{stem}.png"), channel_order=order).u8_bgr)
                    if same != sha256_array(prep.u8_bgr):
                        raise SystemExit(f"{fid}: the converter's decode differs from the shipped sample's")
                else:
                    logits = ref.run_numpy(prep.network_input(), taps=taps)["logits"]
                res = postprocess(prep.src_hw, logits=logits)
                arrays[f"logits_{order}"] = logits.astype(np.float32)
                arrays[f"cls_{order}"] = res.class_map
                r = {"fg_fraction_src": res.foreground_fraction, "class_fraction_320x640":
                     class_fractions(res.class_map), "margins": margin_stats(logits)}
                gold = PUBLIC / "golden" / fid / "ort_fp32.npz"
                if gold.is_file():
                    with np.load(gold) as z:
                        r["vs_research_ort"] = {"logits_pcc": pcc(logits, z[f"logits_{order}"]),
                                                "logits_max_abs": float(np.abs(logits - z[f"logits_{order}"]).max()),
                                                "class_map_agreement":
                                                    float((res.class_map == z[f"cls_{order}"]).mean())}
                if taps is not None:
                    out = LARGE / "taps" / f"{fid}.{order}.npz"
                    out.parent.mkdir(parents=True, exist_ok=True)
                    np.savez(out, **{k: taps[k] for k in config.MODULE_TAPS})
                    index["files"][str(out.relative_to(LARGE))] = {"bytes": out.stat().st_size,
                                                                   "sha256": sha256_file(out),
                                                                   "taps": list(config.MODULE_TAPS)}
                r["seconds"] = round(time.perf_counter() - t1, 1)
                rec["orders"][order] = r
                print(f"{fid} {order}: fg {res.foreground_fraction:.4f} {r.get('vs_research_ort')} {r['seconds']} s",
                      flush=True)
            meta = json.loads((d / "meta.json").read_text()) if (d / "meta.json").is_file() else {}
            rec["dataset"] = meta.get("dataset")
            rec["may_ship_in_bundle"] = meta.get("may_ship_in_bundle")
            rec["source_hw"] = list(bgr.shape[:2])
            out = pub / f"{fid}.npz"
            np.savez_compressed(out, **arrays)
            index["files"][str(out.relative_to(LARGE))] = {"bytes": out.stat().st_size, "sha256": sha256_file(out)}
            index["public"][fid] = rec
    comma10k_goldens(ref, index, common)
    if index["files"]:
        worst = [(fid, o, r["vs_research_ort"]) for fid, rec in index["public"].items()
                 for o, r in rec["orders"].items() if "vs_research_ort" in r]
        if worst:
            index["vs_research_ort_min"] = {
                "logits_pcc": min(w[2]["logits_pcc"] for w in worst),
                "class_map_agreement": min(w[2]["class_map_agreement"] for w in worst)}
        (LARGE / "index.json").write_text(json.dumps(index, indent=1) + "\n")
    (SMALL / "README.md").write_text(SMALL_README)
    print(f"done in {time.perf_counter() - t0:.0f} s", flush=True)
    return 0


def comma10k_goldens(ref, index: dict, common: dict) -> None:
    """The PLAN.md 2.2 comma10k gate's FP32 side: the six comma10k images of the SPEC (MIT; 5 of them SceneSeg
    training images; shipped in samples/comma10k/) with the training ROI crop, both orders: class maps (research/sceneseg/goldens/comma10k/) and
    per-image IoU against the ground truth (index ``comma10k``; SPEC section 3: foreground 0.729 bgr / 0.765 rgb)."""
    from tt_sceneseg.host.postprocess import postprocess
    from tt_sceneseg.host.preprocess import prepare_bgr
    from tt_sceneseg.tests.sceneseg_eval import comma10k_gt, comma10k_ious, comma10k_samples, load_rgb, roi70

    samples = comma10k_samples()
    if not samples:
        print("comma10k images not found: skipped", flush=True)
        return
    out_dir = LARGE / "comma10k"
    out_dir.mkdir(parents=True, exist_ok=True)
    rec = {"protocol": "training ROI crop (top 70 %, last column dropped), class map resized to the crop with "
                       "nearest neighbour, per-image IoU, mean over the images", "images": {}, "mean": {}}
    for name, (img_path, mask_path) in samples.items():
        crop = roi70(load_rgb(img_path)[:, :, ::-1])
        gt = roi70(comma10k_gt(load_rgb(mask_path)))
        arrays, r = {"gt": gt}, {"crop_hw": list(crop.shape[:2]), "image_sha256": sha256_file(img_path),
                                 "mask_sha256": sha256_file(mask_path)}
        for order in config.CHANNEL_ORDERS:
            prep = prepare_bgr(crop, channel_order=order)
            logits = ref.run_numpy(prep.network_input())["logits"]
            res = postprocess(prep.src_hw, logits=logits)
            arrays[f"cls_{order}"] = res.class_map
            r[order] = comma10k_ious(res.class_map, gt)
        np.savez_compressed(out_dir / f"{name[:4]}.npz", **arrays)
        index["files"][f"comma10k/{name[:4]}.npz"] = {"bytes": (out_dir / f"{name[:4]}.npz").stat().st_size,
                                                      "sha256": sha256_file(out_dir / f"{name[:4]}.npz")}
        rec["images"][name] = r
        print(f"comma10k {name[:4]}: {r}", flush=True)
    for order in config.CHANNEL_ORDERS:
        rec["mean"][order] = {c: float(np.nanmean([r[order][c] for r in rec["images"].values()]))
                              for c in config.CLASS_NAMES}
    index["comma10k"] = rec
    print(f"comma10k mean: {rec['mean']}", flush=True)
    # the shipped copy (tests/goldens/comma10k/, since the release docs): the class maps + GT and this record with
    # the files' sha256, so the comma10k gate and test_reference_host run outside the workspace
    small = SMALL / "comma10k"
    small.mkdir(parents=True, exist_ok=True)
    files = {}
    for name in rec["images"]:
        src = out_dir / f"{name[:4]}.npz"
        (small / src.name).write_bytes(src.read_bytes())
        files[src.name] = index["files"][f"comma10k/{src.name}"]
    (small / "index.json").write_text(json.dumps({**rec, "files": files}, indent=1, sort_keys=True) + "\n")


SMALL_README = """# Small goldens of the fp32 CPU reference (shipped)

Written by `code/scripts/ref_golden.py` (research venv, CPU) with `tt_sceneseg.reference` on the 9 shipped samples,
both channel orders (`bgr` deployed, `rgb` training order), and on the six shipped comma10k images (`comma10k/`). The large ones (module taps of the gate samples, logits of
the public frames) stay in the workspace under `research/sceneseg/goldens/` and are never shipped.

| file | content |
|---|---|
| `<stem>.npz` | `sample_yx` int32 (2, 4096) fixed pixels; per order `<order>.class_map` uint8 (320, 640) (the node's argmax, lowest-index ties), `<order>.mask` uint8 at the source size (the published mono8 mask), `<order>.logits_samples` float32 (3, 4096) at `sample_yx` |
| `comma10k/<id>.npz` | the six shipped comma10k images (training ROI crop): `gt` uint8 at the crop's size (the comma10k mask mapped to background / foreground / road like `process_comma10k.py`), `cls_bgr` / `cls_rgb` uint8 (320, 640) fp32 class maps |
| `comma10k/index.json` | the protocol, per-image and mean ground-truth IoUs of the fp32 reference per order (the FP32 side of the comma10k gate), the images' and masks' sha256 and every `.npz`'s sha256 |
| `<stem>.json` | sha256 of the sample, the uint8 squash and the network input; logits statistics; class fractions; decision margins (fraction of pixels within 0.05 / 0.25 / 1.0 logit of a class flip); for `highway_normal_1` every tap's shape / mean / std / absmax / sha256 |

Tests: `tests/test_reference_host.py` (the fp32 forward on another torch build vs these),
`tests/test_host_postprocess.py` (the class decision of the stored logits samples). The TT gates compare with the
CPU reference, never with a TT run.
"""


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