c151bdb819
The cargo-class push deliberately excluded the neural/semantic work to avoid reverting origin; origin never independently changed neural, so this brings it forward cleanly and wires it up. - neural/mod.rs declares the full decoder family; rans.rs gains the categorical decoder it needs (glc/mlicpp/elic/hyperprior/int_synth/ cmic/dcvc/snac + the JPEG-AI entropy/codestream/σ/value/filter port). - runtime.rs uses the shared deterministic ort_session loader. - Restores the semantic/generative path (semantic.rs, min_sim/sim_metric) behind the `semantic` feature, with the bins/tests that construct ModelRequirement updated accordingly. - Re-adds byte-exact-conformance and the tools/ export/encode scripts. Green together: wai-rs all features (75 lib + integration tests) + clippy, byte-exact-conformance (8), wai-web wasm, wai-spz-gen. Co-Authored-By: Claude Opus 4.8 <[email protected]>
103 lines
4.3 KiB
Python
103 lines
4.3 KiB
Python
"""Golden-fixture harness for the Rust JPEG AI mANS port (jpegai_mans.rs).
|
|
|
|
Builds the ISO/IEC 6048-1 reference software mANS extension (BSD-3, public:
|
|
gitlab.com/wg1/jpeg-ai/jpeg-ai-reference-software) and produces
|
|
/tmp/jpegai_mans_goldens.json: decode tables built by the reference utils,
|
|
plus y (sgm, incl. escapes) and z (factorize, incl. escape-63) bitstreams
|
|
encoded by the REFERENCE encoder. The Rust test
|
|
`jpegai_mans_matches_reference_software` then byte-compares its own table
|
|
builder and decodes both streams. Build the ext first (clone repo to
|
|
/tmp/jpeg-ai-ref, omp shim, pybind11 — see tools/README or memory notes);
|
|
z values are 0..62 (63 is the internal escape marker).
|
|
"""
|
|
import sys, json
|
|
import numpy as np
|
|
sys.path.insert(0, '/tmp/jpegai_ext')
|
|
sys.path.insert(0, '/tmp/jpeg-ai-ref/src/codec/entropy_coding/lib_wrappers/mans')
|
|
import ans
|
|
from utils import get_cdf_matrix, get_encode_transitions, get_state_maps, get_decode_transitions
|
|
|
|
rng = np.random.RandomState(11)
|
|
# 4 distributions shaped like the normative set: short pdf rows (mass-256 pmfs
|
|
# over interleaved symbol order) + the mandatory 1 at index bound*2-1.
|
|
pdf_rows = [
|
|
[128, 64, 32, 16, 8, 4, 2, 1], # sharp, bound 4
|
|
[64, 64, 48, 32, 24, 12, 8, 2, 1, 1], # medium, bound 5
|
|
[32, 32, 32, 32, 28, 24, 20, 16, 12, 8, 8, 4, 4, 2, 1, 1], # wide, bound 8
|
|
[16]*14 + [8]*2 + [4]*3 + [2]*1 + [1, 1], # very wide, bound 11
|
|
]
|
|
bounds = np.array([4, 5, 8, 11], dtype=np.uint8)
|
|
NSYM = 256
|
|
pmf = np.zeros([4, NSYM], dtype=np.int64)
|
|
for i, row in enumerate(pdf_rows):
|
|
row = list(row)
|
|
s = sum(row)
|
|
assert s <= 256, s
|
|
row[0] += 256 - s # normalize mass to exactly 256
|
|
pmf[i, :len(row)] = row
|
|
pmf[i, bounds[i]*2 - 1] = max(pmf[i, bounds[i]*2 - 1], 1)
|
|
# re-fix mass after the bound-slot guarantee
|
|
excess = pmf[i].sum() - 256
|
|
pmf[i, 0] -= excess
|
|
cdf = get_cdf_matrix(pmf)
|
|
enc_t = get_encode_transitions(pmf, cdf)
|
|
dec_t = get_decode_transitions(pmf, cdf)
|
|
smaps = get_state_maps(cdf)
|
|
|
|
# --- y stream (sgm): 200 elements, some masked off, some outbound escapes ---
|
|
N = 200
|
|
sigma = rng.randint(0, 4, N).astype(np.uint8)
|
|
masks = (rng.rand(N) > 0.15)
|
|
r = np.zeros(N, dtype=np.int16)
|
|
for i in range(N):
|
|
b = int(bounds[sigma[i]])
|
|
if rng.rand() < 0.08:
|
|
r[i] = rng.randint(b, 200) * (1 if rng.rand() < 0.5 else -1) # escape
|
|
else:
|
|
r[i] = rng.randint(-(b-1), b) # in-table
|
|
r[~masks] = 0
|
|
enc = ans.ANSEncoder(1 << 20, 1)
|
|
enc.set_sgm_transitions(enc_t.data, bounds.data, smaps.data)
|
|
r_enc = r.copy()
|
|
enc.encode_sgm(sigma, r_enc, masks)
|
|
size = enc.close()
|
|
mem = np.empty([size], dtype=np.uint8); enc.get_memory(mem)
|
|
ts = np.empty([1], dtype=np.uint32); enc.get_thread_sizes(ts)
|
|
off = np.cumsum(ts).astype(np.uint32)
|
|
# reference decode (sanity)
|
|
dec = ans.ANSDecoder(mem, off)
|
|
dec.set_sgm_transitions(dec_t.data, bounds.data)
|
|
r_dec = np.zeros(N, dtype=np.int16)
|
|
dec.decode_sgm(sigma, r_dec, masks)
|
|
assert (r_dec[masks] == r[masks]).all(), "reference round-trip failed"
|
|
|
|
# --- z stream (factorize): 2 channels x 96, incl. escape value 63 ---
|
|
def zcdf(rng):
|
|
p = rng.randint(1, 9, 63).astype(np.int64)
|
|
c = np.cumsum(p)
|
|
return ((c * 255 + (c[-1] >> 1)) // c[-1]).astype(np.uint8)
|
|
zc = np.stack([zcdf(rng), zcdf(rng)])
|
|
z = rng.randint(0, 63, [2, 96]).astype(np.uint8) # 0..62; 63 is the internal escape
|
|
enc2 = ans.ANSEncoder(1 << 20, 1)
|
|
enc2.set_sgm_transitions(enc_t.data, bounds.data, smaps.data)
|
|
z_enc = z.copy()
|
|
enc2.encode_factorize(zc, z_enc)
|
|
size2 = enc2.close()
|
|
mem2 = np.empty([size2], dtype=np.uint8); enc2.get_memory(mem2)
|
|
ts2 = np.empty([1], dtype=np.uint32); enc2.get_thread_sizes(ts2)
|
|
off2 = np.cumsum(ts2).astype(np.uint32)
|
|
dec2 = ans.ANSDecoder(mem2, off2)
|
|
dec2.set_sgm_transitions(dec_t.data, bounds.data)
|
|
z_dec = np.zeros_like(z)
|
|
dec2.decode_factorize(zc, z_dec)
|
|
assert (z_dec == z).all(), "z round-trip failed"
|
|
|
|
json.dump({
|
|
"pmf": pmf.tolist(), "cdf": cdf.tolist(), "bounds": bounds.tolist(),
|
|
"decode_transitions": dec_t.tolist(),
|
|
"y": {"memory": mem.tolist(), "offsets": off.tolist(),
|
|
"sigma": sigma.tolist(), "masks": masks.tolist(), "expected": r.tolist()},
|
|
"z": {"memory": mem2.tolist(), "offsets": off2.tolist(),
|
|
"cdfs": zc.tolist(), "expected": z.tolist()},
|
|
}, open("/tmp/jpegai_mans_goldens.json", "w"))
|
|
print(f"goldens OK: y stream {size}B ({int(masks.sum())} coded, escapes incl), z stream {size2}B (2ch x 96, escape 63 incl)")
|