3b6fd99989
WAI is a container + capability-dispatch standard for media. It does not re-implement codecs; it dispatches to SOTA standard libraries (PNG, FLAC, zstd as the mandatory floor; AVIF, JPEG-XL, Opus, AV1, XZ as the recommended modern set) and adds an envelope that lets a neural shared-prior path coexist with the model-free floor. The capability a file requires is named, not supplied. Two paths open every conforming file: the neural condition (sink regenerates from a shared ambient prior) and the zeroth condition (registered SOTA codec menu). Apache-2.0. Adds: - wai/SPEC.md — WAI v1.0 spec (envelope, manifest, capability dispatch, registered capabilities including v1.1 multi-rendition extension) - wai/wai-rs/ — Rust reference implementation: lib + cdylib + staticlib, C FFI, 11 capability + envelope round-trip tests passing - wai/wai-web/ — wasm-bindgen wrapper for browser sinks - wai/wai-neural/ — Python tooling for neural-condition models - wai/tools/ — codec export + encode helpers (Python) - wai/corpus/CORPUS.md — third-party test corpus reacquisition guide Updates README.md to list four standards instead of three. Known consumer: CommunicationOS uses WAI for media transport over MoQT/MLS sessions. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
102 lines
3.8 KiB
Python
102 lines
3.8 KiB
Python
"""WAI source-side neural encoder for `wai.neural.bmshj2018`.
|
|
|
|
bmshj2018-factorized (Ballé-Minnen-Singh-Hwang-Johnston 2018) from
|
|
CompressAI's model zoo. Encodes via CompressAI's rANS coder — the
|
|
exact byte format upstream uses — and ships the bitstream in the WAI
|
|
envelope.
|
|
|
|
Wire format (v0.4):
|
|
<IIIIBB> H, W, L, quality, C, S_log2
|
|
L bytes CompressAI rANS bitstream (one batch element only)
|
|
|
|
The sink loads `cdfs_q<quality>.json` once per (model, quality), runs
|
|
the matching rANS decoder, dequantizes int symbols back to float
|
|
latents (adding per-channel medians), and runs g_s through ONNX.
|
|
|
|
Usage:
|
|
python tools/wai_bmshj2018_encode.py <input.png> -o <output.wai>
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import struct
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from compressai.zoo import models
|
|
from PIL import Image
|
|
|
|
WAI_MAGIC = b"WAI1"
|
|
QUALITY = 3 # CompressAI quality knob (1..8); must match the shipped CDF JSON
|
|
|
|
|
|
def encode(img_path: str, out_path: str, size: int = 256) -> None:
|
|
img = Image.open(img_path).convert("RGB")
|
|
orig_w, orig_h = img.size
|
|
img = img.resize((size, size), Image.LANCZOS)
|
|
arr = np.asarray(img).astype(np.float32) / 255.0
|
|
x = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
|
|
|
|
print(f"loading bmshj2018-factorized (quality={QUALITY})…")
|
|
net = models["bmshj2018-factorized"](quality=QUALITY, pretrained=True).eval()
|
|
net.update(force=True)
|
|
|
|
with torch.no_grad():
|
|
out = net.compress(x) # {strings: [[bytes]], shape: latent HW}
|
|
# PSNR sanity (decode via the same model)
|
|
x_hat = net.decompress(out["strings"], out["shape"])["x_hat"].clamp(0, 1)
|
|
psnr = 10 * np.log10(1.0 / float(((x_hat - x) ** 2).mean()))
|
|
|
|
rans_bytes = out["strings"][0][0]
|
|
H_lat, W_lat = out["shape"]
|
|
assert H_lat == W_lat, f"non-square latent {H_lat}x{W_lat}"
|
|
s_log2 = int(np.log2(H_lat))
|
|
assert (1 << s_log2) == H_lat
|
|
C = int(net.entropy_bottleneck._quantized_cdf.shape[0])
|
|
|
|
header = struct.pack("<IIIIBB", size, size, len(rans_bytes), QUALITY, C, s_log2)
|
|
payload = header + rans_bytes
|
|
|
|
manifest = {
|
|
"wai": "1.0", "media": "image", "intent": "replicate",
|
|
"model_requirement": {
|
|
"capability": "wai.neural.bmshj2018",
|
|
"fallback": "wai.image.jxl",
|
|
},
|
|
"conditioning": {"kind": "bmshj2018_rans"},
|
|
"target": {"H": size, "W": size, "model_quality": QUALITY,
|
|
"orig_h": orig_h, "orig_w": orig_w,
|
|
"encoder_psnr_db": round(psnr, 2)},
|
|
}
|
|
mb = json.dumps(manifest, separators=(",", ":")).encode()
|
|
with open(out_path, "wb") as f:
|
|
f.write(WAI_MAGIC)
|
|
f.write(struct.pack("<I", len(mb))); f.write(mb)
|
|
f.write(struct.pack("<I", len(payload))); f.write(payload)
|
|
|
|
envelope_size = Path(out_path).stat().st_size
|
|
raw_rgb = size * size * 3
|
|
print(f"input image: {orig_w}x{orig_h} resized to {size}x{size} "
|
|
f"({raw_rgb:,} B raw RGB)")
|
|
print(f" latent shape: [C={C}, S={H_lat}]")
|
|
print(f" rANS bytes: {len(rans_bytes):,}")
|
|
print(f" envelope-side reconstruction PSNR: {psnr:.2f} dB")
|
|
print(f"WAI envelope: {envelope_size:>10,} B "
|
|
f"({raw_rgb/envelope_size:6.1f}x compressed vs raw RGB)")
|
|
print(f"model on the wire: 0 (bmshj2018 weights stay at the sink)")
|
|
print(f"capability declared:{manifest['model_requirement']['capability']:>30}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("input")
|
|
p.add_argument("-o", "--output", required=True)
|
|
p.add_argument("--size", type=int, default=256,
|
|
help="Square side length (must be /16).")
|
|
a = p.parse_args()
|
|
if a.size % 16 != 0:
|
|
sys.exit("--size must be a multiple of 16")
|
|
encode(a.input, a.output, a.size)
|