51ab07aff1
Composes every validated stage into one decode of test256.bin and compares the
result to the reference test256_dec.png:
codestream ─┬ luma: z→i_s→σ/mask→SORP→y_q→resi→ψ→MCM→ŷ→LSBS→decode_y
└ chroma: z→i_s→σ/mask→SORS→y_q→resi→ψ→(no MCM)→ŷ→LSBS→decode_uv
YUV → EFElinear → eICCI → EFEnonlinear → LEF → BT.709 → RGB
Result: max|Δ| = 1, 4.45% of samples differ. Every differing sample is off by
exactly one, and the fraction sits at the PTQ floor already measured per plane —
luma ~1.4%, chroma ~3.4%, so ~4.4% across the combined RGB is expected. This is
not bit-equality with the reference PNG and is not meant to be: ψ, the MCM stages
and synthesis are 16-bit PTQ. The contract the integer path keeps is determinism —
byte-identical output on every machine — bought with ~1 LSB of accuracy.
Composing the stages is what earned this test its place. Each stage passed in
isolation, yet the composition surfaced four wiring faults no per-stage test could
see, three of them the exact class this test exists to catch:
* hyper_scale / hyper_decoder take z at the 4×4 hyper resolution, not the 16×16
latent plane — a hardcoded 16 indexed off the end of the weights;
* reconstruct_yhat / LSBS take the 16×16 plane dims, which the function halves
internally — passing 2× double-counted;
* eICCI runs in [0,1]; the chain handed it [0,255], the folded BN and ReLUs
misbehaved, and it scored worse than a no-op (already fixed in the eICCI
commit, re-confirmed composed here);
* EFElinear applies `filters_best2`, NOT `filters_best` — the naming is inverted,
and the wrong group has a ~−16 luma cross-gain that blew U/V to −3882.
wai_jpegai_build_e2e_inputs.py extracts every runtime input — the model weights
for both planes, the normative tables, the LSBS and LEF op-point tables, and the
reference RGB — straight from the shipped assets, so the test is reproducible from
the assets alone rather than from hand-made scratch files. Verified: regenerating
the inputs with it and re-running gives the identical 4.45%.
Co-Authored-By: Claude Opus 4.8 <[email protected]>
84 lines
3.5 KiB
Python
84 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Build every runtime input the whole-decoder test (`jpegai_e2e_full`) needs, from
|
|
the shipped JPEG-AI assets. Nothing here is derived or fitted — it is a straight
|
|
extraction of model weights and normative tables, so the end-to-end test is
|
|
reproducible rather than depending on hand-made scratch files.
|
|
|
|
python wai_jpegai_build_e2e_inputs.py <assets_dir> <out_dir>
|
|
|
|
<assets_dir> is the jpegai-vm-assets root. Writes:
|
|
<out_dir>/weights/{hs,hyp}_{y,uv}.json, mcm{0..3}_y.json, synth_{y,uv}.json, icci.json
|
|
<out_dir>/norm.json (pmf, bounds, per-channel z size, factorized CDFs, gain
|
|
vectors, LSBS tables, LEF op-point tables, reference RGB)
|
|
|
|
The weight exports call `wai_jpegai_export_weights.py` (same directory); norm.json is
|
|
assembled here. Requires numpy + Pillow (only for reading the reference PNG).
|
|
|
|
Model selection is signalled by the picture (icci_model_signalled_idx → Y=model_6,
|
|
UV=model_3); those two wrappers live under e2e/pfchain and are what the exporter's
|
|
`icci` mode reads.
|
|
"""
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
EXPORT = os.path.join(HERE, "wai_jpegai_export_weights.py")
|
|
|
|
|
|
def export(mode, *args):
|
|
subprocess.run([sys.executable, EXPORT, mode, *args], check=True)
|
|
|
|
|
|
def main(assets, out):
|
|
W = os.path.join(out, "weights")
|
|
os.makedirs(W, exist_ok=True)
|
|
M = os.path.join(assets, "model/CCS_SGMM/tools_0")
|
|
M16 = os.path.join(assets, "model16/CCS_SGMM/tools_0")
|
|
|
|
# --- model weights ---
|
|
for pl in ("y", "uv"):
|
|
export("int", f"{M}/model_{pl}/common_modules/hyper_scale_decoder.onnx", f"{W}/hs_{pl}.json")
|
|
export("ptq", f"{M}/model_{pl}/common_modules/hyper_decoder.onnx", f"{W}/hyp_{pl}.json")
|
|
export("ptq", f"{M16}/model_{pl}/synthesis16.onnx", f"{W}/synth_{pl}.json")
|
|
for i in range(4):
|
|
export("ptq", f"{M}/model_y/common_modules/MCM/stage{i}.onnx", f"{W}/mcm{i}_y.json")
|
|
export("icci", f"{assets}/e2e/pfchain/icci_y_m6.onnx",
|
|
f"{assets}/e2e/pfchain/icci_uv_m3.onnx", f"{W}/icci.json")
|
|
|
|
# --- normative tables + per-image reference RGB ---
|
|
fac = json.load(open(f"{assets}/jpegai_factorized_cdfs.json"))
|
|
norm = json.load(open(f"{assets}/jpegai_normative_tables.json"))
|
|
mc = json.load(open(f"{assets}/model_constants.json"))
|
|
bounds = [int(x) for x in open(f"{assets}/me-tANS/bounds.csv") if x.strip()]
|
|
|
|
def cdfs(pl):
|
|
cd = fac[pl]
|
|
return [int(x) for row in cd for x in row] if isinstance(cd[0], list) else [int(x) for x in cd]
|
|
|
|
def gain(pl):
|
|
p = f"{M}/model_{pl}/common_modules/quantizer/gain_unit_mlog.csv"
|
|
return [int(v) for v in np.loadtxt(p, delimiter=",", dtype=int)]
|
|
|
|
rgb = np.array(Image.open(f"{assets}/e2e/test256_dec.png").convert("RGB"))
|
|
data = {
|
|
"pmf": norm["pmf"], "bounds": bounds, "per_channel": fac["per_channel"],
|
|
"cdfs_y": cdfs("y"), "cdfs_uv": cdfs("uv"), "gain_y": gain("y"), "gain_uv": gain("uv"),
|
|
"lsbs_t0": [int(v) for v in np.load(f"{assets}/e2e/lsbs/lsbs_table0.npy")],
|
|
"lsbs_t1": [int(v) for v in np.load(f"{assets}/e2e/lsbs/lsbs_table1.npy")],
|
|
"lef_thr0": mc["lef_thr"][0], "lef_mag0": mc["lef_mag"][0],
|
|
"ref_rgb": rgb.ravel().tolist(),
|
|
}
|
|
json.dump(data, open(f"{out}/norm.json", "w"))
|
|
print(f"wrote {W}/*.json (11 weight files) and {out}/norm.json")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) != 3:
|
|
print(__doc__)
|
|
sys.exit(1)
|
|
main(sys.argv[1], sys.argv[2])
|