#!/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) ``.npz`` = the 320x640 class map, the published mask at the source size and the 3 logits at 4096 fixed pixels (``sample_yx``), and ``.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..reference.json`` (``ReferencePipeline.reference_body``: the ``/predict`` body without timings). - LARGE, never shipped (``research/sceneseg/goldens/``): ``taps/..npz`` = the module taps of ``config.MODULE_TAPS`` (fp32; what the TT module gates compare and teacher-force) for the gate samples, and ``public/.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//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 | |---|---| | `.npz` | `sample_yx` int32 (2, 4096) fixed pixels; per order `.class_map` uint8 (320, 640) (the node's argmax, lowest-index ties), `.mask` uint8 at the source size (the published mono8 mask), `.logits_samples` float32 (3, 4096) at `sample_yx` | | `comma10k/.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 | | `.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())