open-standards/wai/wai-rs/src/jpegai_decode.rs

//! 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 })
}