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]>
114 lines
4.5 KiB
Python
114 lines
4.5 KiB
Python
"""Encode an image to a `wai.neural.mbt2018_mean` WAI envelope, decoded by
|
|
`wai-rs/src/neural/hyperprior.rs`.
|
|
|
|
Mean-scale hyperprior: the envelope carries TWO rANS bitstreams — the
|
|
factorized hyper-latent `z`, then the Gaussian-conditional main latent `y`.
|
|
CompressAI's rANS is the coder `wai-rs/src/neural/rans.rs` ports (precision=16,
|
|
bypass=4) and flushes 32-bit words, so both bitstreams are multiples of 4 bytes,
|
|
exactly as the decoder requires.
|
|
|
|
Wire payload (little-endian, must match hyperprior.rs):
|
|
<IIIIHHBBH> H, W, Lz, Ly, Cz, Cy, Sz_log2, Sy_log2, q
|
|
Lz bytes z bitstream
|
|
Ly bytes y bitstream
|
|
|
|
Container: WAI1 | u32 manifest_len | manifest JSON | u32 payload_len | payload.
|
|
|
|
Note on `intent`: this is a LOSSY codec, so "replicate" here means *every sink
|
|
reconstructs the identical (lossy) image*, not the original. That cross-sink
|
|
identity only holds if the decode is integer-exact (see the determinism work in
|
|
container.rs); float ONNX synthesis is "exact on this build" — so we DON'T pin a
|
|
prior, and `decode_envelope_strict` will (correctly) flag a replicate claim as
|
|
unsound until an integer-exact model is registered. Lenient decode works today.
|
|
|
|
Run (needs torch + compressai + Pillow):
|
|
python tools/wai_mbt2018_mean_encode.py --quality 3 --size 256 \
|
|
--image in.png --out kodim.mbt2018_mean.wai
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import struct
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from compressai.zoo import models
|
|
from PIL import Image
|
|
import numpy as np
|
|
|
|
MAGIC = b"WAI1"
|
|
|
|
|
|
def log2_exact(n: int) -> int:
|
|
e = int(round(math.log2(n)))
|
|
if (1 << e) != n:
|
|
raise ValueError(f"latent side {n} is not a power of two; decoder uses S_log2")
|
|
return e
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--quality", type=int, default=3)
|
|
ap.add_argument("--size", type=int, default=256, help="square crop/resize side")
|
|
ap.add_argument("--image", required=True)
|
|
ap.add_argument("--out", required=True)
|
|
ap.add_argument("--model-dir", required=True,
|
|
help="dir with the EXPORTED hyper_synthesis.onnx — the encoder MUST compute "
|
|
"the entropy params with the same h_s the decoder runs, or float "
|
|
"round-off flips scale buckets and desyncs the rANS stream.")
|
|
args = ap.parse_args()
|
|
|
|
net = models["mbt2018-mean"](quality=args.quality, pretrained=True).eval()
|
|
net.update(force=True)
|
|
gc = net.gaussian_conditional
|
|
|
|
img = Image.open(args.image).convert("RGB").resize((args.size, args.size), Image.LANCZOS)
|
|
x = torch.from_numpy(np.asarray(img, dtype=np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0)
|
|
|
|
import onnxruntime as ort
|
|
hs = ort.InferenceSession(str(Path(args.model_dir) / "hyper_synthesis.onnx"),
|
|
providers=["CPUExecutionProvider"])
|
|
|
|
with torch.no_grad():
|
|
y = net.g_a(x)
|
|
z = net.h_a(y)
|
|
z_strings = net.entropy_bottleneck.compress(z)
|
|
z_hat = net.entropy_bottleneck.decompress(z_strings, z.size()[-2:])
|
|
# Entropy params from the EXPORTED ONNX h_s — identical to the decoder.
|
|
in_name = hs.get_inputs()[0].name
|
|
params = torch.from_numpy(hs.run(None, {in_name: z_hat.numpy()})[0])
|
|
scales, means = params.chunk(2, 1)
|
|
indexes = gc.build_indexes(scales)
|
|
y_strings = gc.compress(y, indexes, means=means)
|
|
y_bytes = y_strings[0]
|
|
z_bytes = z_strings[0]
|
|
_, Cy, Sy, _ = y.shape
|
|
_, Cz, Sz, _ = z.shape
|
|
if len(z_bytes) % 4 or len(y_bytes) % 4:
|
|
raise RuntimeError("rANS bitstream not a multiple of 4 bytes — decoder seeds 32-bit words")
|
|
|
|
H = W = args.size
|
|
header = struct.pack("<IIIIHHBBH", H, W, len(z_bytes), len(y_bytes),
|
|
Cz, Cy, log2_exact(Sz), log2_exact(Sy), args.quality)
|
|
payload = header + z_bytes + y_bytes
|
|
|
|
manifest = {
|
|
"wai": "1.0",
|
|
"media": "image",
|
|
"intent": "replicate",
|
|
"model_requirement": {"capability": "wai.neural.mbt2018_mean"},
|
|
"conditioning": {"kind": "mbt2018_mean"},
|
|
"target": {"w": W, "h": H},
|
|
}
|
|
mb = json.dumps(manifest, separators=(",", ":")).encode()
|
|
blob = MAGIC + struct.pack("<I", len(mb)) + mb + struct.pack("<I", len(payload)) + payload
|
|
Path(args.out).write_bytes(blob)
|
|
print(f"wrote {len(blob):,} B → {args.out} "
|
|
f"(z={len(z_bytes)}B, y={len(y_bytes)}B, latents y[1,{Cy},{Sy},{Sy}] z[1,{Cz},{Sz},{Sz}])")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|