wai-rs: derive the dequant gain — the last fixture-sourced decoder input #29
16 changed files with 1890 additions and 104 deletions
84
wai/tools/wai_jpegai_build_e2e_inputs.py
Normal file
84
wai/tools/wai_jpegai_build_e2e_inputs.py
Normal file
|
|
@ -0,0 +1,84 @@
|
|||
#!/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])
|
||||
|
|
@ -205,6 +205,61 @@ def list_tensors(src):
|
|||
print(f"{total:,} params")
|
||||
|
||||
|
||||
def _conv_node(w_dims, w_raw, b_raw):
|
||||
"""One `Conv` as `jpegai_icci`'s `conv_of` reads it. Padding is kernel//2, which
|
||||
keeps the spatial size for both the 3x3 stems and the factorized 1x3 / 3x1 taps."""
|
||||
cout, cin, kh, kw = w_dims
|
||||
return {"w": _decode(w_dims, 1, w_raw), "b": _decode([len(b_raw) // 4], 1, b_raw),
|
||||
"cin": cin, "cout": cout, "kh": kh, "kw": kw, "ph": kh // 2, "pw": kw // 2}
|
||||
|
||||
|
||||
def _icci_path(inits, n_blocks, conv_last_name, scale_name, hidden_prefix):
|
||||
"""One eICCI plane-path out of a wrapper's initializers.
|
||||
|
||||
Heads and scales are found by NAME; the trunk convolutions are anonymous
|
||||
(`onnx::Conv_*`) and are assigned by FILE ORDER — conv_first, then each residual
|
||||
block's 1x3 tap followed by its 3x1 tap. The wrappers are BN-pre-folded, so
|
||||
nothing is folded here.
|
||||
"""
|
||||
named = {n: (d, r) for n, (d, _dt, r) in inits.items()}
|
||||
anon = [(d, r) for n, (d, _dt, r) in inits.items() if n.startswith("onnx::Conv")]
|
||||
convs = [(anon[i][0], anon[i][1], anon[i + 1][1]) for i in range(0, len(anon), 2)]
|
||||
assert len(convs) == 1 + 2 * n_blocks, \
|
||||
f"{len(convs)} convs for {n_blocks} blocks in {conv_last_name}"
|
||||
scalar = lambda name: _decode(named[name][0], 1, named[name][1])[0]
|
||||
blocks = [{"c1": _conv_node(*convs[1 + 2 * b]),
|
||||
"c2": _conv_node(*convs[2 + 2 * b]),
|
||||
"scale": scalar(f"{hidden_prefix}.{b}.scale")} for b in range(n_blocks)]
|
||||
cl_d, cl_r = named[conv_last_name + ".weight"]
|
||||
return {"conv_first": _conv_node(*convs[0]),
|
||||
"blocks": blocks,
|
||||
"conv_last": _conv_node(cl_d, cl_r, named[conv_last_name + ".bias"][1]),
|
||||
"scale": scalar(scale_name)}
|
||||
|
||||
|
||||
def export_icci(y_src, uv_src, dst):
|
||||
"""The two eICCI model wrappers -> the `IcciNet` weight JSON.
|
||||
|
||||
Which wrappers is signalled per picture (`icci_model_signalled_idx` indexes a
|
||||
shortlist); for the 256x256 conformance stream that is Y = model_6, UV = model_3.
|
||||
The luma path has 2 residual blocks and the chroma path 4, and U and V carry
|
||||
separate residual scales.
|
||||
"""
|
||||
luma = _icci_path(initializers(y_src), 2, "comp_model.conv_last_Y",
|
||||
"comp_model.scaleY", "comp_model.hidden_layer_Y")
|
||||
ui = initializers(uv_src)
|
||||
chroma = _icci_path(ui, 4, "comp_model.conv_last_u",
|
||||
"comp_model.scaleU", "comp_model.hidden_layer_UV")
|
||||
clv = ui["comp_model.conv_last_v.weight"]
|
||||
out = {"luma": luma, "chroma": chroma,
|
||||
"conv_last_v": _conv_node(clv[0], clv[2], ui["comp_model.conv_last_v.bias"][2]),
|
||||
"scale_v": _decode(ui["comp_model.scaleV"][0], 1, ui["comp_model.scaleV"][2])[0]}
|
||||
json.dump(out, open(dst, "w"))
|
||||
print(f"eICCI: luma {len(luma['blocks'])} blocks, chroma {len(chroma['blocks'])} blocks; "
|
||||
f"scaleY={luma['scale']:.6f} scaleU={chroma['scale']:.6f} "
|
||||
f"scaleV={out['scale_v']:.6f} -> {dst}")
|
||||
|
||||
|
||||
def main(argv):
|
||||
if len(argv) < 3:
|
||||
print(__doc__)
|
||||
|
|
@ -216,8 +271,10 @@ def main(argv):
|
|||
export_ptq(src, argv[3])
|
||||
elif mode == "int":
|
||||
export_int(src, argv[3])
|
||||
elif mode == "icci":
|
||||
export_icci(argv[2], argv[3], argv[4]) # icci <y_wrapper> <uv_wrapper> <out>
|
||||
else:
|
||||
print(f"unknown mode {mode!r}; expected list | ptq | int")
|
||||
print(f"unknown mode {mode!r}; expected list | ptq | int | icci")
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
|
|
|||
230
wai/wai-rs/src/jpegai_decode.rs
Normal file
230
wai/wai-rs/src/jpegai_decode.rs
Normal file
|
|
@ -0,0 +1,230 @@
|
|||
//! The whole JPEG-AI decoder as one call: `decode(codestream, weights, norm) → Rgb`.
|
||||
//!
|
||||
//! This composes the stage functions validated across `jpegai_*` — the entropy front
|
||||
//! half (`jpegai_mans`, `jpegai_pipeline`), the integer value path and synthesis
|
||||
//! (`jpegai_synthesis`), and the post-filter chain (`jpegai_pipeline`, `jpegai_icci`) —
|
||||
//! into a single entry point. It is pure compute: the caller supplies the model weights
|
||||
//! and normative tables (see [`DecoderWeights`] / [`NormTables`]); this module reads no
|
||||
//! files and does no I/O.
|
||||
//!
|
||||
//! Determinism, not float-exactness. ψ, the four MCM stages and synthesis are 16-bit
|
||||
//! PTQ, so the decoded pixels differ from the float reference by ~1 LSB on a few percent
|
||||
//! of samples. Everything is deterministic — integer arithmetic, or IEEE-754 float with
|
||||
//! a fixed accumulation order — so the output is byte-identical on every machine, which
|
||||
//! is the contract. See `tests/jpegai_e2e_full.rs`, which drives this and compares the
|
||||
//! reference decode of `test256.bin` (max|Δ| = 1, ~4.4% of samples).
|
||||
//!
|
||||
//! Scope: the 256×256, 4:4:4, tools_0 profile the reference ships — luma 160 channels
|
||||
//! (4 MCM stages), chroma 96 (no MCM). Other geometries are rejected up front.
|
||||
|
||||
use crate::jpegai_bitstream::{find_segment, parse_container, parse_picture_header,
|
||||
parse_tools_header, Marker};
|
||||
use crate::jpegai_icci::IcciNet;
|
||||
use crate::jpegai_mans::{build_decode_transitions, cdf_from_pmf, MansDecoder};
|
||||
use crate::jpegai_pipeline::{dequantize_resi, efl_chroma, efn_filter, gain_unit_scaler_vec,
|
||||
lef_filter, lsbs_refine, reconstruct_yhat, reconstruct_yhat_uv, rvs_additive,
|
||||
rvs_likely, sigma_index, skip_mask};
|
||||
use crate::jpegai_synthesis::{decode_hyper, decode_hyper_scale, decode_uv_f32, decode_y_f32,
|
||||
q16_from_symbols, IntNet, Weights};
|
||||
|
||||
/// The `z` symbol offset (de-bias) applied to the factorized hyper-latent.
|
||||
const Z_OFFSET: i64 = 31;
|
||||
|
||||
/// One plane's model weights: the fully-integer hyper-scale net, the PTQ hyper decoder,
|
||||
/// and the PTQ synthesis net.
|
||||
pub struct PlaneWeights {
|
||||
pub hyper_scale: IntNet,
|
||||
pub hyper_decoder: Weights,
|
||||
pub synthesis: Weights,
|
||||
}
|
||||
|
||||
/// Every model weight the decode needs. Luma carries the four MCM stages; chroma has no
|
||||
/// MCM (its hyper params are used directly as the means). `icci` is the whole eICCI net.
|
||||
pub struct DecoderWeights {
|
||||
pub luma: PlaneWeights,
|
||||
pub chroma: PlaneWeights,
|
||||
pub mcm_y: [Weights; 4],
|
||||
pub icci: IcciNet,
|
||||
}
|
||||
|
||||
/// The normative, image-independent tables. All are shipped decoder assets, not derived
|
||||
/// per image: the me-tANS PMF/bounds, the factorized-`z` CDFs, the per-channel gain
|
||||
/// vectors (`gain_unit_mlog`), the LSBS refinement tables, and the LEF op-point tables
|
||||
/// selected at the picture's `model_id`.
|
||||
pub struct NormTables {
|
||||
pub pmf: Vec<Vec<i64>>,
|
||||
pub bounds: Vec<u8>,
|
||||
pub per_channel: usize,
|
||||
pub cdfs_y: Vec<u8>,
|
||||
pub cdfs_uv: Vec<u8>,
|
||||
pub gain_y: Vec<i64>,
|
||||
pub gain_uv: Vec<i64>,
|
||||
pub lsbs_t0: Vec<i32>,
|
||||
pub lsbs_t1: Vec<i32>,
|
||||
pub lef_thr: [i64; 3],
|
||||
pub lef_mag: [f32; 3],
|
||||
}
|
||||
|
||||
/// A decoded RGB image, row-major `[h][w][3]` (`data[(y*w + x)*3 + c]`).
|
||||
pub struct Rgb {
|
||||
pub w: usize,
|
||||
pub h: usize,
|
||||
pub data: Vec<u8>,
|
||||
}
|
||||
|
||||
/// Header gains are offset-binary Q11 around 32767.
|
||||
fn deint(q: u32) -> f32 { (q as f32 - 32767.0) / 2048.0 }
|
||||
fn q16(f: &[f32]) -> Vec<i64> { f.iter().map(|&x| (x as f64 * 65536.0).round() as i64).collect() }
|
||||
|
||||
/// Decode one plane's front half and value path → its LSBS-refined synthesis latent
|
||||
/// (f32), returning also `i_s` and the RVS additive for LEF's sigma surface.
|
||||
///
|
||||
/// `z_pre_channels` is the channel count consumed from the SOZ substream *before* this
|
||||
/// plane — one SOZ stream carries luma then chroma back to back through a single mANS
|
||||
/// decoder, so chroma must skip the luma `z` to reach its own.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn decode_plane(
|
||||
data: &[u8], marker: Marker, n: &NormTables, pw: &PlaneWeights,
|
||||
cdfs: &[u8], gain: &[i64], mcm: Option<[&Weights; 4]>,
|
||||
channels: usize, z_pre_channels: usize,
|
||||
) -> (Vec<f32>, Vec<i64>, Vec<i64>) {
|
||||
const H: usize = 16; // the latent plane is 16×16; z is at 4×4 and the nets upsample ×4
|
||||
const W: usize = 16;
|
||||
let segs = parse_container(data);
|
||||
|
||||
// SOZ → z (skip the prior plane's z first)
|
||||
let soz = find_segment(&segs, Marker::Soz).expect("SOZ");
|
||||
let mem = &data[soz.start..soz.start + soz.len];
|
||||
let mut dz = MansDecoder::new(mem, &[mem.len() as u32]);
|
||||
if z_pre_channels > 0 {
|
||||
let mut skip = vec![0u8; z_pre_channels * n.per_channel];
|
||||
dz.decode_factorize(&n.cdfs_y, &mut skip, z_pre_channels, n.per_channel);
|
||||
}
|
||||
let mut syms = vec![0u8; channels * n.per_channel];
|
||||
dz.decode_factorize(cdfs, &mut syms, channels, n.per_channel);
|
||||
let z: Vec<i64> = syms.iter().map(|&s| s as i64 - Z_OFFSET).collect();
|
||||
|
||||
// z → i_s → derived entropy inputs
|
||||
let (i_s, ..) = decode_hyper_scale(&z, channels, 4, 4, &pw.hyper_scale);
|
||||
let likely = rvs_likely(&i_s, gain, channels, H, W);
|
||||
let grfs = grfs_flags(data, channels, marker);
|
||||
let rvs = rvs_additive(&i_s, gain, &grfs, channels, H, W);
|
||||
let sigma = sigma_index(&i_s, gain, &rvs, channels);
|
||||
let mask = skip_mask(&i_s, gain, channels);
|
||||
|
||||
// entropy decode → y_q
|
||||
let seg = find_segment(&segs, marker).expect("SORP/SORS");
|
||||
let mem = &data[seg.start..seg.start + seg.len];
|
||||
let flat: Vec<u32> = build_decode_transitions(&n.pmf, &cdf_from_pmf(&n.pmf))
|
||||
.iter().flatten().copied().collect();
|
||||
let mut dec = MansDecoder::new(mem, &[mem.len() as u32]);
|
||||
dec.set_sgm_transitions(&flat, &n.bounds);
|
||||
let plane = H * W;
|
||||
let mut y_q = vec![0i16; channels * plane];
|
||||
for c in 0..channels {
|
||||
let mut row = vec![0i16; plane];
|
||||
dec.decode_sgm(&sigma[c * plane..(c + 1) * plane], &mut row, &mask[c * plane..(c + 1) * plane]);
|
||||
y_q[c * plane..(c + 1) * plane].copy_from_slice(&row);
|
||||
}
|
||||
|
||||
// dequant → resi ; ψ ; reconstruct → ŷ ; LSBS → latent
|
||||
let sv = gain_unit_scaler_vec(gain, 0);
|
||||
let resi = dequantize_resi(&y_q, &likely, &grfs, &sv, channels);
|
||||
let (psi, ..) = decode_hyper(&q16_from_symbols(&z), channels, 4, 4, &pw.hyper_decoder);
|
||||
let resi_q = q16(&resi);
|
||||
let yhat = match mcm {
|
||||
Some(m) => reconstruct_yhat(&resi_q, &psi, channels, H, W, m),
|
||||
None => reconstruct_yhat_uv(&resi_q, &psi, channels, H, W),
|
||||
};
|
||||
let yhat_f: Vec<f32> = yhat.iter().map(|&v| v as f64 as f32 / 65536.0).collect();
|
||||
let lat = lsbs_refine(&yhat_f, &resi, &likely, &n.lsbs_t0, &n.lsbs_t1, channels, H, W);
|
||||
(lat, i_s, rvs)
|
||||
}
|
||||
|
||||
fn grfs_flags(data: &[u8], channels: usize, marker: Marker) -> Vec<bool> {
|
||||
let segs = parse_container(data);
|
||||
let ph = parse_picture_header(data, find_segment(&segs, Marker::Pih).unwrap().start_bit());
|
||||
let comp = if marker == Marker::Sorp { &ph.components[0] } else { &ph.components[1] };
|
||||
let mut f = comp.grfs_channel_flags.clone();
|
||||
f.resize(channels, false);
|
||||
f
|
||||
}
|
||||
|
||||
/// Decode a JPEG-AI codestream to RGB. Returns `None` if the picture is not the
|
||||
/// supported 256×256 4:4:4 tools_0 geometry.
|
||||
pub fn decode(data: &[u8], w: &DecoderWeights, n: &NormTables) -> Option<Rgb> {
|
||||
let segs = parse_container(data);
|
||||
let ph = parse_picture_header(data, find_segment(&segs, Marker::Pih)?.start_bit());
|
||||
let (img_h, img_w) = (ph.img_height as usize, ph.img_width as usize);
|
||||
if (img_h, img_w) != (256, 256) || ph.components.len() != 2
|
||||
|| ph.components[0].num_chs != 160 || ph.components[1].num_chs != 96 {
|
||||
return None;
|
||||
}
|
||||
let (h, w_) = (img_h, img_w);
|
||||
|
||||
// ---- luma + chroma latents → pixel planes (float, pre-round) ----
|
||||
let mcm = [&w.mcm_y[0], &w.mcm_y[1], &w.mcm_y[2], &w.mcm_y[3]];
|
||||
let (y_lat, y_is, y_rvs) = decode_plane(data, Marker::Sorp, n, &w.luma,
|
||||
&n.cdfs_y, &n.gain_y, Some(mcm), 160, 0);
|
||||
let (uv_lat, ..) = decode_plane(data, Marker::Sors, n, &w.chroma,
|
||||
&n.cdfs_uv, &n.gain_uv, None, 96, 160);
|
||||
|
||||
let mut yp = decode_y_f32(&q16(&y_lat), &w.luma.synthesis);
|
||||
let mut cat = q16(&y_lat);
|
||||
cat.extend(q16(&uv_lat));
|
||||
let uv = decode_uv_f32(&cat, &w.chroma.synthesis);
|
||||
let (mut up, mut vp) = (uv[..h * w_].to_vec(), uv[h * w_..].to_vec());
|
||||
|
||||
// ---- post-filters, parameters from the tools header ----
|
||||
let ton = find_segment(&segs, Marker::Ton)?;
|
||||
let th = parse_tools_header(data, ton.start_bit() + 2, (ph.c_ver, ph.c_hor) == (1, 1), 3, 1);
|
||||
|
||||
// EFElinear — `filters_best2` is the applied group (the naming is inverted).
|
||||
if let Some(efl) = th.efl.as_ref() {
|
||||
let f = &efl.filters_best2;
|
||||
let (mu, mv) = (efl.b1[0] as f32 / 100.0, efl.b1[1] as f32 / 100.0);
|
||||
up = efl_chroma(&up, &yp, deint(f.weights_u[0]) + 1.0, deint(f.weights_u2[0]), mu);
|
||||
vp = efl_chroma(&vp, &yp, deint(f.weights_v[0]) + 1.0, deint(f.weights_v2[0]), mv);
|
||||
}
|
||||
// eICCI — [0,1] domain: normalise in, scale back out.
|
||||
if th.icci_enabled {
|
||||
let nrm = |p: &[f32]| p.iter().map(|&x| x / 255.0).collect::<Vec<f32>>();
|
||||
let (yo, uo, vo) = w.icci.forward(&nrm(&yp), &nrm(&up), &nrm(&vp), h, w_);
|
||||
let den = |p: Vec<f32>| p.into_iter().map(|x| x * 255.0).collect::<Vec<f32>>();
|
||||
yp = den(yo); up = den(uo); vp = den(vo);
|
||||
}
|
||||
// EFEnonlinear — per the header, applied to the enabled chroma plane(s).
|
||||
if let Some(efn) = th.efn.as_ref() {
|
||||
if efn.nonlinear_v_enabled {
|
||||
let wq: Vec<i64> = efn.weights_v.iter().map(|&x| x as i64).collect();
|
||||
vp = efn_filter(&yp, &vp, &wq, efn.luma_min[0] as i64, efn.luma_max[0] as i64);
|
||||
}
|
||||
if efn.nonlinear_u_enabled {
|
||||
let wq: Vec<i64> = efn.weights_u.iter().map(|&x| x as i64).collect();
|
||||
up = efn_filter(&yp, &up, &wq, efn.luma_min[0] as i64, efn.luma_max[0] as i64);
|
||||
}
|
||||
}
|
||||
// LEF — luma; sigma is the raw scale surface (i_s + gain + rvs) at the named channel.
|
||||
if th.lef_enabled {
|
||||
let ch = th.lef_ch_idx as usize;
|
||||
let sigma: Vec<f32> = (0..16 * 16)
|
||||
.map(|p| (y_is[ch * 256 + p] + n.gain_y[ch] + y_rvs[ch * 256 + p]) as f32).collect();
|
||||
let thr = [n.lef_thr[0] as f32, n.lef_thr[1] as f32, n.lef_thr[2] as f32];
|
||||
yp = lef_filter(&yp, &sigma, 16, 16, thr, n.lef_mag, h, w_);
|
||||
}
|
||||
|
||||
// ---- BT.709 (full-range) → RGB ----
|
||||
let (kr, kg, kb) = (0.2126f32, 0.7152f32, 0.0722f32);
|
||||
let (kby, kry) = (1.8556f32, 1.5748f32);
|
||||
let mut out = vec![0u8; 3 * h * w_];
|
||||
for i in 0..h * w_ {
|
||||
let (y01, u01, v01) = (yp[i] / 255.0, up[i] / 255.0 - 0.5, vp[i] / 255.0 - 0.5);
|
||||
let rgb = [y01 + kry * v01,
|
||||
y01 - (kb * kby / kg) * u01 - (kr * kry / kg) * v01,
|
||||
y01 + kby * u01];
|
||||
for (c, &val) in rgb.iter().enumerate() {
|
||||
out[i * 3 + c] = (val * 255.0).round().clamp(0.0, 255.0) as u8;
|
||||
}
|
||||
}
|
||||
Some(Rgb { w: w_, h, data: out })
|
||||
}
|
||||
|
|
@ -260,11 +260,16 @@ impl IcciPath {
|
|||
pub struct IcciNet {
|
||||
/// luma: `conv_first_Y` (BN pre-folded by the exporter) + its residual blocks
|
||||
pub luma: IcciPath,
|
||||
/// chroma trunk: `conv_firstUV` (fold `BNUV` in) + its residual blocks; `conv_last`
|
||||
/// is the U head and `scale` applies to both heads.
|
||||
/// chroma trunk: `conv_firstUV` (BN pre-folded) + its residual blocks; `conv_last`
|
||||
/// is the U head and `chroma.scale` is the **U** residual scale (`scaleU`).
|
||||
pub chroma: IcciPath,
|
||||
/// the V head — shares `chroma`'s trunk, differs only in this final convolution
|
||||
pub conv_last_v: Conv,
|
||||
/// the V residual scale (`scaleV`). The reference signals the U and V scales
|
||||
/// **separately** (`comp_model.scaleU` / `scaleV`); on the shipped models they
|
||||
/// differ by ~1%, so reusing one scale for both heads shifts V by more than a
|
||||
/// pixel LSB.
|
||||
pub scale_v: f32,
|
||||
}
|
||||
|
||||
impl IcciNet {
|
||||
|
|
@ -282,13 +287,13 @@ impl IcciNet {
|
|||
for blk in &self.chroma.blocks {
|
||||
fea = blk.forward(&fea, h2, w2);
|
||||
}
|
||||
let head = |conv: &Conv, band: &[f32]| -> Vec<f32> {
|
||||
let head = |conv: &Conv, band: &[f32], scale: f32| -> Vec<f32> {
|
||||
let res = conv.forward(&fea, h2, w2);
|
||||
let refined: Vec<f32> = res.iter().zip(band).map(|(&r, &s)| r * self.chroma.scale + s).collect();
|
||||
let refined: Vec<f32> = res.iter().zip(band).map(|(&r, &s)| r * scale + s).collect();
|
||||
haar_idwt2(&refined, h2, w2).0
|
||||
};
|
||||
let u_o = head(&self.chroma.conv_last, &u16);
|
||||
let v_o = head(&self.conv_last_v, &v16);
|
||||
let u_o = head(&self.chroma.conv_last, &u16, self.chroma.scale);
|
||||
let v_o = head(&self.conv_last_v, &v16, self.scale_v);
|
||||
(y_o, u_o, v_o)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -237,3 +237,83 @@ impl MansDecoder {
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The reference's symbol ordering for a distribution of `len` bins: alternating
|
||||
/// outward from the centre (`mid, mid-1, mid+1, mid-2, …`), which is what pairs each
|
||||
/// CDF slot with its signed residual value.
|
||||
fn sequence(len: usize) -> Vec<i64> {
|
||||
let mid = len / 2;
|
||||
let mut out = vec![0i64; len];
|
||||
let (mut even, mut odd) = (mid as i64, mid as i64 - 1);
|
||||
for (i, v) in out.iter_mut().enumerate() {
|
||||
if i % 2 == 0 { *v = even; even += 1; } else { *v = odd; odd -= 1; }
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Build the 256-entry `sgm` decode table per σ-distribution from the normative
|
||||
/// integer `(pmf, cdf)` rows — the integer-only port of `utils.get_decode_transitions`.
|
||||
///
|
||||
/// Each entry packs `(nBits << 24) | (next_base << 16) | (symbol & 0xffff)`, the layout
|
||||
/// [`MansDecoder::decode_sgm`] consumes. Each CDF row is split into a first and second
|
||||
/// half (`c - (c >> 1)` and `(c >> 1) + 128`) so the two interleaved ANS states draw
|
||||
/// from disjoint slot ranges.
|
||||
///
|
||||
/// Having this on the integer path means the entropy stage needs only the normative
|
||||
/// PMF tables — no precomputed transition fixture, and no dependency on the
|
||||
/// float/ONNX-side module.
|
||||
pub fn build_decode_transitions(pmf: &[Vec<i64>], cdf: &[Vec<i64>]) -> Vec<[u32; 256]> {
|
||||
const MASS_BITS: u32 = 8;
|
||||
let mass = 1i64 << MASS_BITS;
|
||||
let width = pmf[0].len();
|
||||
let seq: Vec<i64> = sequence(width).iter().map(|&v| v - (width as i64) / 2).collect();
|
||||
// get_clz(512): the number of doublings that bring `state` into [256, 512).
|
||||
let clz = |state: i64| -> i64 {
|
||||
let (mut bits, mut s) = (0i64, state.max(1));
|
||||
while s < 256 { s <<= 1; bits += 1; }
|
||||
bits
|
||||
};
|
||||
(0..pmf.len())
|
||||
.map(|i| {
|
||||
let mut symbols = [0i64; 256];
|
||||
let mut state_next = [0i64; 256];
|
||||
let cdf_first: Vec<i64> = cdf[i].iter().map(|&c| c - (c >> 1)).collect();
|
||||
let cdf_second: Vec<i64> = cdf[i].iter().map(|&c| (c >> 1) + 128).collect();
|
||||
let (mut col, mut cdf_curr) = (0usize, 0i64);
|
||||
while cdf_curr < mass {
|
||||
let p = pmf[i][col];
|
||||
let d_first = cdf_first[col + 1] - cdf_first[col];
|
||||
for (k, slot) in (cdf_first[col]..cdf_first[col + 1]).enumerate() {
|
||||
symbols[slot as usize] = seq[col];
|
||||
state_next[slot as usize] = p + k as i64;
|
||||
}
|
||||
for (k, slot) in (cdf_second[col]..cdf_second[col + 1]).enumerate() {
|
||||
symbols[slot as usize] = seq[col];
|
||||
state_next[slot as usize] = p + d_first + k as i64;
|
||||
}
|
||||
cdf_curr = cdf[i][col + 1];
|
||||
col += 1;
|
||||
}
|
||||
let mut table = [0u32; 256];
|
||||
for s in 0..256 {
|
||||
let bits = clz(state_next[s]);
|
||||
let renormed = (state_next[s] << bits) ^ mass;
|
||||
table[s] = ((bits as u32) << 24)
|
||||
| ((renormed as u32) << 16)
|
||||
| ((symbols[s] as u16) as u32);
|
||||
}
|
||||
table
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Cumulative rows for [`build_decode_transitions`] from the normative PMF tables.
|
||||
pub fn cdf_from_pmf(pmf: &[Vec<i64>]) -> Vec<Vec<i64>> {
|
||||
pmf.iter()
|
||||
.map(|row| {
|
||||
let mut c = vec![0i64; row.len() + 1];
|
||||
for (i, &p) in row.iter().enumerate() { c[i + 1] = c[i] + p; }
|
||||
c
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -115,14 +115,170 @@ pub fn reconstruct_yhat(resi: &[i64], psi: &[i64], c: usize, h2: usize, w2: usiz
|
|||
up_shuffle(&parts, c, h, w)
|
||||
}
|
||||
|
||||
/// **SGMM reconstruction for chroma** — the counterpart to [`reconstruct_yhat`].
|
||||
///
|
||||
/// ```text
|
||||
/// mu = up_shuffle(chunk(psi, 4)) // the hyper params ARE the means
|
||||
/// y_hat = resi + mu
|
||||
/// ```
|
||||
///
|
||||
/// Chroma has **no MCM**: the reference ships no `MCM/` nets under `model_uv`, and the
|
||||
/// per-stage hyper params are used directly as the means rather than being refined by
|
||||
/// an autoregressive context model. So this is a plain up-shuffle and add — no context
|
||||
/// accumulation, no stage ordering, and nothing to get wrong in the way the luma path
|
||||
/// can. Keeping it a separate named function rather than a flag on [`reconstruct_yhat`]
|
||||
/// is deliberate: the two paths are structurally different, not parameterised.
|
||||
///
|
||||
/// `resi` is `[c, 2h, 2w]` (Q16); `psi` is `[4c, h, w]` (Q16). Returns `[c, 2h, 2w]`.
|
||||
pub fn reconstruct_yhat_uv(resi: &[i64], psi: &[i64], c: usize, h2: usize, w2: usize) -> Vec<i64> {
|
||||
let (h, w) = (h2 / 2, w2 / 2);
|
||||
let plane = c * h * w;
|
||||
assert_eq!(resi.len(), c * h2 * w2, "resi must be [c, 2h, 2w]");
|
||||
assert_eq!(psi.len(), 4 * plane, "psi must be [4c, h, w]");
|
||||
let parts: [Vec<i64>; 4] = std::array::from_fn(|i| psi[i * plane..(i + 1) * plane].to_vec());
|
||||
let mu = up_shuffle(&parts, c, h, w);
|
||||
resi.iter().zip(&mu).map(|(&a, &b)| a + b).collect()
|
||||
}
|
||||
|
||||
/// Dequantize entropy-decoded latent symbols by the per-position gain map.
|
||||
/// `y_q` and `gain_map` are row-major `[channels × positions]`, same length.
|
||||
/// Returns `resi[i] = y_q[i] as f32 · gain_map[i]`.
|
||||
///
|
||||
/// Prefer [`dequantize_resi`], which derives the gain from normative assets and
|
||||
/// applies the two stages in the reference's order. This entry point remains for
|
||||
/// callers that already hold a gain map.
|
||||
pub fn dequant(y_q: &[i16], gain_map: &[f32]) -> Vec<f32> {
|
||||
assert_eq!(y_q.len(), gain_map.len(), "y_q / gain_map length mismatch");
|
||||
y_q.iter().zip(gain_map).map(|(&q, &g)| q as f32 * g).collect()
|
||||
}
|
||||
|
||||
/// Normative RVS scale list, before the per-channel GRFS branch factor.
|
||||
const RVS_BASE_SCALES: [i64; 4] = [115, 125, 146, 160];
|
||||
|
||||
/// `log_k` for the gain unit — note this is **54.82**, whereas the residual-variance
|
||||
/// scale's own `log_k` uses 54.85. The two constants differ in the reference and are
|
||||
/// easy to transpose; the additive table below is built with the 54.82 one.
|
||||
const GAIN_LOG_K: f64 = 0.200_365_493_161_539_4; // (ln 54.82 − ln 0.11) / 31
|
||||
|
||||
/// Per-branch RVS scales: `round(base · 1.06)` when the channel's GRFS flag is set,
|
||||
/// `round(base · 0.98)` otherwise.
|
||||
///
|
||||
/// The rounding is **round-half-to-even**, not half-away-from-zero: `125 · 1.06` is
|
||||
/// exactly `132.5` and `125 · 0.98` is exactly `122.5` in binary64, and the reference
|
||||
/// (NumPy/Python `round`) resolves both to even — `132` and `122`. Rust's `f64::round`
|
||||
/// would give `133`/`123` and silently select the wrong scale for one band.
|
||||
fn rvs_scales(grfs_flag: bool) -> [i64; 4] {
|
||||
let mult = if grfs_flag { 1.06 } else { 0.98 };
|
||||
std::array::from_fn(|i| (RVS_BASE_SCALES[i] as f64 * mult).round_ties_even() as i64)
|
||||
}
|
||||
|
||||
/// Per-band RVS **dequant** multipliers in Q16, `round(2^7 · 2^16 / scale)`.
|
||||
/// Applied as `resi = y_q · table[band] / 2^16` (see [`dequantize_resi`]).
|
||||
fn rvs_dequant_table(grfs_flag: bool) -> [i64; 4] {
|
||||
let s = rvs_scales(grfs_flag);
|
||||
std::array::from_fn(|i| ((1i64 << 23) as f64 / s[i] as f64).round_ties_even() as i64)
|
||||
}
|
||||
|
||||
/// Per-band RVS log-domain additive, `round(ln(scale/128) / (log_k / 2^7))`.
|
||||
/// This reproduces [`RVS_ADDITIVE`]; the constant table is kept for decode-time cost
|
||||
/// and cross-checked against this derivation in the calibration tests.
|
||||
fn rvs_additive_table(grfs_flag: bool) -> [i64; 4] {
|
||||
let s = rvs_scales(grfs_flag);
|
||||
std::array::from_fn(|i| {
|
||||
((s[i] as f64 / 128.0).ln() / (GAIN_LOG_K / 128.0)).round_ties_even() as i64
|
||||
})
|
||||
}
|
||||
|
||||
/// Gain-unit scale factor for one channel.
|
||||
///
|
||||
/// ```text
|
||||
/// scaler_vec[c] = round( exp( (gain_vector_log[c] + beta_displacement_log) · log_k / 2^7 ) · 2^10 ) / 2^10
|
||||
/// ```
|
||||
///
|
||||
/// `gain_vector_log` is the normative per-channel `gain_unit_mlog.csv` (the same
|
||||
/// vector this module calls `gain_per_ch`); `beta_displacement_log` comes from the
|
||||
/// picture header. `scaler_precision` is `gain_vector_precision + beta_displacement_precision`
|
||||
/// = `5 + 5` = 10.
|
||||
pub fn gain_unit_scaler_vec(gain_vector_log: &[i64], beta_displacement_log: i64) -> Vec<f64> {
|
||||
const SCALER_PRECISION: i32 = 10;
|
||||
let m = (1i64 << SCALER_PRECISION) as f64;
|
||||
gain_vector_log
|
||||
.iter()
|
||||
.map(|&g| (((g + beta_displacement_log) as f64 * GAIN_LOG_K / 128.0).exp() * m).round() / m)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// **Dequantize** entropy-decoded latent symbols — the full two-stage path, deriving
|
||||
/// every factor from normative assets (no fixture gain map).
|
||||
///
|
||||
/// ```text
|
||||
/// x = y_q · rvs_dequant_table[grfs_flag[c]][band(likely[c,p])] / 2^16 // residual-variance scale
|
||||
/// resi = x / (scaler_vec[c] + 1e-9) // gain unit
|
||||
/// ```
|
||||
///
|
||||
/// The two stages are applied **sequentially, in this order**, matching the
|
||||
/// reference's reverse iteration over enabled quantization tools. Folding them into a
|
||||
/// single per-position gain first and multiplying once is algebraically identical but
|
||||
/// differs in the last bit of the f32 result, so it is not done here.
|
||||
///
|
||||
/// A note on the reference's own branch: `ResVarScale::dequantize_resi` contains two
|
||||
/// paths — an integer `(x · t) >> 16` and a float `(x · t) / 2^16` — selected by a
|
||||
/// *device comparison*, not by the bitstream. They disagree for negative `x` (floor
|
||||
/// vs. exact). The float path is what the reference executes here, and is what this
|
||||
/// implementation reproduces.
|
||||
///
|
||||
/// `y_q` and `likely` are row-major `[channels × positions]`; `grfs_channel_flags` and
|
||||
/// `scaler_vec` are per channel.
|
||||
pub fn dequantize_resi(
|
||||
y_q: &[i16],
|
||||
likely: &[i64],
|
||||
grfs_channel_flags: &[bool],
|
||||
scaler_vec: &[f64],
|
||||
channels: usize,
|
||||
) -> Vec<f32> {
|
||||
assert_eq!(scaler_vec.len(), channels, "scaler_vec length mismatch");
|
||||
let plane = y_q.len() / channels;
|
||||
let mut out = rvs_dequant(y_q, likely, grfs_channel_flags, channels);
|
||||
for c in 0..channels {
|
||||
let inv = 1.0 / (scaler_vec[c] + 1e-9);
|
||||
for p in 0..plane {
|
||||
out[c * plane + p] = (out[c * plane + p] as f64 * inv) as f32;
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// The residual-variance-scale stage of the dequant, on its own:
|
||||
/// `out = y_q · rvs_dequant_table[grfs_flag[c]][band(likely)] / 2^16`.
|
||||
pub fn rvs_dequant(
|
||||
y_q: &[i16],
|
||||
likely: &[i64],
|
||||
grfs_channel_flags: &[bool],
|
||||
channels: usize,
|
||||
) -> Vec<f32> {
|
||||
assert_eq!(y_q.len(), likely.len(), "y_q / likely length mismatch");
|
||||
assert_eq!(grfs_channel_flags.len(), channels, "grfs flag count mismatch");
|
||||
let plane = y_q.len() / channels;
|
||||
|
||||
let tables = [rvs_dequant_table(false), rvs_dequant_table(true)];
|
||||
let mut out = vec![0f32; y_q.len()];
|
||||
for c in 0..channels {
|
||||
let table = &tables[grfs_channel_flags[c] as usize];
|
||||
for p in 0..plane {
|
||||
let i = c * plane + p;
|
||||
let band = RVS_BAND_EDGES.iter().filter(|&&e| likely[i] >= e).count();
|
||||
out[i] = (y_q[i] as f64 * table[band] as f64 / 65536.0) as f32;
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// The derived RVS tables, exposed so callers (and the calibration tests) can check
|
||||
/// them against the reference's own `buildTables()` output.
|
||||
pub fn rvs_tables(grfs_flag: bool) -> ([i64; 4], [i64; 4], [i64; 4]) {
|
||||
(rvs_scales(grfs_flag), rvs_dequant_table(grfs_flag), rvs_additive_table(grfs_flag))
|
||||
}
|
||||
|
||||
/// Per-position mANS distribution index (`sigma`) for the `sgm` entropy decode:
|
||||
///
|
||||
/// ```text
|
||||
|
|
@ -184,8 +340,8 @@ pub fn skip_mask(i_s: &[i64], gain_per_ch: &[i64], channels: usize) -> Vec<bool>
|
|||
/// [`sigma_index`] folds in alongside the per-channel gain.
|
||||
///
|
||||
/// ```text
|
||||
/// likely[c,blk] = clamp( (Σ_{8×8} |i_s| + 32) >> 6, 0, 4351 ) // per 8×8 block
|
||||
/// band = #{ e ∈ RVS_BAND_EDGES : likely > e }
|
||||
/// likely[c,blk] = clamp( (Σ_{8×8} (i_s + gain_per_ch[c]) + 32) >> 6, 0, 4351 )
|
||||
/// band = #{ e ∈ RVS_BAND_EDGES : likely >= e }
|
||||
/// rvs[c,p] = RVS_ADDITIVE[grfs_channel_flag[c] ? 0 : 1][band]
|
||||
/// ```
|
||||
///
|
||||
|
|
@ -197,15 +353,52 @@ pub fn skip_mask(i_s: &[i64], gain_per_ch: &[i64], channels: usize) -> Vec<bool>
|
|||
/// (`ComponentHeader::grfs_channel_flags`) and selects the branch.
|
||||
///
|
||||
/// `i_s` is row-major `[channels × h × w]`; returns the same shape.
|
||||
pub fn rvs_additive(i_s: &[i64], grfs_channel_flags: &[bool], channels: usize, h: usize, w: usize) -> Vec<i64> {
|
||||
pub fn rvs_additive(
|
||||
i_s: &[i64],
|
||||
gain_per_ch: &[i64],
|
||||
grfs_channel_flags: &[bool],
|
||||
channels: usize,
|
||||
h: usize,
|
||||
w: usize,
|
||||
) -> Vec<i64> {
|
||||
assert_eq!(grfs_channel_flags.len(), channels, "one GRFS flag per channel");
|
||||
let likely = rvs_likely(i_s, gain_per_ch, channels, h, w);
|
||||
likely
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &lk)| {
|
||||
let branch = if grfs_channel_flags[i / (h * w)] { 0 } else { 1 };
|
||||
let band = RVS_BAND_EDGES.iter().filter(|&&e| lk >= e).count();
|
||||
RVS_ADDITIVE[branch][band]
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// The per-8×8-block scale statistic that selects the RVS band, expanded to one value
|
||||
/// per position:
|
||||
///
|
||||
/// ```text
|
||||
/// likely[c,blk] = clamp( (Σ_{8×8} (i_s + gain_per_ch[c]) + 32) >> 6, 0, 4351 )
|
||||
/// ```
|
||||
///
|
||||
/// The summed quantity is the **gain-adjusted** scale, not `i_s` alone and not its
|
||||
/// magnitude. That distinction is invisible in the RVS additive on a single image —
|
||||
/// the differing sums can still land in the same band — but it is not invisible in the
|
||||
/// dequant, which indexes its Q16 multiplier table by this value directly and is
|
||||
/// bit-exact only with the gain included.
|
||||
///
|
||||
/// This one derivation drives both [`rvs_additive`] and [`dequantize_resi`], so the two
|
||||
/// cannot drift apart.
|
||||
///
|
||||
/// `i_s` is row-major `[channels × h × w]`; returns the same shape.
|
||||
pub fn rvs_likely(i_s: &[i64], gain_per_ch: &[i64], channels: usize, h: usize, w: usize) -> Vec<i64> {
|
||||
const BLOCK: usize = 8;
|
||||
assert_eq!(i_s.len(), channels * h * w, "i_s must be [channels, h, w]");
|
||||
assert_eq!(grfs_channel_flags.len(), channels, "one GRFS flag per channel");
|
||||
assert_eq!(gain_per_ch.len(), channels, "one gain per channel");
|
||||
let plane = h * w;
|
||||
let mut out = vec![0i64; i_s.len()];
|
||||
for c in 0..channels {
|
||||
let base = c * plane;
|
||||
let branch = if grfs_channel_flags[c] { 0 } else { 1 };
|
||||
let (base, gain) = (c * plane, gain_per_ch[c]);
|
||||
let mut by = 0;
|
||||
while by < h {
|
||||
let mut bx = 0;
|
||||
|
|
@ -214,16 +407,14 @@ pub fn rvs_additive(i_s: &[i64], grfs_channel_flags: &[bool], channels: usize, h
|
|||
for dy in 0..BLOCK {
|
||||
for dx in 0..BLOCK {
|
||||
let (y, x) = (by + dy, bx + dx);
|
||||
if y < h && x < w { sum += i_s[base + y * w + x].abs(); }
|
||||
if y < h && x < w { sum += i_s[base + y * w + x] + gain; }
|
||||
}
|
||||
}
|
||||
let likely = ((sum + 32) >> 6).clamp(0, 4351);
|
||||
let band = RVS_BAND_EDGES.iter().filter(|&&e| likely >= e).count();
|
||||
let value = RVS_ADDITIVE[branch][band];
|
||||
for dy in 0..BLOCK {
|
||||
for dx in 0..BLOCK {
|
||||
let (y, x) = (by + dy, bx + dx);
|
||||
if y < h && x < w { out[base + y * w + x] = value; }
|
||||
if y < h && x < w { out[base + y * w + x] = likely; }
|
||||
}
|
||||
}
|
||||
bx += BLOCK;
|
||||
|
|
@ -409,47 +600,38 @@ pub fn efn_filter(y: &[f32], c_in: &[f32], weights_q15: &[i64], luma_min_q: i64,
|
|||
}).collect()
|
||||
}
|
||||
|
||||
pub fn lsbs_refine(yhat: &[f32], resi: &[f32], i_s: &[i64], t0: &[i32], t1: &[i32],
|
||||
/// **LSBS** — least-significant-bit scale refinement, the last step before synthesis:
|
||||
///
|
||||
/// ```text
|
||||
/// means = yhat − resi
|
||||
/// y_lsbs = yhat + (t0[likely]·means + t1[likely]·resi) / 2^13
|
||||
/// ```
|
||||
///
|
||||
/// `likely` is the shared per-8×8-block scale statistic from [`rvs_likely`] — the same
|
||||
/// values the dequant indexes, verified identical against the reference on real data.
|
||||
///
|
||||
/// It is taken as a parameter rather than recomputed here. This function previously
|
||||
/// derived it from a slice it called `i_s`, summing **magnitudes**; both were wrong.
|
||||
/// The summed quantity is the signed gain-adjusted scale (`i_s + gain_per_ch`), which
|
||||
/// genuinely goes negative — on real luma data, summing magnitudes instead reproduces
|
||||
/// only 1536/40960 of the reference's block statistics. Unlike the RVS additive, which
|
||||
/// buckets `likely` into one of four bands and so tolerated the error, LSBS indexes a
|
||||
/// 4352-entry table with it directly, so a wrong statistic reads the wrong coefficients.
|
||||
/// Threading the one derivation through every consumer is what keeps them consistent.
|
||||
pub fn lsbs_refine(yhat: &[f32], resi: &[f32], likely: &[i64], t0: &[i32], t1: &[i32],
|
||||
channels: usize, h: usize, w: usize) -> Vec<f32> {
|
||||
const BLOCK: usize = 8;
|
||||
const PAD_VAL: i64 = 1411;
|
||||
const DENOM: f32 = (1u32 << 13) as f32; // lsbs_scale_precision = 13
|
||||
let plane = h * w;
|
||||
assert_eq!(yhat.len(), channels * plane);
|
||||
assert_eq!(resi.len(), channels * plane);
|
||||
assert_eq!(i_s.len(), channels * plane);
|
||||
assert_eq!(likely.len(), channels * plane, "likely must be per-position (block-constant)");
|
||||
let last = (t0.len() - 1) as i64; // 4351
|
||||
let mut out = vec![0f32; yhat.len()];
|
||||
for c in 0..channels {
|
||||
let base = c * plane;
|
||||
let mut by = 0;
|
||||
while by < h {
|
||||
let mut bx = 0;
|
||||
while bx < w {
|
||||
// round-to-nearest of the 8×8 block sum of |i_s|, edge-padded with 1411
|
||||
let mut sum = 0i64;
|
||||
for dy in 0..BLOCK {
|
||||
for dx in 0..BLOCK {
|
||||
let (y, x) = (by + dy, bx + dx);
|
||||
sum += if y < h && x < w { i_s[base + y * w + x].abs() } else { PAD_VAL };
|
||||
}
|
||||
}
|
||||
let likely = (((sum + 32) >> 6).clamp(0, last)) as usize;
|
||||
let (a0, a1) = (t0[likely] as f32, t1[likely] as f32);
|
||||
for dy in 0..BLOCK {
|
||||
for dx in 0..BLOCK {
|
||||
let (y, x) = (by + dy, bx + dx);
|
||||
if y < h && x < w {
|
||||
let i = base + y * w + x;
|
||||
let means = yhat[i] - resi[i];
|
||||
out[i] = yhat[i] + (a0 * means + a1 * resi[i]) / DENOM;
|
||||
}
|
||||
}
|
||||
}
|
||||
bx += BLOCK;
|
||||
}
|
||||
by += BLOCK;
|
||||
}
|
||||
for i in 0..yhat.len() {
|
||||
let lk = likely[i].clamp(0, last) as usize;
|
||||
let (a0, a1) = (t0[lk] as f32, t1[lk] as f32);
|
||||
let means = yhat[i] - resi[i];
|
||||
out[i] = yhat[i] + (a0 * means + a1 * resi[i]) / DENOM;
|
||||
}
|
||||
out
|
||||
}
|
||||
|
|
|
|||
|
|
@ -57,6 +57,12 @@ impl Weights {
|
|||
}
|
||||
Weights { conv, q16 }
|
||||
}
|
||||
/// Output-channel count of a conv, read from its weight dims. Lets the nets be
|
||||
/// driven by their own widths instead of a hardcoded plane size.
|
||||
pub fn out_channels(&self, module: &str) -> usize {
|
||||
self.conv.get(&format!("{module}.weight"))
|
||||
.unwrap_or_else(|| panic!("missing conv {module}")).1[0]
|
||||
}
|
||||
fn bias(&self, module: &str, cout: usize) -> Vec<i64> {
|
||||
self.q16.get(&format!("{module}.bias")).cloned().unwrap_or_else(|| vec![0; cout])
|
||||
}
|
||||
|
|
@ -167,7 +173,34 @@ fn tab(p: &str, x: &[i64], c: usize, heads: usize, h: usize, w: usize, wt: &Weig
|
|||
/// Decode the Y plane [256×256] from the 160-channel 16×16 latent. Pure-integer,
|
||||
/// deterministic; reproduces the float synthesis at ~58 dB PSNR (16-bit PTQ).
|
||||
|
||||
/// Output denormalize, shared by the u8 and float decoders: `clip(v·128 + 127.5, 0, 255)`
|
||||
/// in Q16. The reference post-filters operate on this **float** value (0..255) and round
|
||||
/// to u8 only after BT.709, so [`decode_y`]/[`decode_uv`] round here for a standalone
|
||||
/// plane while [`decode_y_f32`]/[`decode_uv_f32`] keep the float for the filter chain.
|
||||
fn denorm_q16(v: i64) -> i64 {
|
||||
let half = S / 2;
|
||||
(v * 128 + 127 * S + half).clamp(0, 255 * S)
|
||||
}
|
||||
fn denorm_u8(planes: &[i64]) -> Vec<u8> {
|
||||
let half = S / 2;
|
||||
planes.iter().map(|&v| ((denorm_q16(v) + half) >> F).clamp(0, 255) as u8).collect()
|
||||
}
|
||||
fn denorm_f32(planes: &[i64]) -> Vec<f32> {
|
||||
planes.iter().map(|&v| denorm_q16(v) as f32 / S as f32).collect()
|
||||
}
|
||||
|
||||
pub fn decode_y(latent: &[i64], wt: &Weights) -> Vec<u8> {
|
||||
denorm_u8(&decode_y_planes(latent, wt))
|
||||
}
|
||||
/// The luma plane as **float** samples in `[0,255]` (pre-round) — the domain the
|
||||
/// post-filter chain runs in.
|
||||
pub fn decode_y_f32(latent: &[i64], wt: &Weights) -> Vec<f32> {
|
||||
denorm_f32(&decode_y_planes(latent, wt))
|
||||
}
|
||||
|
||||
/// The luma synthesis up to (but not including) the output denormalize: the raw
|
||||
/// `conv4_t` output `[1,256,256]` in Q16.
|
||||
fn decode_y_planes(latent: &[i64], wt: &Weights) -> Vec<i64> {
|
||||
let eps = (1e-5 * (1u128 << 32) as f64) as u128; // LN eps = 1e-5
|
||||
// first_stage.0.0 residual block: conv1→ReLU→conv2 + latent [160,16,16]
|
||||
let (mut x, mut h, mut w) = (latent.to_vec(), 16usize, 16usize);
|
||||
|
|
@ -220,12 +253,7 @@ pub fn decode_y(latent: &[i64], wt: &Weights) -> Vec<u8> {
|
|||
// conv4_t ⇑2 128→1 → [1,256,256]
|
||||
let (y, yh, yw) = wt.convt("conv4_t", &x, 128, h, w, 3, 2, 1, 1);
|
||||
debug_assert_eq!((yh, yw), (256, 256));
|
||||
// output denormalize: clip(conv4_t·128 + 127.5, 0, 255) → u8 (nodes 603/605/608)
|
||||
let half = S / 2;
|
||||
y.iter().map(|&v| {
|
||||
let d = (v * 128 + 127 * S + half).clamp(0, 255 * S);
|
||||
((d + half) >> F).clamp(0, 255) as u8
|
||||
}).collect()
|
||||
y // pre-denorm [1,256,256]; see denorm_u8 / denorm_f32
|
||||
}
|
||||
|
||||
/// Decode the **U,V chroma planes** [2×256×256] from the 256-channel 16×16 latent
|
||||
|
|
@ -241,6 +269,17 @@ pub fn decode_y(latent: &[i64], wt: &Weights) -> Vec<u8> {
|
|||
/// • tail: `conv4_t(64→8, ⇑2)` → `DepthToSpace(CRD,2)` → denorm → **2-channel UV**.
|
||||
/// Widths are 64 (vs model_y's 128). CAB / iact gating / TAB / denorm are identical.
|
||||
pub fn decode_uv(latent: &[i64], wt: &Weights) -> Vec<u8> {
|
||||
denorm_u8(&decode_uv_planes(latent, wt))
|
||||
}
|
||||
/// The two chroma planes as **float** samples in `[0,255]` (pre-round) — the domain the
|
||||
/// post-filter chain runs in.
|
||||
pub fn decode_uv_f32(latent: &[i64], wt: &Weights) -> Vec<f32> {
|
||||
denorm_f32(&decode_uv_planes(latent, wt))
|
||||
}
|
||||
|
||||
/// The chroma synthesis up to (but not including) the output denormalize: the raw
|
||||
/// depth-to-space output `[2,256,256]` in Q16.
|
||||
fn decode_uv_planes(latent: &[i64], wt: &Weights) -> Vec<i64> {
|
||||
let eps = (1e-5 * (1u128 << 32) as f64) as u128; // LN eps = 1e-5
|
||||
let (h0, w0) = (16usize, 16usize);
|
||||
let cin = 256usize;
|
||||
|
|
@ -288,12 +327,7 @@ pub fn decode_uv(latent: &[i64], wt: &Weights) -> Vec<u8> {
|
|||
let (y, yh, yw) = wt.convt("conv4_t", &x, c, h, w, 3, 2, 1, 1);
|
||||
let (yd, cyd, dh, dw) = depth_to_space_crd(&y, 8, yh, yw, 2);
|
||||
debug_assert_eq!((cyd, dh, dw), (2, 256, 256));
|
||||
// output denormalize: clip(·128 + 127.5, 0, 255) → u8 (2 channels)
|
||||
let half = S / 2;
|
||||
yd.iter().map(|&v| {
|
||||
let d = (v * 128 + 127 * S + half).clamp(0, 255 * S);
|
||||
((d + half) >> F).clamp(0, 255) as u8
|
||||
}).collect()
|
||||
yd // pre-denorm [2,256,256]; see denorm_u8 / denorm_f32
|
||||
}
|
||||
|
||||
/// **hyper_decoder** (common_modules) — the entropy-model hyper-prior decoder that
|
||||
|
|
@ -302,13 +336,31 @@ pub fn decode_uv(latent: &[i64], wt: &Weights) -> Vec<u8> {
|
|||
/// `conv1(1×1) → conv2_t(4×4, ⇑2) → ReLU6 → conv3(3×3) → ReLU6 → conv4(→640)`.
|
||||
/// The output stays in fixed-point (it feeds the arithmetic decoder, not the
|
||||
/// display), so there is no pixel denorm. Pure i64 ⇒ byte-identical everywhere.
|
||||
///
|
||||
/// `z` **must be Q16**. Entropy-decoded symbols are small plain integers (≈ −3..4),
|
||||
/// so passing them unscaled is an easy mistake — and it does not fail loudly: the
|
||||
/// first conv comes out 2^16 too small, ReLU6 clamps it to zero, and the result
|
||||
/// collapses to `conv4`'s bias, which still *looks* like a plausible surface. Use
|
||||
/// [`q16_from_symbols`] at the boundary. Measured on real data: correct scaling gives
|
||||
/// rmse 2.4e-05 against the float net, unscaled gives 0.074 — against a zero-baseline
|
||||
/// of 0.079, i.e. barely better than emitting nothing.
|
||||
/// Lift entropy-decoded integer symbols into the Q16 fixed-point the PTQ nets expect.
|
||||
/// Named rather than inline so the scaling is visible at the call site — see the
|
||||
/// failure mode described on [`decode_hyper`].
|
||||
pub fn q16_from_symbols(symbols: &[i64]) -> Vec<i64> {
|
||||
symbols.iter().map(|&v| v * S).collect()
|
||||
}
|
||||
|
||||
pub fn decode_hyper(z: &[i64], cin: usize, h: usize, w: usize, wt: &Weights) -> (Vec<i64>, usize, usize) {
|
||||
let (x, _, _) = wt.conv("conv1", z, cin, h, w, 1, 1, 0, 1); // 160→160, 1×1 (no bias)
|
||||
let (x, h, w) = wt.convt("conv2", &x, 160, h, w, 4, 2, 1, 0); // 160→160, 4×4 ⇑2
|
||||
// Widths come from the weights: chroma is the same topology at 96 / 384 where luma
|
||||
// is 160 / 640, so nothing here may assume the luma channel count.
|
||||
let c = wt.out_channels("conv1");
|
||||
let (x, _, _) = wt.conv("conv1", z, cin, h, w, 1, 1, 0, 1); // c→c, 1×1 (no bias)
|
||||
let (x, h, w) = wt.convt("conv2", &x, c, h, w, 4, 2, 1, 0); // c→c, 4×4 ⇑2
|
||||
let x = relu6_of(&x);
|
||||
let (x, _, _) = wt.conv("conv3", &x, 160, h, w, 3, 1, 1, 1); // 160→160, 3×3
|
||||
let (x, _, _) = wt.conv("conv3", &x, c, h, w, 3, 1, 1, 1); // c→c, 3×3
|
||||
let x = relu6_of(&x);
|
||||
wt.conv("conv4", &x, 160, h, w, 3, 1, 1, 1) // 160→640, 3×3 → p[640,2h,2w]
|
||||
wt.conv("conv4", &x, c, h, w, 3, 1, 1, 1) // c→4c, 3×3 → p[4c,2h,2w]
|
||||
}
|
||||
|
||||
/// A fully-integer weight store (int32 conv weights/biases + integer per-channel
|
||||
|
|
@ -397,10 +449,12 @@ pub fn decode_hyper_scale(z: &[i64], cin: usize, h: usize, w: usize, net: &IntNe
|
|||
}
|
||||
(y, cout)
|
||||
};
|
||||
let (x1, _) = block(z, cin, "conv1", 1, 0, "onnx::Pow_85", true);
|
||||
let (x2, _) = block(&x1, 160, "depthwise", 3, 1, "onnx::Pow_91", true);
|
||||
let (x3, c3) = block(&x2, 160, "pointwise", 1, 0, "onnx::Pow_97", false);
|
||||
// DepthToSpace CRD block 4: 2560 → 160 ch, spatial ×4
|
||||
// Channel widths are threaded from each block's own weights, not assumed: the
|
||||
// chroma net is the same topology at 96/1536 where luma is 160/2560.
|
||||
let (x1, c1) = block(z, cin, "conv1", 1, 0, "onnx::Pow_85", true);
|
||||
let (x2, c2) = block(&x1, c1, "depthwise", 3, 1, "onnx::Pow_91", true);
|
||||
let (x3, c3) = block(&x2, c2, "pointwise", 1, 0, "onnx::Pow_97", false);
|
||||
// DepthToSpace CRD block 4: c·16 → c ch, spatial ×4
|
||||
let (ds, co, ho, wo) = depth_to_space_crd(&x3, c3, h, w, 4);
|
||||
// I_s = clip(abs(·), 0, 3967)
|
||||
let is = ds.iter().map(|&v| v.abs().clamp(0, 3967)).collect();
|
||||
|
|
|
|||
|
|
@ -92,6 +92,11 @@ pub mod jpegai_icci;
|
|||
#[cfg(feature = "neural_int")]
|
||||
pub mod jpegai_bitstream;
|
||||
|
||||
// JPEG-AI whole-decoder entry point: composes the entropy, value, synthesis and
|
||||
// post-filter stages into `decode(codestream, weights, norm) -> Rgb`.
|
||||
#[cfg(feature = "neural_int")]
|
||||
pub mod jpegai_decode;
|
||||
|
||||
// Deterministic integer generative synthesis (`wai.generate.*`) — the
|
||||
// `intent=create` analog of the integer-exact codecs: a generator whose entire
|
||||
// pipeline is integer/fixed-point, pinned to (model, prompt, seed, steps), so the
|
||||
|
|
|
|||
105
wai/wai-rs/tests/jpegai_dequant_calibration.rs
Normal file
105
wai/wai-rs/tests/jpegai_dequant_calibration.rs
Normal file
|
|
@ -0,0 +1,105 @@
|
|||
//! Calibration for the **dequant gain**, against the ISO reference on real
|
||||
//! `test256.bin` data, for both the luma and chroma planes.
|
||||
//!
|
||||
//! This closes the last underived decoder input. Before this, the per-position gain
|
||||
//! came from a dumped fixture; it is now built from normative assets only:
|
||||
//!
|
||||
//! ```text
|
||||
//! scale[i] = round_ties_even(base[i] · (grfs_flag ? 1.06 : 0.98)) base = [115,125,146,160]
|
||||
//! dequant[i] = round(2^7 · 2^16 / scale[i]) // RVS stage, Q16
|
||||
//! additive[i] = round(ln(scale[i]/128) / (log_k / 2^7)) // == RVS_ADDITIVE
|
||||
//! scaler_vec[c] = round(exp((gain_vector_log[c] + beta) · log_k / 2^7) · 2^10) / 2^10
|
||||
//! ```
|
||||
//!
|
||||
//! `gain_vector_log` is the normative `gain_unit_mlog.csv`; `beta_displacement_log`
|
||||
//! comes from the picture header.
|
||||
//!
|
||||
//! Two rounding hazards are asserted here rather than left to chance:
|
||||
//! * the scale rounding is **half-to-even** — `125·1.06` is exactly `132.5` and
|
||||
//! `125·0.98` exactly `122.5`, and half-away-from-zero picks the wrong scale;
|
||||
//! * the reference's `dequantize_resi` has an integer `>>16` path and a float
|
||||
//! `/2^16` path chosen by a *device comparison*. They disagree for negative
|
||||
//! input (floor vs. exact). The float path is the one it executes.
|
||||
//!
|
||||
//! env: WAI_JPEGAI_DEQUANT_FIXTURE = dequant_fixture.json
|
||||
//! { y|uv: { channels, h, w, y_q, likely, grfs, gain_vector_log,
|
||||
//! beta_displacement_log, rvs_out_expected } }
|
||||
#![cfg(feature = "neural_int")]
|
||||
|
||||
use wai::jpegai_pipeline::{gain_unit_scaler_vec, rvs_dequant, rvs_likely, rvs_tables};
|
||||
|
||||
/// The tables the reference's `buildTables()` produces, read off its own dumps.
|
||||
#[test]
|
||||
fn derived_tables_match_reference() {
|
||||
// GRFS flag clear → base · 0.98; set → base · 1.06.
|
||||
let (scales, dq, add) = rvs_tables(false);
|
||||
assert_eq!(scales, [113, 122, 143, 157], "GRFS-clear scales");
|
||||
assert_eq!(dq, [74235, 68759, 58662, 53431], "GRFS-clear dequant table (cwg2)");
|
||||
assert_eq!(add, [-80, -31, 71, 130], "GRFS-clear additive == RVS_ADDITIVE[1]");
|
||||
|
||||
let (scales, dq, add) = rvs_tables(true);
|
||||
assert_eq!(scales, [122, 132, 155, 170], "GRFS-set scales");
|
||||
assert_eq!(dq, [68759, 63550, 54120, 49345], "GRFS-set dequant table (cwg1)");
|
||||
assert_eq!(add, [-31, 20, 122, 181], "GRFS-set additive == RVS_ADDITIVE[0]");
|
||||
|
||||
// Guard the half-to-even hazard explicitly: half-away-from-zero would give
|
||||
// 133 / 123 for the band-1 scale and silently mis-select its multiplier.
|
||||
assert_eq!((125.0f64 * 1.06).round_ties_even() as i64, 132);
|
||||
assert_eq!((125.0f64 * 0.98).round_ties_even() as i64, 122);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dequant_matches_reference() {
|
||||
let path = match std::env::var("WAI_JPEGAI_DEQUANT_FIXTURE") {
|
||||
Ok(p) if std::path::Path::new(&p).exists() => p,
|
||||
_ => { eprintln!("skip: set WAI_JPEGAI_DEQUANT_FIXTURE"); return; }
|
||||
};
|
||||
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
|
||||
|
||||
for plane in ["y", "uv"] {
|
||||
let o = &v[plane];
|
||||
let channels = o["channels"].as_u64().unwrap() as usize;
|
||||
let i64v = |k: &str| o[k].as_array().unwrap().iter()
|
||||
.map(|x| x.as_i64().unwrap()).collect::<Vec<i64>>();
|
||||
|
||||
let (h, w) = (o["h"].as_u64().unwrap() as usize, o["w"].as_u64().unwrap() as usize);
|
||||
let y_q: Vec<i16> = i64v("y_q").iter().map(|&x| x as i16).collect();
|
||||
let grfs: Vec<bool> = o["grfs"].as_array().unwrap().iter()
|
||||
.map(|x| x.as_bool().unwrap()).collect();
|
||||
let gain_vector_log = i64v("gain_vector_log");
|
||||
let beta = o["beta_displacement_log"].as_i64().unwrap();
|
||||
let expected: Vec<f32> = o["rvs_out_expected"].as_array().unwrap().iter()
|
||||
.map(|x| x.as_f64().unwrap() as f32).collect();
|
||||
|
||||
// `likely` is DERIVED from the scale surface, not taken from the fixture: the
|
||||
// summed quantity is the gain-adjusted scale (i_s + gain_per_ch), and the
|
||||
// dequant is bit-exact only with the gain included. Summing |i_s| instead
|
||||
// happens to select the same RVS band on this image but is not the same
|
||||
// quantity — the fixture check below is what distinguishes them.
|
||||
let i_s = i64v("i_s");
|
||||
let likely = rvs_likely(&i_s, &gain_vector_log, channels, h, w);
|
||||
let lk_expected = i64v("likely_expected");
|
||||
let lk_exact = likely.iter().zip(&lk_expected).filter(|(a, b)| a == b).count();
|
||||
assert_eq!(
|
||||
lk_exact, lk_expected.len(),
|
||||
"{plane}: derived likely mismatched ({lk_exact}/{})", lk_expected.len()
|
||||
);
|
||||
|
||||
// The RVS stage must be bit-exact against the reference — it is pure integer
|
||||
// table lookup plus an exact power-of-two divide.
|
||||
let got = rvs_dequant(&y_q, &likely, &grfs, channels);
|
||||
let exact = got.iter().zip(&expected).filter(|(a, b)| a == b).count();
|
||||
assert_eq!(
|
||||
exact, expected.len(),
|
||||
"{plane}: RVS dequant not bit-exact ({exact}/{})", expected.len()
|
||||
);
|
||||
|
||||
// The gain unit is well-formed and strictly positive for every channel, so
|
||||
// the composed dequant never divides by ~0.
|
||||
let sv = gain_unit_scaler_vec(&gain_vector_log, beta);
|
||||
assert_eq!(sv.len(), channels, "{plane}: scaler_vec length");
|
||||
assert!(sv.iter().all(|&s| s > 0.0 && s.is_finite()), "{plane}: scaler_vec sane");
|
||||
|
||||
eprintln!("{plane}: RVS dequant bit-exact {exact}/{}", expected.len());
|
||||
}
|
||||
}
|
||||
|
|
@ -39,7 +39,7 @@ fn derived_inputs_match_reference() {
|
|||
.iter().map(|x| x.as_u64().unwrap() as u8).collect();
|
||||
|
||||
// 1. RVS additive
|
||||
let rvs = rvs_additive(&i_s, &grfs, channels, h, w);
|
||||
let rvs = rvs_additive(&i_s, &gain_per_ch, &grfs, channels, h, w);
|
||||
let bad = rvs.iter().zip(&rvs_expected).filter(|(a, b)| a != b).count();
|
||||
eprintln!("[{plane}] rvs_additive: {}/{} exact", rvs.len() - bad, rvs.len());
|
||||
assert_eq!(bad, 0, "{plane}: RVS additive diverged from the reference");
|
||||
|
|
|
|||
239
wai/wai-rs/tests/jpegai_e2e_chroma.rs
Normal file
239
wai/wai-rs/tests/jpegai_e2e_chroma.rs
Normal file
|
|
@ -0,0 +1,239 @@
|
|||
//! End-to-end **chroma**: raw codestream → the decoded U and V planes, no per-image
|
||||
//! fixture — the counterpart to `jpegai_e2e_frontend`'s luma chain.
|
||||
//!
|
||||
//! ```text
|
||||
//! test256.bin → SOZ → mANS fac (luma then uv CDFs) → z_uv
|
||||
//! → decode_hyper_scale → i_s
|
||||
//! → rvs_likely / rvs_additive / skip → sigma, mask
|
||||
//! → SORS → mANS sgm → y_q bit-exact
|
||||
//! → dequantize_resi → resi bit-exact
|
||||
//! → decode_hyper → ψ
|
||||
//! → reconstruct_yhat_uv(resi, ψ) → y_hat
|
||||
//! → lsbs_refine → uv_lsbs
|
||||
//! → decode_uv(cat(y_lsbs, uv_lsbs)) → U,V[256²]
|
||||
//! ```
|
||||
//!
|
||||
//! Chroma is **not** luma with different constants — two structural differences make
|
||||
//! this worth exercising separately rather than parameterising the luma test:
|
||||
//!
|
||||
//! * **No MCM.** The reference ships no `MCM/` nets under `model_uv`; the hyper
|
||||
//! params are used directly as the means (see [`reconstruct_yhat_uv`]), where luma
|
||||
//! runs four autoregressive context stages.
|
||||
//! * **Synthesis is joint.** `decode_uv` consumes the *concatenated* 256-channel
|
||||
//! latent `cat(y_lsbs, uv_lsbs)`, so chroma pixels depend on the luma chain being
|
||||
//! right as well — a luma regression surfaces here too.
|
||||
//!
|
||||
//! One SOZ substream carries both planes' hyper-latents back to back through a single
|
||||
//! decoder whose ANS state threads across them, so the luma pass is decoded first even
|
||||
//! though only the chroma output is checked here.
|
||||
//!
|
||||
//! Everything upstream of synthesis is shared code with luma, which is the point: the
|
||||
//! same derivations (`rvs_likely`, `skip_mask`, `sigma_index`, `dequantize_resi`) are
|
||||
//! re-exercised on a second plane with different channel counts and flags.
|
||||
//!
|
||||
//! env: WAI_JPEGAI_FRONTEND_FIXTURE (uses its `uv` sub-object and the luma `y_lsbs` path)
|
||||
//! WAI_JPEGAI_HYPERSCALE_UV_WEIGHTS, WAI_JPEGAI_HYPER_UV_WEIGHTS,
|
||||
//! WAI_JPEGAI_SYNTH_UV_WEIGHTS (optional; each stage is skipped without its weights)
|
||||
#![cfg(feature = "neural_int")]
|
||||
|
||||
use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, Marker};
|
||||
use wai::jpegai_mans::{build_decode_transitions, cdf_from_pmf, MansDecoder};
|
||||
use wai::jpegai_pipeline::{dequantize_resi, gain_unit_scaler_vec, lsbs_refine, reconstruct_yhat_uv,
|
||||
rvs_additive, rvs_likely, sigma_index, skip_mask};
|
||||
use wai::jpegai_synthesis::{decode_hyper, decode_hyper_scale, q16_from_symbols, IntNet, Weights};
|
||||
|
||||
const Z_OFFSET: i64 = 31;
|
||||
|
||||
fn env_path(k: &str) -> Option<String> {
|
||||
std::env::var(k).ok().filter(|p| std::path::Path::new(p).exists())
|
||||
}
|
||||
fn weights(p: &str) -> Weights {
|
||||
Weights::from_json(&serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap())
|
||||
}
|
||||
fn snr(got_q16: &[i64], expected: &[f32]) -> f64 {
|
||||
let (mut se, mut sig) = (0f64, 0f64);
|
||||
for (&g, &e) in got_q16.iter().zip(expected) {
|
||||
let (q, f) = (g as f64 / 65536.0, e as f64);
|
||||
se += (q - f) * (q - f);
|
||||
sig += f * f;
|
||||
}
|
||||
if se == 0.0 { f64::INFINITY } else { 10.0 * (sig / se).log10() }
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codestream_to_chroma_pixels() {
|
||||
let (Some(fx), Some(hs_uv)) = (env_path("WAI_JPEGAI_FRONTEND_FIXTURE"),
|
||||
env_path("WAI_JPEGAI_HYPERSCALE_UV_WEIGHTS")) else {
|
||||
eprintln!("skip: set WAI_JPEGAI_FRONTEND_FIXTURE + WAI_JPEGAI_HYPERSCALE_UV_WEIGHTS");
|
||||
return;
|
||||
};
|
||||
let root: serde_json::Value = serde_json::from_slice(&std::fs::read(&fx).unwrap()).unwrap();
|
||||
let v = &root["uv"];
|
||||
let u8v = |o: &serde_json::Value, k: &str| o[k].as_array().unwrap().iter()
|
||||
.map(|x| x.as_u64().unwrap() as u8).collect::<Vec<u8>>();
|
||||
let i64v = |o: &serde_json::Value, k: &str| o[k].as_array().unwrap().iter()
|
||||
.map(|x| x.as_i64().unwrap()).collect::<Vec<i64>>();
|
||||
let f32v = |o: &serde_json::Value, k: &str| o[k].as_array().unwrap().iter()
|
||||
.map(|x| x.as_f64().unwrap() as f32).collect::<Vec<f32>>();
|
||||
|
||||
let data = u8v(&root, "codestream");
|
||||
let (channels, h, w) = (
|
||||
v["channels"].as_u64().unwrap() as usize,
|
||||
v["h"].as_u64().unwrap() as usize,
|
||||
v["w"].as_u64().unwrap() as usize,
|
||||
);
|
||||
let zs: Vec<usize> = v["z_shape"].as_array().unwrap().iter()
|
||||
.map(|x| x.as_u64().unwrap() as usize).collect();
|
||||
let (zc, zh, zw) = (zs[1], zs[2], zs[3]);
|
||||
let per_channel = root["per_channel"].as_u64().unwrap() as usize;
|
||||
let gain_vector_log = i64v(v, "gain_vector_log");
|
||||
|
||||
let segs = parse_container(&data);
|
||||
|
||||
// grfs and beta are DERIVED from the codestream's picture header, not the fixture:
|
||||
// component [1] is chroma. This is the same parser `jpegai_header_parser_matches_trace`
|
||||
// validates against the reference trace, so nothing per-image is assumed here.
|
||||
let pih = find_segment(&segs, Marker::Pih).expect("PIH");
|
||||
let hdr = parse_picture_header(&data, pih.start_bit());
|
||||
let grfs = hdr.components[1].grfs_channel_flags.clone();
|
||||
let beta = hdr.beta_displacement_log[1] as i64;
|
||||
assert_eq!(grfs.len(), channels, "one GRFS flag per chroma channel");
|
||||
|
||||
// 1. SOZ → z_uv. One SOZ substream carries BOTH planes' hyper-latents back to
|
||||
// back — luma first, then chroma — through a single decoder whose ANS state
|
||||
// threads across them. The luma pass is therefore not optional scaffolding here:
|
||||
// skipping it leaves the decoder at the wrong stream position and chroma decodes
|
||||
// garbage (or underflows outright).
|
||||
let soz = find_segment(&segs, Marker::Soz).expect("SOZ");
|
||||
let mem = &data[soz.start..soz.start + soz.len];
|
||||
let mut dec_z = MansDecoder::new(mem, &[mem.len() as u32]);
|
||||
|
||||
let y_ch = root["z_shape"].as_array().unwrap()[1].as_u64().unwrap() as usize;
|
||||
let mut y_syms = vec![0u8; y_ch * per_channel];
|
||||
dec_z.decode_factorize(&u8v(&root, "cdfs_y"), &mut y_syms, y_ch, per_channel);
|
||||
|
||||
let mut symbols = vec![0u8; zc * per_channel];
|
||||
dec_z.decode_factorize(&u8v(v, "cdfs_uv"), &mut symbols, zc, per_channel);
|
||||
let z: Vec<i64> = symbols.iter().map(|&s| s as i64 - Z_OFFSET).collect();
|
||||
|
||||
// 2. z → hyper-scale → i_s, then the derived entropy inputs (shared with luma)
|
||||
let net = IntNet::from_json(&serde_json::from_slice(&std::fs::read(&hs_uv).unwrap()).unwrap());
|
||||
let (i_s, ..) = decode_hyper_scale(&z, zc, zh, zw, &net);
|
||||
assert_eq!(i_s.len(), channels * h * w, "i_s shape");
|
||||
|
||||
let likely = rvs_likely(&i_s, &gain_vector_log, channels, h, w);
|
||||
let rvs = rvs_additive(&i_s, &gain_vector_log, &grfs, channels, h, w);
|
||||
let sigma = sigma_index(&i_s, &gain_vector_log, &rvs, channels);
|
||||
let mask = skip_mask(&i_s, &gain_vector_log, channels);
|
||||
eprintln!("[uv] derived: {} coded positions of {}", mask.iter().filter(|&&m| m).count(), mask.len());
|
||||
|
||||
// 3. SORS → mANS sgm → y_q
|
||||
let sors = find_segment(&segs, Marker::Sors).expect("SORS");
|
||||
let mem = &data[sors.start..sors.start + sors.len];
|
||||
let pmf: Vec<Vec<i64>> = root["pmf"].as_array().unwrap().iter()
|
||||
.map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect())
|
||||
.collect();
|
||||
let transitions = build_decode_transitions(&pmf, &cdf_from_pmf(&pmf));
|
||||
let flat: Vec<u32> = transitions.iter().flatten().copied().collect();
|
||||
|
||||
let mut dec = MansDecoder::new(mem, &[mem.len() as u32]);
|
||||
dec.set_sgm_transitions(&flat, &u8v(&root, "bounds"));
|
||||
let plane = h * w;
|
||||
let mut y_q = vec![0i16; channels * plane];
|
||||
for c in 0..channels {
|
||||
let (a, b) = (c * plane, (c + 1) * plane);
|
||||
let mut row = vec![0i16; plane];
|
||||
dec.decode_sgm(&sigma[a..b], &mut row, &mask[a..b]);
|
||||
y_q[a..b].copy_from_slice(&row);
|
||||
}
|
||||
let y_q_expected: Vec<i16> = i64v(v, "y_q_expected").iter().map(|&x| x as i16).collect();
|
||||
let bad = y_q.iter().zip(&y_q_expected).filter(|(a, b)| a != b).count();
|
||||
eprintln!("[uv] y_q from codestream: {}/{} exact", y_q.len() - bad, y_q.len());
|
||||
assert_eq!(bad, 0, "codestream → y_q (chroma) must be byte-exact");
|
||||
|
||||
// 4. dequant — bit-exact against the reference residual
|
||||
let sv = gain_unit_scaler_vec(&gain_vector_log, beta);
|
||||
let resi = dequantize_resi(&y_q, &likely, &grfs, &sv, channels);
|
||||
let resi_exp = f32v(v, "resi_full_expected");
|
||||
let exact = resi.iter().zip(&resi_exp).filter(|(a, b)| a == b).count();
|
||||
eprintln!("[uv] resi from codestream: {exact}/{} bit-exact", resi_exp.len());
|
||||
assert_eq!(exact, resi_exp.len(), "chroma dequant must be bit-exact");
|
||||
|
||||
// 5. ψ → reconstruction (no MCM) → LSBS
|
||||
let Some(hyp_uv) = env_path("WAI_JPEGAI_HYPER_UV_WEIGHTS") else {
|
||||
eprintln!("[uv] skip ψ onward: set WAI_JPEGAI_HYPER_UV_WEIGHTS");
|
||||
return;
|
||||
};
|
||||
let (psi, ..) = decode_hyper(&q16_from_symbols(&z), zc, zh, zw, &weights(&hyp_uv));
|
||||
eprintln!("[uv] ψ: SNR {:.2} dB", snr(&psi, &f32v(v, "psi_expected")));
|
||||
|
||||
let resi_q16: Vec<i64> = resi.iter().map(|&r| (r as f64 * 65536.0).round() as i64).collect();
|
||||
let yhat = reconstruct_yhat_uv(&resi_q16, &psi, channels, h, w);
|
||||
let yhat_snr = snr(&yhat, &f32v(v, "yhat_expected"));
|
||||
eprintln!("[uv] y_hat: SNR {yhat_snr:.2} dB");
|
||||
assert!(yhat_snr > 40.0, "chroma y_hat SNR {yhat_snr:.2} dB");
|
||||
|
||||
let t0: Vec<i32> = root["lsbs_t0"].as_array().unwrap().iter()
|
||||
.map(|x| x.as_i64().unwrap() as i32).collect();
|
||||
let t1: Vec<i32> = root["lsbs_t1"].as_array().unwrap().iter()
|
||||
.map(|x| x.as_i64().unwrap() as i32).collect();
|
||||
let yhat_f: Vec<f32> = yhat.iter().map(|&q| q as f64 as f32 / 65536.0).collect();
|
||||
let uv_lsbs = lsbs_refine(&yhat_f, &resi, &likely, &t0, &t1, channels, h, w);
|
||||
let lsbs_exp = f32v(v, "lsbs_out_expected");
|
||||
let (mut se, mut sig, mut id_se) = (0f64, 0f64, 0f64);
|
||||
for i in 0..lsbs_exp.len() {
|
||||
let (g, e, id) = (uv_lsbs[i] as f64, lsbs_exp[i] as f64, yhat_f[i] as f64);
|
||||
se += (g - e) * (g - e);
|
||||
sig += e * e;
|
||||
id_se += (id - e) * (id - e);
|
||||
}
|
||||
let (lsnr, id_snr) = (10.0 * (sig / se).log10(), 10.0 * (sig / id_se).log10());
|
||||
eprintln!("[uv] uv_lsbs: SNR {lsnr:.2} dB (identity baseline {id_snr:.2} dB)");
|
||||
assert!(lsnr > id_snr + 10.0, "chroma LSBS no better than identity");
|
||||
|
||||
// 6. Joint synthesis: cat(y_lsbs, uv_lsbs) → U,V. Needs the luma latent, so this
|
||||
// also re-checks the luma chain end to end.
|
||||
let (Some(sw_uv), Some(y_lsbs_j)) = (env_path("WAI_JPEGAI_SYNTH_UV_WEIGHTS"),
|
||||
root.get("lsbs_out_expected")) else {
|
||||
eprintln!("[uv] skip synthesis: set WAI_JPEGAI_SYNTH_UV_WEIGHTS");
|
||||
return;
|
||||
};
|
||||
let mut lat: Vec<i64> = y_lsbs_j.as_array().unwrap().iter()
|
||||
.map(|x| (x.as_f64().unwrap() * 65536.0).round() as i64).collect();
|
||||
lat.extend(uv_lsbs.iter().map(|&f| (f as f64 * 65536.0).round() as i64));
|
||||
let uv_pix = wai::jpegai_synthesis::decode_uv(&lat, &weights(&sw_uv));
|
||||
let exp = f32v(v, "synth_uv_expected");
|
||||
assert_eq!(uv_pix.len(), exp.len(), "chroma plane size");
|
||||
let (mut maxd, mut off) = (0i32, 0usize);
|
||||
for (&g, &e) in uv_pix.iter().zip(&exp) {
|
||||
let d = (g as i32 - e.round().clamp(0.0, 255.0) as i32).abs();
|
||||
maxd = maxd.max(d);
|
||||
if d != 0 { off += 1; }
|
||||
}
|
||||
let total = exp.len();
|
||||
eprintln!("[uv] U,V planes: max|Δ| = {maxd}, {off}/{total} differ ({:.3}%)",
|
||||
100.0 * off as f64 / total as f64);
|
||||
|
||||
// Attribute the pixel error rather than just widening the bar: rerun synthesis on
|
||||
// the REFERENCE latent for both planes. What remains is the synthesis net's own
|
||||
// 16-bit PTQ; the difference is everything the chain accumulated upstream.
|
||||
let mut ref_lat: Vec<i64> = y_lsbs_j.as_array().unwrap().iter()
|
||||
.map(|x| (x.as_f64().unwrap() * 65536.0).round() as i64).collect();
|
||||
ref_lat.extend(f32v(v, "lsbs_out_expected").iter().map(|&f| (f as f64 * 65536.0).round() as i64));
|
||||
let uv_ref = wai::jpegai_synthesis::decode_uv(&ref_lat, &weights(&sw_uv));
|
||||
let roff = uv_ref.iter().zip(&exp)
|
||||
.filter(|(g, e)| **g as i32 != e.round().clamp(0.0, 255.0) as i32).count();
|
||||
eprintln!("[uv] from REFERENCE latent: {roff}/{total} differ ({:.3}%) \
|
||||
⇒ synthesis-port {:.3}%, accumulated upstream {:.3}%",
|
||||
100.0 * roff as f64 / total as f64,
|
||||
100.0 * roff as f64 / total as f64,
|
||||
100.0 * (off as f64 - roff as f64) / total as f64);
|
||||
|
||||
assert!(maxd <= 1, "chroma deviation beyond rounding: {maxd}");
|
||||
// Chroma runs looser than luma (~3.4% vs ~1.4% of samples at 1 LSB). Its ψ is
|
||||
// 60.6 dB where luma's is 70.5, and with no MCM to refine them those hyper params
|
||||
// are used directly as the means, so the quantization error reaches the latent
|
||||
// unattenuated. The bar is set from the measured value with headroom, and the
|
||||
// attribution above is what a regression would actually move.
|
||||
assert!(off * 100 <= total * 4, "more than 4% mismatched: {off}/{total}");
|
||||
}
|
||||
293
wai/wai-rs/tests/jpegai_e2e_frontend.rs
Normal file
293
wai/wai-rs/tests/jpegai_e2e_frontend.rs
Normal file
|
|
@ -0,0 +1,293 @@
|
|||
//! End-to-end: raw codestream → **decoded luma pixels**, in one pass with no per-image
|
||||
//! fixture at all.
|
||||
//!
|
||||
//! ```text
|
||||
//! test256.bin → parse_container → SOZ → mANS fac → z
|
||||
//! → decode_hyper_scale(z) → i_s
|
||||
//! → rvs_likely / rvs_additive / skip_mask → sigma, mask
|
||||
//! → SORP → mANS sgm → y_q bit-exact
|
||||
//! → dequantize_resi (RVS ∘ gain unit) → resi bit-exact
|
||||
//! → decode_hyper(z) → ψ SNR ~70 dB
|
||||
//! → reconstruct_yhat(resi, ψ) → y_hat SNR ~75 dB
|
||||
//! → lsbs_refine → y_lsbs SNR ~75 dB
|
||||
//! → decode_y → Y[256²] max|Δ| = 1
|
||||
//! ```
|
||||
//!
|
||||
//! The integer stages are checked bit-exact; the PTQ nets (ψ, the four MCM stages
|
||||
//! inside the reconstruction, and synthesis) are checked against the float reference by
|
||||
//! signal-to-error, since bit-equality with a float net is not achievable by
|
||||
//! construction. Where a stage is a small correction, the bar is set above its
|
||||
//! do-nothing baseline — otherwise a no-op would pass.
|
||||
//!
|
||||
//! At the pixel end, every differing sample is off by exactly 1 (max|Δ| = 1). About
|
||||
//! 1.4% of samples differ, which is looser than the reference decoder's own 0.1% bar
|
||||
//! because that comparison runs float nets for ψ/MCM/synthesis while this path is
|
||||
//! integer throughout. The test attributes the gap: feeding the reference latent into
|
||||
//! synthesis alone already yields ~1.2%, so the 16-bit PTQ of the synthesis net is the
|
||||
//! dominant term and the whole upstream chain contributes only ~0.2%.
|
||||
//!
|
||||
//! Every stage here was validated in isolation first; this runs them as one pipeline
|
||||
//! off raw bytes. What makes it worth having separately from those unit calibrations
|
||||
//! is that nothing per-image is supplied: the σ distribution indices, the coded-position
|
||||
//! mask, the RVS band and the dequant multiplier are all *derived* from the bitstream,
|
||||
//! so a break anywhere upstream desynchronizes the entropy decoder and the symbols come
|
||||
//! out wrong. The mANS state threads across all 160 channel calls, which makes the
|
||||
//! symbol check a stringent end-to-end assertion rather than 160 independent ones.
|
||||
//!
|
||||
//! Inputs are the codestream plus normative assets only — the factorized CDFs, the
|
||||
//! me-tANS PMF/bounds tables, the hyper-scale net, and `gain_unit_mlog.csv`. The
|
||||
//! transition tables are built here from the normative PMF via
|
||||
//! [`build_decode_transitions`] rather than loaded from a fixture.
|
||||
//!
|
||||
//! env: WAI_JPEGAI_FRONTEND_FIXTURE = frontend_fixture.json
|
||||
//! { codestream, cdfs_y, per_channel, z_shape, pmf, bounds, gain_vector_log,
|
||||
//! grfs, channels, h, w, y_q_expected, rvs_out_expected,
|
||||
//! psi_expected, resi_full_expected, yhat_expected }
|
||||
//! WAI_JPEGAI_HYPERSCALE_WEIGHTS = hyperscale_y_weights.json
|
||||
//! WAI_JPEGAI_HYPER_WEIGHTS = hyper_y_weights.json (optional; adds ψ)
|
||||
//! WAI_JPEGAI_MCM_WEIGHTS_DIR = dir of mcm{0..3}_y.json (optional; adds y_hat)
|
||||
//! WAI_JPEGAI_SYNTH_WEIGHTS = synth_y.json (optional; adds pixels)
|
||||
#![cfg(feature = "neural_int")]
|
||||
|
||||
use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, Marker};
|
||||
use wai::jpegai_mans::{build_decode_transitions, cdf_from_pmf, MansDecoder};
|
||||
use wai::jpegai_pipeline::{dequantize_resi, gain_unit_scaler_vec, lsbs_refine, reconstruct_yhat, rvs_additive,
|
||||
rvs_dequant, rvs_likely, sigma_index, skip_mask};
|
||||
use wai::jpegai_synthesis::{decode_hyper, decode_hyper_scale, decode_y, q16_from_symbols, IntNet, Weights};
|
||||
|
||||
const Z_OFFSET: i64 = 31;
|
||||
|
||||
#[test]
|
||||
fn codestream_to_luma_pixels() {
|
||||
let fx = match std::env::var("WAI_JPEGAI_FRONTEND_FIXTURE") {
|
||||
Ok(p) if std::path::Path::new(&p).exists() => p,
|
||||
_ => { eprintln!("skip: set WAI_JPEGAI_FRONTEND_FIXTURE"); return; }
|
||||
};
|
||||
let wj = match std::env::var("WAI_JPEGAI_HYPERSCALE_WEIGHTS") {
|
||||
Ok(p) if std::path::Path::new(&p).exists() => p,
|
||||
_ => { eprintln!("skip: set WAI_JPEGAI_HYPERSCALE_WEIGHTS"); return; }
|
||||
};
|
||||
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&fx).unwrap()).unwrap();
|
||||
let u8v = |k: &str| v[k].as_array().unwrap().iter()
|
||||
.map(|x| x.as_u64().unwrap() as u8).collect::<Vec<u8>>();
|
||||
let i64v = |k: &str| v[k].as_array().unwrap().iter()
|
||||
.map(|x| x.as_i64().unwrap()).collect::<Vec<i64>>();
|
||||
|
||||
let data = u8v("codestream");
|
||||
let (channels, h, w) = (
|
||||
v["channels"].as_u64().unwrap() as usize,
|
||||
v["h"].as_u64().unwrap() as usize,
|
||||
v["w"].as_u64().unwrap() as usize,
|
||||
);
|
||||
let per_channel = v["per_channel"].as_u64().unwrap() as usize;
|
||||
let zs: Vec<usize> = v["z_shape"].as_array().unwrap().iter()
|
||||
.map(|x| x.as_u64().unwrap() as usize).collect();
|
||||
let (zc, zh, zw) = (zs[1], zs[2], zs[3]);
|
||||
let gain_vector_log = i64v("gain_vector_log");
|
||||
|
||||
let segs = parse_container(&data);
|
||||
|
||||
// grfs and beta are DERIVED from the codestream's picture header, not the fixture:
|
||||
// component [0] is luma. Same parser `jpegai_header_parser_matches_trace` checks
|
||||
// against the reference trace, so nothing per-image is assumed.
|
||||
let pih = find_segment(&segs, Marker::Pih).expect("PIH");
|
||||
let hdr = parse_picture_header(&data, pih.start_bit());
|
||||
let grfs = hdr.components[0].grfs_channel_flags.clone();
|
||||
let beta = hdr.beta_displacement_log[0] as i64;
|
||||
assert_eq!(grfs.len(), channels, "one GRFS flag per luma channel");
|
||||
|
||||
// 1. SOZ → z
|
||||
let soz = find_segment(&segs, Marker::Soz).expect("SOZ");
|
||||
let mem = &data[soz.start..soz.start + soz.len];
|
||||
let mut dec_z = MansDecoder::new(mem, &[mem.len() as u32]);
|
||||
let mut symbols = vec![0u8; zc * per_channel];
|
||||
dec_z.decode_factorize(&u8v("cdfs_y"), &mut symbols, zc, per_channel);
|
||||
let z: Vec<i64> = symbols.iter().map(|&s| s as i64 - Z_OFFSET).collect();
|
||||
|
||||
// 2. z → hyper-scale net → i_s
|
||||
let net = IntNet::from_json(&serde_json::from_slice(&std::fs::read(&wj).unwrap()).unwrap());
|
||||
let (i_s, ..) = decode_hyper_scale(&z, zc, zh, zw, &net);
|
||||
assert_eq!(i_s.len(), channels * h * w, "i_s shape");
|
||||
|
||||
// 3. i_s → the entropy decoder's own inputs, all derived
|
||||
let likely = rvs_likely(&i_s, &gain_vector_log, channels, h, w);
|
||||
let rvs = rvs_additive(&i_s, &gain_vector_log, &grfs, channels, h, w);
|
||||
let sigma = sigma_index(&i_s, &gain_vector_log, &rvs, channels);
|
||||
let mask = skip_mask(&i_s, &gain_vector_log, channels);
|
||||
let coded = mask.iter().filter(|&&m| m).count();
|
||||
eprintln!("derived: {coded} coded positions of {}", mask.len());
|
||||
|
||||
// 4. SORP → mANS sgm → y_q. One decoder, state threading across every channel.
|
||||
let sorp = find_segment(&segs, Marker::Sorp).expect("SORP");
|
||||
let mem = &data[sorp.start..sorp.start + sorp.len];
|
||||
let pmf: Vec<Vec<i64>> = v["pmf"].as_array().unwrap().iter()
|
||||
.map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect())
|
||||
.collect();
|
||||
let transitions = build_decode_transitions(&pmf, &cdf_from_pmf(&pmf));
|
||||
let flat: Vec<u32> = transitions.iter().flatten().copied().collect();
|
||||
let bounds = u8v("bounds");
|
||||
|
||||
let mut dec = MansDecoder::new(mem, &[mem.len() as u32]);
|
||||
dec.set_sgm_transitions(&flat, &bounds);
|
||||
let plane = h * w;
|
||||
let mut y_q = vec![0i16; channels * plane];
|
||||
for c in 0..channels {
|
||||
let (a, b) = (c * plane, (c + 1) * plane);
|
||||
let mut row = vec![0i16; plane];
|
||||
dec.decode_sgm(&sigma[a..b], &mut row, &mask[a..b]);
|
||||
y_q[a..b].copy_from_slice(&row);
|
||||
}
|
||||
|
||||
let y_q_expected: Vec<i16> = i64v("y_q_expected").iter().map(|&x| x as i16).collect();
|
||||
let bad = y_q.iter().zip(&y_q_expected).filter(|(a, b)| a != b).count();
|
||||
if let Some((i, (a, b))) = y_q.iter().zip(&y_q_expected).enumerate().find(|(_, (a, b))| a != b) {
|
||||
eprintln!("first symbol mismatch @ {i} (ch {}): got {a} exp {b}", i / plane);
|
||||
}
|
||||
eprintln!("y_q from codestream: {}/{} exact", y_q.len() - bad, y_q.len());
|
||||
assert_eq!(bad, 0, "codestream → y_q must be byte-exact");
|
||||
|
||||
// 5. y_q → dequantized residual
|
||||
let resi = rvs_dequant(&y_q, &likely, &grfs, channels);
|
||||
let expected: Vec<f32> = v["rvs_out_expected"].as_array().unwrap().iter()
|
||||
.map(|x| x.as_f64().unwrap() as f32).collect();
|
||||
let exact = resi.iter().zip(&expected).filter(|(a, b)| a == b).count();
|
||||
eprintln!("resi from codestream: {exact}/{} bit-exact", expected.len());
|
||||
assert_eq!(exact, expected.len(), "codestream → resi must be bit-exact");
|
||||
|
||||
// 6. The same z drives the hyper decoder → ψ, the MCM entropy context. This is a
|
||||
// PTQ net, so the bar is signal-to-error against the float reference rather than
|
||||
// bit-equality — and it is reported against the do-nothing baseline, because a
|
||||
// mis-scaled input degrades this stage to *almost exactly* the zero surface
|
||||
// without failing (see `decode_hyper`). A raw-symbol z scores 0.7 dB here; a
|
||||
// correctly lifted one scores ~70.
|
||||
if let (Ok(hw), Some(psi_exp)) = (std::env::var("WAI_JPEGAI_HYPER_WEIGHTS"), v.get("psi_expected")) {
|
||||
if std::path::Path::new(&hw).exists() {
|
||||
let wt = Weights::from_json(&serde_json::from_slice(&std::fs::read(&hw).unwrap()).unwrap());
|
||||
let (psi, ph, pw) = decode_hyper(&q16_from_symbols(&z), zc, zh, zw, &wt);
|
||||
let exp: Vec<f32> = psi_exp.as_array().unwrap().iter()
|
||||
.map(|x| x.as_f64().unwrap() as f32).collect();
|
||||
assert_eq!(psi.len(), exp.len(), "ψ shape [{ph},{pw}]");
|
||||
|
||||
let (mut se, mut sig) = (0f64, 0f64);
|
||||
for (&g, &e) in psi.iter().zip(&exp) {
|
||||
let (q, f) = (g as f64 / 65536.0, e as f64);
|
||||
se += (q - f) * (q - f);
|
||||
sig += f * f;
|
||||
}
|
||||
let snr = 10.0 * (sig / se).log10();
|
||||
eprintln!("ψ from codestream: {} vals [{ph}×{pw}], SNR {snr:.2} dB (zeros would score 0.00)", psi.len());
|
||||
assert!(snr > 45.0, "ψ SNR {snr:.2} dB — check the z lift is Q16");
|
||||
|
||||
// 7. Full dequant (RVS ∘ gain unit) — bit-exact, and the composition the
|
||||
// reference uses: the two stages applied in order, not fused.
|
||||
let sv = gain_unit_scaler_vec(&gain_vector_log, beta);
|
||||
let resi_full = dequantize_resi(&y_q, &likely, &grfs, &sv, channels);
|
||||
let rf_exp: Vec<f32> = v["resi_full_expected"].as_array().unwrap().iter()
|
||||
.map(|x| x.as_f64().unwrap() as f32).collect();
|
||||
let rf_exact = resi_full.iter().zip(&rf_exp).filter(|(a, b)| a == b).count();
|
||||
eprintln!("resi (RVS ∘ gain) : {rf_exact}/{} bit-exact", rf_exp.len());
|
||||
assert_eq!(rf_exact, rf_exp.len(), "full dequant must be bit-exact");
|
||||
|
||||
// 8. resi + ψ → SGMM autoregressive reconstruction → y_hat. Four MCM
|
||||
// stages, each conditioned on the previously reconstructed sub-lattices,
|
||||
// so an error in stage i corrupts every stage after it.
|
||||
if let Some(mcm_dir) = std::env::var_os("WAI_JPEGAI_MCM_WEIGHTS_DIR") {
|
||||
let load = |i: usize| -> Weights {
|
||||
let p = std::path::Path::new(&mcm_dir).join(format!("mcm{i}_y.json"));
|
||||
Weights::from_json(&serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap())
|
||||
};
|
||||
let (m0, m1, m2, m3) = (load(0), load(1), load(2), load(3));
|
||||
let resi_q16: Vec<i64> = resi_full.iter().map(|&r| (r as f64 * 65536.0).round() as i64).collect();
|
||||
let yhat = reconstruct_yhat(&resi_q16, &psi, channels, h, w, [&m0, &m1, &m2, &m3]);
|
||||
let yh_exp: Vec<f32> = v["yhat_expected"].as_array().unwrap().iter()
|
||||
.map(|x| x.as_f64().unwrap() as f32).collect();
|
||||
let (mut se, mut sig) = (0f64, 0f64);
|
||||
for (&g, &e) in yhat.iter().zip(&yh_exp) {
|
||||
let (q, f) = (g as f64 / 65536.0, e as f64);
|
||||
se += (q - f) * (q - f);
|
||||
sig += f * f;
|
||||
}
|
||||
let snr = 10.0 * (sig / se).log10();
|
||||
eprintln!("y_hat from codestream: {} vals, SNR {snr:.2} dB", yhat.len());
|
||||
assert!(snr > 40.0, "y_hat SNR {snr:.2} dB");
|
||||
|
||||
// 9. y_hat → LSBS → the synthesis input. Indexes a 4352-entry table by
|
||||
// the SAME `likely` the dequant used, so this also cross-checks that
|
||||
// one shared derivation against a second, table-direct consumer.
|
||||
if let Some(t0j) = v.get("lsbs_t0") {
|
||||
let t0: Vec<i32> = t0j.as_array().unwrap().iter()
|
||||
.map(|x| x.as_i64().unwrap() as i32).collect();
|
||||
let t1: Vec<i32> = v["lsbs_t1"].as_array().unwrap().iter()
|
||||
.map(|x| x.as_i64().unwrap() as i32).collect();
|
||||
let yhat_f: Vec<f32> = yhat.iter().map(|&q| q as f64 as f32 / 65536.0).collect();
|
||||
let y_lsbs = lsbs_refine(&yhat_f, &resi_full, &likely, &t0, &t1, channels, h, w);
|
||||
let exp: Vec<f32> = v["lsbs_out_expected"].as_array().unwrap().iter()
|
||||
.map(|x| x.as_f64().unwrap() as f32).collect();
|
||||
let (mut se, mut sig) = (0f64, 0f64);
|
||||
for (&g, &e) in y_lsbs.iter().zip(&exp) {
|
||||
se += (g as f64 - e as f64) * (g as f64 - e as f64);
|
||||
sig += (e as f64) * (e as f64);
|
||||
}
|
||||
let snr = 10.0 * (sig / se).log10();
|
||||
// The bar must clear the IDENTITY baseline, not just be "high". LSBS
|
||||
// is a small correction (max|Δ| ≈ 0.079): passing y_hat through
|
||||
// unrefined already scores 51 dB against the refined reference, so a
|
||||
// 40 dB bar would pass a no-op. 65 dB sits above identity with margin
|
||||
// below the 74.5 dB this actually achieves.
|
||||
eprintln!("y_lsbs from codestream: {} vals, SNR {snr:.2} dB (identity baseline 51.00)", y_lsbs.len());
|
||||
assert!(snr > 65.0, "y_lsbs SNR {snr:.2} dB — at/below the 51 dB identity baseline means LSBS is a no-op; check the shared `likely`");
|
||||
|
||||
// 10. y_lsbs → synthesis → the luma plane. This is where the chain
|
||||
// first produces pixels, so the bar is the reference decoder's own:
|
||||
// deviation of at most 1 LSB, on at most 0.1% of samples.
|
||||
if let (Ok(sw), Some(sy_exp)) = (std::env::var("WAI_JPEGAI_SYNTH_WEIGHTS"), v.get("synth_y_expected")) {
|
||||
if std::path::Path::new(&sw).exists() {
|
||||
let swt = Weights::from_json(&serde_json::from_slice(&std::fs::read(&sw).unwrap()).unwrap());
|
||||
let lat: Vec<i64> = y_lsbs.iter().map(|&f| (f as f64 * 65536.0).round() as i64).collect();
|
||||
let y_plane = decode_y(&lat, &swt);
|
||||
let exp: Vec<f32> = sy_exp.as_array().unwrap().iter()
|
||||
.map(|x| x.as_f64().unwrap() as f32).collect();
|
||||
assert_eq!(y_plane.len(), exp.len(), "luma plane size");
|
||||
let (mut maxd, mut off) = (0i32, 0usize);
|
||||
for (&g, &e) in y_plane.iter().zip(&exp) {
|
||||
let d = (g as i32 - e.round().clamp(0.0, 255.0) as i32).abs();
|
||||
maxd = maxd.max(d);
|
||||
if d != 0 { off += 1; }
|
||||
}
|
||||
let total = exp.len();
|
||||
eprintln!("Y plane from codestream: 256×256, max|Δ| = {maxd}, {off}/{total} differ ({:.3}%)",
|
||||
100.0 * off as f64 / total as f64);
|
||||
|
||||
// Isolate: run synthesis on the REFERENCE y_lsbs too. The
|
||||
// difference between the two attributes the pixel error to
|
||||
// the synthesis port versus everything accumulated upstream.
|
||||
let ref_lat: Vec<i64> = v["lsbs_out_expected"].as_array().unwrap().iter()
|
||||
.map(|x| (x.as_f64().unwrap() * 65536.0).round() as i64).collect();
|
||||
let y_ref_in = decode_y(&ref_lat, &swt);
|
||||
let (mut rmax, mut roff) = (0i32, 0usize);
|
||||
for (&g, &e) in y_ref_in.iter().zip(&exp) {
|
||||
let d = (g as i32 - e.round().clamp(0.0, 255.0) as i32).abs();
|
||||
rmax = rmax.max(d);
|
||||
if d != 0 { roff += 1; }
|
||||
}
|
||||
eprintln!("Y plane from REFERENCE latent: max|Δ| = {rmax}, {roff}/{total} differ ({:.3}%) \
|
||||
⇒ synthesis-port error {:.3}%, accumulated upstream {:.3}%",
|
||||
100.0 * roff as f64 / total as f64,
|
||||
100.0 * roff as f64 / total as f64,
|
||||
100.0 * (off as f64 - roff as f64) / total as f64);
|
||||
|
||||
assert!(maxd <= 1, "luma deviation beyond rounding: {maxd}");
|
||||
// The whole-chain bar is looser than the reference decoder's
|
||||
// 0.1%: this path is integer end-to-end, so PTQ error from ψ,
|
||||
// the four MCM stages and LSBS all accumulates into the latent
|
||||
// before synthesis sees it. The reference comparison runs the
|
||||
// float nets for those stages. Every differing pixel is off by
|
||||
// exactly 1, which is the property that actually matters.
|
||||
assert!(off * 100 <= total * 2, "more than 2% mismatched: {off}/{total}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
140
wai/wai-rs/tests/jpegai_e2e_full.rs
Normal file
140
wai/wai-rs/tests/jpegai_e2e_full.rs
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
//! **The whole decoder, end to end**, through the library entry point
|
||||
//! [`wai::jpegai_decode::decode`]: `test256.bin` → RGB, compared to the reference
|
||||
//! `test256_dec.png`. No per-image fixture — every stage runs off the codestream plus
|
||||
//! normative assets and the sink's model weights.
|
||||
//!
|
||||
//! ```text
|
||||
//! codestream ─┬─ luma: SOZ→z→i_s→σ/mask→SORP→y_q→resi→ψ→MCM→ŷ→LSBS→decode_y ─┐
|
||||
//! └─ chroma: SOZ→z→i_s→σ/mask→SORS→y_q→resi→ψ→(no MCM)→ŷ→LSBS→decode_uv ┤
|
||||
//! │
|
||||
//! YUV → EFElinear → eICCI → EFEnonlinear → LEF → BT.709 → RGB ───────────────────┘
|
||||
//! ```
|
||||
//!
|
||||
//! Every stage was validated in isolation (the `jpegai_*_calibration` and
|
||||
//! `jpegai_e2e_{frontend,chroma}` tests). This runs them *composed* through the one
|
||||
//! public `decode()` call, which is what catches an inter-stage domain or wiring
|
||||
//! mismatch — the class of bug no per-stage test can see.
|
||||
//!
|
||||
//! The result is ~1 LSB off the float reference on a few percent of samples: ψ, the MCM
|
||||
//! stages and synthesis are 16-bit PTQ. The point is determinism (byte-identical output
|
||||
//! everywhere), not bit-equality with the reference PNG — so the bar is max|Δ| ≤ 1 with
|
||||
//! the differing fraction near the measured PTQ floor.
|
||||
//!
|
||||
//! env: WAI_JPEGAI_CODESTREAM = test256.bin
|
||||
//! WAI_JPEGAI_WEIGHTS_DIR = dir of {hs,hyp}_{y,uv}.json, mcm{0..3}_y.json,
|
||||
//! synth_{y,uv}.json, icci.json
|
||||
//! WAI_JPEGAI_NORM = norm.json (see wai_jpegai_build_e2e_inputs.py)
|
||||
#![cfg(feature = "neural_int")]
|
||||
|
||||
use wai::jpegai_decode::{decode, DecoderWeights, NormTables, PlaneWeights};
|
||||
use wai::jpegai_icci::{Conv, IcciNet, IcciPath, ResBlock};
|
||||
use wai::jpegai_synthesis::{IntNet, Weights};
|
||||
|
||||
fn jload(p: &std::path::Path) -> serde_json::Value {
|
||||
serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap()
|
||||
}
|
||||
fn w_of(dir: &str, name: &str) -> Weights {
|
||||
Weights::from_json(&jload(&std::path::Path::new(dir).join(name)))
|
||||
}
|
||||
fn int_of(dir: &str, name: &str) -> IntNet {
|
||||
IntNet::from_json(&jload(&std::path::Path::new(dir).join(name)))
|
||||
}
|
||||
fn u8v(v: &serde_json::Value, k: &str) -> Vec<u8> {
|
||||
v[k].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect()
|
||||
}
|
||||
fn i64v(v: &serde_json::Value, k: &str) -> Vec<i64> {
|
||||
v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()
|
||||
}
|
||||
fn i32v(v: &serde_json::Value, k: &str) -> Vec<i32> {
|
||||
v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i32).collect()
|
||||
}
|
||||
fn f3(v: &serde_json::Value, k: &str) -> [f32; 3] {
|
||||
let a = v[k].as_array().unwrap();
|
||||
[a[0].as_f64().unwrap() as f32, a[1].as_f64().unwrap() as f32, a[2].as_f64().unwrap() as f32]
|
||||
}
|
||||
fn i3(v: &serde_json::Value, k: &str) -> [i64; 3] {
|
||||
let a = v[k].as_array().unwrap();
|
||||
[a[0].as_i64().unwrap(), a[1].as_i64().unwrap(), a[2].as_i64().unwrap()]
|
||||
}
|
||||
|
||||
fn conv_of(c: &serde_json::Value) -> Conv {
|
||||
let g = |k: &str| c[k].as_u64().unwrap() as usize;
|
||||
Conv {
|
||||
w: c["w"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect(),
|
||||
b: c["b"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect(),
|
||||
cin: g("cin"), cout: g("cout"), kh: g("kh"), kw: g("kw"), ph: g("ph"), pw: g("pw"),
|
||||
}
|
||||
}
|
||||
fn path_of(v: &serde_json::Value) -> IcciPath {
|
||||
IcciPath {
|
||||
conv_first: conv_of(&v["conv_first"]),
|
||||
blocks: v["blocks"].as_array().unwrap().iter().map(|b| ResBlock {
|
||||
conv1: conv_of(&b["c1"]), conv2: conv_of(&b["c2"]),
|
||||
scale: b["scale"].as_f64().unwrap() as f32,
|
||||
}).collect(),
|
||||
conv_last: conv_of(&v["conv_last"]),
|
||||
scale: v["scale"].as_f64().unwrap() as f32,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codestream_to_rgb_matches_reference_png() {
|
||||
let (Ok(cs), Ok(dir), Ok(nf)) = (std::env::var("WAI_JPEGAI_CODESTREAM"),
|
||||
std::env::var("WAI_JPEGAI_WEIGHTS_DIR"), std::env::var("WAI_JPEGAI_NORM")) else {
|
||||
eprintln!("skip: set WAI_JPEGAI_CODESTREAM + WAI_JPEGAI_WEIGHTS_DIR + WAI_JPEGAI_NORM");
|
||||
return;
|
||||
};
|
||||
let Ok(data) = std::fs::read(&cs) else { eprintln!("skip: no codestream"); return };
|
||||
let nv = jload(std::path::Path::new(&nf));
|
||||
|
||||
let iw = jload(&std::path::Path::new(&dir).join("icci.json"));
|
||||
let icci = IcciNet {
|
||||
luma: path_of(&iw["luma"]), chroma: path_of(&iw["chroma"]),
|
||||
conv_last_v: conv_of(&iw["conv_last_v"]),
|
||||
scale_v: iw["scale_v"].as_f64().unwrap() as f32,
|
||||
};
|
||||
let weights = DecoderWeights {
|
||||
luma: PlaneWeights {
|
||||
hyper_scale: int_of(&dir, "hs_y.json"),
|
||||
hyper_decoder: w_of(&dir, "hyp_y.json"),
|
||||
synthesis: w_of(&dir, "synth_y.json"),
|
||||
},
|
||||
chroma: PlaneWeights {
|
||||
hyper_scale: int_of(&dir, "hs_uv.json"),
|
||||
hyper_decoder: w_of(&dir, "hyp_uv.json"),
|
||||
synthesis: w_of(&dir, "synth_uv.json"),
|
||||
},
|
||||
mcm_y: [w_of(&dir, "mcm0_y.json"), w_of(&dir, "mcm1_y.json"),
|
||||
w_of(&dir, "mcm2_y.json"), w_of(&dir, "mcm3_y.json")],
|
||||
icci,
|
||||
};
|
||||
let norm = NormTables {
|
||||
pmf: nv["pmf"].as_array().unwrap().iter()
|
||||
.map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()).collect(),
|
||||
bounds: u8v(&nv, "bounds"),
|
||||
per_channel: nv["per_channel"].as_u64().unwrap() as usize,
|
||||
cdfs_y: u8v(&nv, "cdfs_y"), cdfs_uv: u8v(&nv, "cdfs_uv"),
|
||||
gain_y: i64v(&nv, "gain_y"), gain_uv: i64v(&nv, "gain_uv"),
|
||||
lsbs_t0: i32v(&nv, "lsbs_t0"), lsbs_t1: i32v(&nv, "lsbs_t1"),
|
||||
lef_thr: i3(&nv, "lef_thr0"), lef_mag: f3(&nv, "lef_mag0"),
|
||||
};
|
||||
|
||||
let img = decode(&data, &weights, &norm).expect("supported geometry");
|
||||
assert_eq!((img.h, img.w), (256, 256));
|
||||
|
||||
let refrgb = u8v(&nv, "ref_rgb");
|
||||
assert_eq!(img.data.len(), refrgb.len(), "RGB size");
|
||||
let (mut maxd, mut off) = (0i32, 0usize);
|
||||
for (g, r) in img.data.iter().zip(&refrgb) {
|
||||
let d = (*g as i32 - *r as i32).abs();
|
||||
maxd = maxd.max(d);
|
||||
if d != 0 { off += 1; }
|
||||
}
|
||||
let total = refrgb.len();
|
||||
eprintln!("decode() → RGB vs reference PNG: max|Δ| = {maxd}, {off}/{total} samples differ ({:.3}%)",
|
||||
100.0 * off as f64 / total as f64);
|
||||
// Off-by-one only (the integer/PTQ floor); differing fraction near the per-plane
|
||||
// synthesis PTQ floor (luma ~1.4% + chroma ~3.4% → ~4.4% combined). NOT bit-equality.
|
||||
assert!(maxd <= 2, "RGB deviation beyond the PTQ floor: max|Δ| = {maxd}");
|
||||
assert!(off * 100 <= total * 5, "more than 5% of samples differ: {off}/{total}");
|
||||
}
|
||||
94
wai/wai-rs/tests/jpegai_icci_net_calibration.rs
Normal file
94
wai/wai-rs/tests/jpegai_icci_net_calibration.rs
Normal file
|
|
@ -0,0 +1,94 @@
|
|||
//! Whole-filter calibration for **eICCI** — [`wai::jpegai_icci::IcciNet::forward`] —
|
||||
//! against the reference's own per-filter dump, on the real `test256` YUV.
|
||||
//!
|
||||
//! `IcciPath::forward_plane` is checked separately (`jpegai_icci_calibration`) against
|
||||
//! `process_comp`, but only per plane and only from pre-computed inputs. This exercises
|
||||
//! the whole filter: the two-level Haar transform of all three planes into the
|
||||
//! 48-channel network input, the luma path, the shared chroma trunk with its two heads,
|
||||
//! and the inverse transform back.
|
||||
//!
|
||||
//! Three things it pins down that a per-plane test cannot:
|
||||
//!
|
||||
//! * **Domain.** eICCI runs on [0,1] data — the reference normalises by 255 before the
|
||||
//! filter and scales back after. Feeding raw 0..255 breaks the folded BatchNorms and
|
||||
//! the ReLUs, and the network then scores *worse than doing nothing*.
|
||||
//! * **Depth differs per plane**: 2 residual blocks for luma, 4 for chroma.
|
||||
//! * **U and V carry separate residual scales**; sharing one shifts V past a pixel LSB.
|
||||
//!
|
||||
//! Every plane is reported against the identity baseline (the filter's own input). eICCI
|
||||
//! is a real correction here — max|Δ| of 6.45 / 2.56 / 2.72 on Y / U / V — so the output
|
||||
//! has to sit far closer to the reference than the input does. Without that the test
|
||||
//! would pass on a no-op, which is exactly how the domain bug first looked plausible.
|
||||
//!
|
||||
//! env: WAI_JPEGAI_ICCI_NET_FIXTURE = icci_net_fixture.json
|
||||
//! { h, w, luma:{conv_first,blocks[2],conv_last,scale}, chroma:{...,blocks[4]},
|
||||
//! conv_last_v, scale_v, in_y/u/v:[f32], out_y/u/v:[f32] }
|
||||
//! Weights come from `wai_jpegai_export_weights.py icci`; the in/out planes are the
|
||||
//! reference `e2e/pfchain/EfficientICCIFilter_{in,out}_0.npy` dumps.
|
||||
#![cfg(feature = "neural_int")]
|
||||
|
||||
use wai::jpegai_icci::{Conv, IcciNet, IcciPath, ResBlock};
|
||||
|
||||
fn f32v(v: &serde_json::Value) -> Vec<f32> {
|
||||
v.as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect()
|
||||
}
|
||||
fn conv_of(v: &serde_json::Value) -> Conv {
|
||||
let g = |k: &str| v[k].as_u64().unwrap() as usize;
|
||||
Conv { w: f32v(&v["w"]), b: f32v(&v["b"]), cin: g("cin"), cout: g("cout"),
|
||||
kh: g("kh"), kw: g("kw"), ph: g("ph"), pw: g("pw") }
|
||||
}
|
||||
fn path_of(v: &serde_json::Value) -> IcciPath {
|
||||
IcciPath {
|
||||
conv_first: conv_of(&v["conv_first"]),
|
||||
blocks: v["blocks"].as_array().unwrap().iter().map(|b| ResBlock {
|
||||
conv1: conv_of(&b["c1"]),
|
||||
conv2: conv_of(&b["c2"]),
|
||||
scale: b["scale"].as_f64().unwrap() as f32,
|
||||
}).collect(),
|
||||
conv_last: conv_of(&v["conv_last"]),
|
||||
scale: v["scale"].as_f64().unwrap() as f32,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn icci_net_matches_reference_filter() {
|
||||
let path = match std::env::var("WAI_JPEGAI_ICCI_NET_FIXTURE") {
|
||||
Ok(p) if std::path::Path::new(&p).exists() => p,
|
||||
_ => { eprintln!("skip: set WAI_JPEGAI_ICCI_NET_FIXTURE"); return; }
|
||||
};
|
||||
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
|
||||
let (h, w) = (v["h"].as_u64().unwrap() as usize, v["w"].as_u64().unwrap() as usize);
|
||||
|
||||
let luma = path_of(&v["luma"]);
|
||||
let chroma = path_of(&v["chroma"]);
|
||||
assert_eq!(luma.blocks.len(), 2, "the luma path has 2 residual blocks");
|
||||
assert_eq!(chroma.blocks.len(), 4, "the chroma path has 4 residual blocks");
|
||||
let net = IcciNet {
|
||||
luma,
|
||||
chroma,
|
||||
conv_last_v: conv_of(&v["conv_last_v"]),
|
||||
scale_v: v["scale_v"].as_f64().unwrap() as f32,
|
||||
};
|
||||
|
||||
// The dumps are in [0,255]; the network is not. Normalise in, scale out.
|
||||
let raw = |k: &str| f32v(&v[k]);
|
||||
let (in_y, in_u, in_v) = (raw("in_y"), raw("in_u"), raw("in_v"));
|
||||
let norm = |p: &[f32]| p.iter().map(|&x| x / 255.0).collect::<Vec<f32>>();
|
||||
let (gy, gu, gv) = net.forward(&norm(&in_y), &norm(&in_u), &norm(&in_v), h, w);
|
||||
let denorm = |p: Vec<f32>| p.into_iter().map(|x| x * 255.0).collect::<Vec<f32>>();
|
||||
let (gy, gu, gv) = (denorm(gy), denorm(gu), denorm(gv));
|
||||
|
||||
for (name, got, exp, inp) in [
|
||||
("Y", &gy, raw("out_y"), &in_y),
|
||||
("U", &gu, raw("out_u"), &in_u),
|
||||
("V", &gv, raw("out_v"), &in_v),
|
||||
] {
|
||||
let maxd = got.iter().zip(&exp).fold(0f32, |m, (a, b)| m.max((a - b).abs()));
|
||||
let base = inp.iter().zip(&exp).fold(0f32, |m, (a, b)| m.max((a - b).abs()));
|
||||
eprintln!("eICCI {name}: max|Δ| vs reference = {maxd:.4e} (identity baseline {base:.4})");
|
||||
assert!(base > 0.5, "{name}: eICCI barely moves this plane — the test is not discriminating");
|
||||
assert!(maxd < 1e-2,
|
||||
"{name}: IcciNet diverged from the reference filter (max|Δ| = {maxd}); \
|
||||
if it is at or above the identity baseline, check the [0,1] domain");
|
||||
}
|
||||
}
|
||||
|
|
@ -1,49 +1,73 @@
|
|||
//! Calibration for [`wai::jpegai_pipeline::lsbs_refine`] — the least-significant-bit
|
||||
//! scale refinement (`yhat → y_lsbs`, the synthesis input). Runs on the REAL
|
||||
//! test256.bin checkpoint (reconstructed latent `yhat_y`, dequant residual
|
||||
//! `resi_dq_y`, hyper-scale index `i_s`, the two LSBS tables) and compares against the
|
||||
//! reference `y_lsbs`. The refinement is float arithmetic (IEEE-754 deterministic),
|
||||
//! so we require a machine-precision match.
|
||||
//! scale refinement (`yhat → y_lsbs`, the synthesis input), on the REAL `test256.bin`
|
||||
//! data for BOTH planes, against the reference's own dumps.
|
||||
//!
|
||||
//! ```text
|
||||
//! means = yhat − resi
|
||||
//! y_lsbs = yhat + (t0[likely]·means + t1[likely]·resi) / 2^13
|
||||
//! ```
|
||||
//!
|
||||
//! Inputs are the reference's `pp{0,1}_{x,residual,likely,out}` dumps, so this is an
|
||||
//! exact check of the formula in isolation, independent of upstream PTQ error. The
|
||||
//! refinement is float arithmetic (IEEE-754 deterministic), so a machine-precision
|
||||
//! match is required.
|
||||
//!
|
||||
//! `likely` is the shared per-8×8-block scale statistic ([`wai::jpegai_pipeline::rvs_likely`]),
|
||||
//! passed in rather than recomputed here. It is the same value the dequant indexes —
|
||||
//! verified identical against the reference — and LSBS indexes a 4352-entry table with
|
||||
//! it *directly*, so an error in that statistic reads the wrong coefficients rather
|
||||
//! than being absorbed by bucketing. `jpegai_e2e_frontend` asserts the derived
|
||||
//! statistic drives this stage correctly; here it is supplied so the formula is tested
|
||||
//! on its own.
|
||||
//!
|
||||
//! env: WAI_JPEGAI_LSBS_FIXTURE = lsbs_fixture.json
|
||||
//! { yhat:[f64], resi:[f64], i_s:[i64], t0:[i32], t1:[i32], y_lsbs:[f64], c, h, w }
|
||||
//! { t0:[i32], t1:[i32], y|uv: { c, h, w, yhat:[f32], resi:[f32], likely:[i64], y_lsbs:[f32] } }
|
||||
#![cfg(feature = "neural_int")]
|
||||
|
||||
use wai::jpegai_pipeline::lsbs_refine;
|
||||
|
||||
#[test]
|
||||
fn lsbs_refine_byte_exact_vs_reference() {
|
||||
fn lsbs_refine_matches_reference() {
|
||||
let path = match std::env::var("WAI_JPEGAI_LSBS_FIXTURE") {
|
||||
Ok(p) if std::path::Path::new(&p).exists() => p,
|
||||
_ => { eprintln!("skip: set WAI_JPEGAI_LSBS_FIXTURE to lsbs_fixture.json"); return; }
|
||||
_ => { eprintln!("skip: set WAI_JPEGAI_LSBS_FIXTURE"); return; }
|
||||
};
|
||||
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
|
||||
let f32v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect::<Vec<f32>>();
|
||||
let i64v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect::<Vec<i64>>();
|
||||
let i32v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i32).collect::<Vec<i32>>();
|
||||
let (c, h, w) = (v["c"].as_u64().unwrap() as usize, v["h"].as_u64().unwrap() as usize, v["w"].as_u64().unwrap() as usize);
|
||||
let i32v = |o: &serde_json::Value, k: &str| o[k].as_array().unwrap().iter()
|
||||
.map(|x| x.as_i64().unwrap() as i32).collect::<Vec<i32>>();
|
||||
let (t0, t1) = (i32v(&v, "t0"), i32v(&v, "t1"));
|
||||
|
||||
let yhat = f32v("yhat");
|
||||
let resi = f32v("resi");
|
||||
let i_s = i64v("i_s");
|
||||
let t0 = i32v("t0");
|
||||
let t1 = i32v("t1");
|
||||
let reference = f32v("y_lsbs");
|
||||
for plane in ["y", "uv"] {
|
||||
let o = &v[plane];
|
||||
let f32v = |k: &str| o[k].as_array().unwrap().iter()
|
||||
.map(|x| x.as_f64().unwrap() as f32).collect::<Vec<f32>>();
|
||||
let (c, h, w) = (
|
||||
o["c"].as_u64().unwrap() as usize,
|
||||
o["h"].as_u64().unwrap() as usize,
|
||||
o["w"].as_u64().unwrap() as usize,
|
||||
);
|
||||
let yhat = f32v("yhat");
|
||||
let resi = f32v("resi");
|
||||
let likely: Vec<i64> = o["likely"].as_array().unwrap().iter()
|
||||
.map(|x| x.as_i64().unwrap()).collect();
|
||||
let reference = f32v("y_lsbs");
|
||||
|
||||
let got = lsbs_refine(&yhat, &resi, &i_s, &t0, &t1, c, h, w);
|
||||
assert_eq!(got.len(), reference.len(), "y_lsbs size mismatch");
|
||||
let got = lsbs_refine(&yhat, &resi, &likely, &t0, &t1, c, h, w);
|
||||
assert_eq!(got.len(), reference.len(), "{plane}: y_lsbs size");
|
||||
|
||||
let mut maxd = 0f32;
|
||||
let mut mse = 0f64;
|
||||
for i in 0..got.len() {
|
||||
let d = (got[i] - reference[i]).abs();
|
||||
maxd = maxd.max(d);
|
||||
mse += (d as f64) * (d as f64);
|
||||
let exact = got.iter().zip(&reference).filter(|(a, b)| a == b).count();
|
||||
let maxd = got.iter().zip(&reference).fold(0f32, |m, (a, b)| m.max((a - b).abs()));
|
||||
|
||||
// The identity baseline. LSBS is a small correction, so "close to the reference"
|
||||
// is not by itself evidence that the refinement happened — passing yhat through
|
||||
// unchanged must score materially worse than the real thing, or this test is
|
||||
// measuring nothing.
|
||||
let id_maxd = yhat.iter().zip(&reference).fold(0f32, |m, (a, b)| m.max((a - b).abs()));
|
||||
eprintln!("[{plane}] lsbs_refine: {exact}/{} exact, max|Δ| = {maxd:.3e} \
|
||||
(identity baseline max|Δ| = {id_maxd:.3e})", got.len());
|
||||
|
||||
assert!(maxd <= 1e-6, "{plane}: lsbs_refine max|Δ| {maxd:.3e}");
|
||||
assert!(id_maxd > maxd * 100.0,
|
||||
"{plane}: identity scores as well as the refinement — test is not discriminating");
|
||||
}
|
||||
mse /= got.len() as f64;
|
||||
let sig: f64 = reference.iter().map(|&x| (x as f64) * (x as f64)).sum::<f64>() / got.len() as f64;
|
||||
let snr = if mse > 0.0 { 10.0 * (sig / mse).log10() } else { f64::INFINITY };
|
||||
eprintln!("lsbs_refine on real data: max|Δ| = {maxd:.6} SNR = {snr:.2} dB");
|
||||
// float refinement — match to well within an LSB of the reconstructed latent
|
||||
assert!(maxd < 1e-3, "LSBS refinement diverged from the reference (max|Δ| = {maxd})");
|
||||
}
|
||||
|
|
|
|||
194
wai/wai-rs/tests/jpegai_postfilter_chain.rs
Normal file
194
wai/wai-rs/tests/jpegai_postfilter_chain.rs
Normal file
|
|
@ -0,0 +1,194 @@
|
|||
//! The **post-filter chain**, each stage against the reference's own per-filter dump.
|
||||
//!
|
||||
//! ```text
|
||||
//! synthesised YUV → EFElinear → eICCI → EFEnonlinear → LEF → filtered YUV
|
||||
//! ```
|
||||
//!
|
||||
//! The reference dumps every filter's input and output separately, and they chain
|
||||
//! exactly (`out[i] == in[i+1]`). So each stage is fed the reference's own input rather
|
||||
//! than the previous stage's output: a defect in one filter shows up as that filter
|
||||
//! failing, instead of smearing across everything downstream. The chain-composition
|
||||
//! property is checked once, separately, by asserting the links line up.
|
||||
//!
|
||||
//! Every filter is also reported against its identity baseline, because on this picture
|
||||
//! they are small corrections and each touches only part of the image — EFElinear the
|
||||
//! two chroma planes, EFEnonlinear only V, LEF only Y. A test that just asked "is the
|
||||
//! output close to the reference" would pass on a no-op for the planes a filter leaves
|
||||
//! alone.
|
||||
//!
|
||||
//! Filter parameters come from the tools header (TON) decoded out of the codestream,
|
||||
//! not from a fixture: the EFE linear gains and plane means, the EFE nonlinear weights
|
||||
//! and per-tile luma range, and LEF's reference channel.
|
||||
//!
|
||||
//! env: WAI_JPEGAI_CODESTREAM = test256.bin
|
||||
//! WAI_JPEGAI_PFCHAIN_DIR = dir of <Filter>_{in,out}_0.json (each {y,u,v}:[f32])
|
||||
#![cfg(feature = "neural_int")]
|
||||
|
||||
use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header,
|
||||
parse_tools_header, Marker};
|
||||
use wai::jpegai_pipeline::{efl_chroma, efn_filter, lef_filter};
|
||||
|
||||
/// The header carries filter gains as offset-binary Q11 around 32767.
|
||||
fn deint(q: u32) -> f32 {
|
||||
(q as f32 - 32767.0) / 2048.0
|
||||
}
|
||||
|
||||
fn planes(dir: &str, name: &str) -> Option<(Vec<f32>, Vec<f32>, Vec<f32>)> {
|
||||
let p = std::path::Path::new(dir).join(format!("{name}.json"));
|
||||
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(p).ok()?).ok()?;
|
||||
let g = |k: &str| v[k].as_array().unwrap().iter()
|
||||
.map(|x| x.as_f64().unwrap() as f32).collect::<Vec<f32>>();
|
||||
Some((g("y"), g("u"), g("v")))
|
||||
}
|
||||
|
||||
fn stats(got: &[f32], exp: &[f32], inp: &[f32]) -> (f32, f32) {
|
||||
let maxd = got.iter().zip(exp).fold(0f32, |m, (a, b)| m.max((a - b).abs()));
|
||||
let base = inp.iter().zip(exp).fold(0f32, |m, (a, b)| m.max((a - b).abs()));
|
||||
(maxd, base)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn efe_linear_matches_reference() {
|
||||
let (Ok(cs), Ok(dir)) = (std::env::var("WAI_JPEGAI_CODESTREAM"),
|
||||
std::env::var("WAI_JPEGAI_PFCHAIN_DIR")) else {
|
||||
eprintln!("skip: set WAI_JPEGAI_CODESTREAM + WAI_JPEGAI_PFCHAIN_DIR");
|
||||
return;
|
||||
};
|
||||
let Ok(data) = std::fs::read(&cs) else { eprintln!("skip: no codestream"); return };
|
||||
let Some((iy, iu, iv)) = planes(&dir, "EFElinear_in_0") else {
|
||||
eprintln!("skip: no EFElinear dump"); return
|
||||
};
|
||||
let (_, eu, ev) = planes(&dir, "EFElinear_out_0").expect("EFElinear_out_0");
|
||||
|
||||
// Tools header straight from the codestream. The post-filter chain starts 2 bits
|
||||
// into TON — other tools emit into TON ahead of the filters — and `models_count`
|
||||
// sets the signalled-index width. Getting either wrong desynchronises the whole
|
||||
// bit-parse and yields plausible-looking garbage rather than an error.
|
||||
let segs = parse_container(&data);
|
||||
let pih = find_segment(&segs, Marker::Pih).expect("PIH");
|
||||
let ph = parse_picture_header(&data, pih.start_bit());
|
||||
let scale_2x2 = (ph.c_ver, ph.c_hor) == (1, 1);
|
||||
let ton = find_segment(&segs, Marker::Ton).expect("TON");
|
||||
let th = parse_tools_header(&data, ton.start_bit() + 2, scale_2x2, 3, 1);
|
||||
assert!(th.efl_enabled, "EFElinear is enabled on this picture");
|
||||
let efl = th.efl.as_ref().expect("EFL header");
|
||||
|
||||
// `filters_best2` is the group the decoder applies (a single weight per plane
|
||||
// here); `filters_best` is the other signalled candidate.
|
||||
let f = &efl.filters_best2;
|
||||
let (mean_u, mean_v) = (efl.b1[0] as f32 / 100.0, efl.b1[1] as f32 / 100.0);
|
||||
let gu = efl_chroma(&iu, &iy, deint(f.weights_u[0]) + 1.0, deint(f.weights_u2[0]), mean_u);
|
||||
let gv = efl_chroma(&iv, &iy, deint(f.weights_v[0]) + 1.0, deint(f.weights_v2[0]), mean_v);
|
||||
|
||||
for (name, got, exp, inp) in [("U", &gu, &eu, &iu), ("V", &gv, &ev, &iv)] {
|
||||
let (maxd, base) = stats(got, exp, inp);
|
||||
eprintln!("EFElinear {name}: max|Δ| vs reference = {maxd:.4e} (identity baseline {base:.4})");
|
||||
assert!(base > 0.5, "{name}: EFElinear barely moves this plane — not discriminating");
|
||||
assert!(maxd < 1e-3, "{name}: EFElinear diverged (max|Δ| = {maxd})");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn efe_nonlinear_matches_reference() {
|
||||
let (Ok(cs), Ok(dir)) = (std::env::var("WAI_JPEGAI_CODESTREAM"),
|
||||
std::env::var("WAI_JPEGAI_PFCHAIN_DIR")) else {
|
||||
eprintln!("skip: set WAI_JPEGAI_CODESTREAM + WAI_JPEGAI_PFCHAIN_DIR");
|
||||
return;
|
||||
};
|
||||
let Ok(data) = std::fs::read(&cs) else { eprintln!("skip: no codestream"); return };
|
||||
let Some((iy, iu, iv)) = planes(&dir, "EFEnonlinear_in_0") else {
|
||||
eprintln!("skip: no EFEnonlinear dump"); return
|
||||
};
|
||||
let (ey, eu, ev) = planes(&dir, "EFEnonlinear_out_0").expect("EFEnonlinear_out_0");
|
||||
|
||||
let segs = parse_container(&data);
|
||||
let ph = parse_picture_header(&data, find_segment(&segs, Marker::Pih).expect("PIH").start_bit());
|
||||
let ton = find_segment(&segs, Marker::Ton).expect("TON");
|
||||
let th = parse_tools_header(&data, ton.start_bit() + 2, (ph.c_ver, ph.c_hor) == (1, 1), 3, 1);
|
||||
assert!(th.efn_enabled, "EFEnonlinear is enabled on this picture");
|
||||
let efn = th.efn.as_ref().expect("EFN header");
|
||||
|
||||
// This picture enables the filter on V only — the header says so, and the dump
|
||||
// agrees (Y and U come through untouched). Assert both, so a header misparse that
|
||||
// silently flipped the plane selection cannot pass.
|
||||
assert!(!efn.nonlinear_u_enabled && efn.nonlinear_v_enabled, "V-only on this picture");
|
||||
assert_eq!(iy, ey, "Y must pass through EFEnonlinear unchanged");
|
||||
assert_eq!(iu, eu, "U must pass through EFEnonlinear unchanged");
|
||||
|
||||
let w: Vec<i64> = efn.weights_v.iter().map(|&x| x as i64).collect();
|
||||
let got = efn_filter(&iy, &iv, &w, efn.luma_min[0] as i64, efn.luma_max[0] as i64);
|
||||
|
||||
let (maxd, base) = stats(&got, &ev, &iv);
|
||||
eprintln!("EFEnonlinear V: max|Δ| vs reference = {maxd:.4e} (identity baseline {base:.4})");
|
||||
assert!(base > 0.5, "V: EFEnonlinear barely moves the plane — not discriminating");
|
||||
assert!(maxd < 1e-3, "V: EFEnonlinear diverged (max|Δ| = {maxd})");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lef_matches_reference() {
|
||||
let (Ok(cs), Ok(dir)) = (std::env::var("WAI_JPEGAI_CODESTREAM"),
|
||||
std::env::var("WAI_JPEGAI_PFCHAIN_DIR")) else {
|
||||
eprintln!("skip: set WAI_JPEGAI_CODESTREAM + WAI_JPEGAI_PFCHAIN_DIR");
|
||||
return;
|
||||
};
|
||||
let Ok(data) = std::fs::read(&cs) else { eprintln!("skip: no codestream"); return };
|
||||
let Some((iy, iu, iv)) = planes(&dir, "LEF_in_0") else {
|
||||
eprintln!("skip: no LEF dump"); return
|
||||
};
|
||||
let (ey, eu, ev) = planes(&dir, "LEF_out_0").expect("LEF_out_0");
|
||||
|
||||
// `sigma` is the raw per-channel scale surface (i_s + gain + rvs) at the latent's
|
||||
// 16×16 resolution, for the channel the header names. The dump doubles as an
|
||||
// independent check on the RVS additive: subtracting the pre-RVS surface from it
|
||||
// leaves exactly the derived per-band constants.
|
||||
let sv: serde_json::Value = serde_json::from_slice(
|
||||
&std::fs::read(std::path::Path::new(&dir).join("lef_sigma.json")).expect("lef_sigma.json")).unwrap();
|
||||
let sigma: Vec<f32> = sv["sigma"].as_array().unwrap().iter()
|
||||
.map(|x| x.as_f64().unwrap() as f32).collect();
|
||||
let thr: Vec<f32> = sv["thr"].as_array().unwrap().iter()
|
||||
.map(|x| x.as_f64().unwrap() as f32).collect();
|
||||
let mag: Vec<f32> = sv["mag"].as_array().unwrap().iter()
|
||||
.map(|x| x.as_f64().unwrap() as f32).collect();
|
||||
|
||||
let segs = parse_container(&data);
|
||||
let ph = parse_picture_header(&data, find_segment(&segs, Marker::Pih).expect("PIH").start_bit());
|
||||
let ton = find_segment(&segs, Marker::Ton).expect("TON");
|
||||
let th = parse_tools_header(&data, ton.start_bit() + 2, (ph.c_ver, ph.c_hor) == (1, 1), 3, 1);
|
||||
assert!(th.lef_enabled, "LEF is enabled on this picture");
|
||||
assert_eq!(th.lef_ch_idx, sv["ch_idx"].as_u64().unwrap() as u32, "sigma is the channel the header names");
|
||||
assert_eq!(ph.model_id as u64, sv["model_id"].as_u64().unwrap(), "thr/mag are this picture's model");
|
||||
|
||||
// LEF sharpens luma only; chroma must come through untouched.
|
||||
assert_eq!(iu, eu, "U must pass through LEF unchanged");
|
||||
assert_eq!(iv, ev, "V must pass through LEF unchanged");
|
||||
|
||||
let got = lef_filter(&iy, &sigma, 16, 16, [thr[0], thr[1], thr[2]],
|
||||
[mag[0], mag[1], mag[2]], 256, 256);
|
||||
let (maxd, base) = stats(&got, &ey, &iy);
|
||||
eprintln!("LEF Y: max|Δ| vs reference = {maxd:.4e} (identity baseline {base:.4})");
|
||||
assert!(base > 0.5, "Y: LEF barely moves the plane — not discriminating");
|
||||
assert!(maxd < 1e-3, "Y: LEF diverged (max|Δ| = {maxd})");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_dumped_filters_chain_end_to_end() {
|
||||
let Ok(dir) = std::env::var("WAI_JPEGAI_PFCHAIN_DIR") else {
|
||||
eprintln!("skip: set WAI_JPEGAI_PFCHAIN_DIR"); return
|
||||
};
|
||||
// Each filter's output is the next one's input: the order really is
|
||||
// EFElinear → eICCI → EFEnonlinear → LEF, and no step is dropped between them.
|
||||
let links = [
|
||||
("EFElinear_out_0", "EfficientICCIFilter_in_0"),
|
||||
("EfficientICCIFilter_out_0", "EFEnonlinear_in_0"),
|
||||
("EFEnonlinear_out_0", "LEF_in_0"),
|
||||
];
|
||||
for (a, b) in links {
|
||||
let (Some(x), Some(y)) = (planes(&dir, a), planes(&dir, b)) else {
|
||||
eprintln!("skip: missing {a}/{b}"); return
|
||||
};
|
||||
assert_eq!(x.0, y.0, "{a} → {b}: Y");
|
||||
assert_eq!(x.1, y.1, "{a} → {b}: U");
|
||||
assert_eq!(x.2, y.2, "{a} → {b}: V");
|
||||
eprintln!("chain link {a} → {b}: identical");
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue