wai/wai-web/src/neural_int.rs

@@ -172,3 +172,57 @@ pub fn decode_jpegai_mans(
     dec.decode_sgm(sigma, &mut out, &masks_b);
     out
 }
+
+// ---- JPEG-AI full decode → RGB, in the browser (PTQ value/synthesis path) ----
+
+/// Decode a full ISO/IEC 6048 JPEG-AI codestream to RGB in the browser
+/// (`wai.image.jpegai`, decode-equivalence). The same decoder a native sink runs
+/// (`wai::jpegai_decode::decode`), compiled to wasm32: integer entropy front end +
+/// deterministic value/synthesis/post-filter path. One pinned reconstruction,
+/// byte-identical on every machine — so the decoded RGB has a portable hash. The 13
+/// byte inputs are the model weight bundle (`tools/wai_jpegai_build_e2e_inputs.py`):
+/// per-plane hyper-scale / hyper-decoder / synthesis, the four luma MCM stages, the
+/// eICCI net, and the normative tables (`norm`).
+#[wasm_bindgen(js_name = decodeJpegAiFull)]
+#[allow(clippy::too_many_arguments)]
+pub fn decode_jpegai_full(
+    codestream: &[u8],
+    hs_y: &[u8], hyp_y: &[u8], synth_y: &[u8],
+    hs_uv: &[u8], hyp_uv: &[u8], synth_uv: &[u8],
+    mcm0: &[u8], mcm1: &[u8], mcm2: &[u8], mcm3: &[u8],
+    icci: &[u8], norm: &[u8],
+) -> Result<DecodedImage, JsError> {
+    use wai::jpegai_decode::{decode, DecoderWeights, NormTables, PlaneWeights};
+    use wai::jpegai_icci::IcciNet;
+    use wai::jpegai_synthesis::{IntNet, Weights};
+
+    let parse = |b: &[u8]| -> Result<serde_json::Value, JsError> {
+        serde_json::from_slice(b).map_err(|e| JsError::new(&e.to_string()))
+    };
+    let w = |b: &[u8]| -> Result<Weights, JsError> { Ok(Weights::from_json(&parse(b)?)) };
+    let net = |b: &[u8]| -> Result<IntNet, JsError> { Ok(IntNet::from_json(&parse(b)?)) };
+
+    let weights = DecoderWeights {
+        luma: PlaneWeights { hyper_scale: net(hs_y)?, hyper_decoder: w(hyp_y)?, synthesis: w(synth_y)? },
+        chroma: PlaneWeights { hyper_scale: net(hs_uv)?, hyper_decoder: w(hyp_uv)?, synthesis: w(synth_uv)? },
+        mcm_y: [w(mcm0)?, w(mcm1)?, w(mcm2)?, w(mcm3)?],
+        icci: IcciNet::from_json(&parse(icci)?),
+    };
+
+    let nv = parse(norm)?;
+    let i64v = |k: &str| nv[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect::<Vec<i64>>();
+    let u8v = |k: &str| nv[k].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect::<Vec<u8>>();
+    let i32v = |k: &str| nv[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i32).collect::<Vec<i32>>();
+    let i3 = |k: &str| { let a = nv[k].as_array().unwrap(); [a[0].as_i64().unwrap(), a[1].as_i64().unwrap(), a[2].as_i64().unwrap()] };
+    let f3 = |k: &str| { let a = nv[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] };
+    let norm_t = NormTables {
+        pmf: nv["pmf"].as_array().ok_or_else(|| JsError::new("norm.pmf"))?.iter()
+            .map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()).collect(),
+        bounds: u8v("bounds"), per_channel: nv["per_channel"].as_u64().unwrap() as usize,
+        cdfs_y: u8v("cdfs_y"), cdfs_uv: u8v("cdfs_uv"), gain_y: i64v("gain_y"), gain_uv: i64v("gain_uv"),
+        lsbs_t0: i32v("lsbs_t0"), lsbs_t1: i32v("lsbs_t1"), lef_thr: i3("lef_thr0"), lef_mag: f3("lef_mag0"),
+    };
+
+    let img = decode(codestream, &weights, &norm_t).ok_or_else(|| JsError::new("unsupported geometry"))?;
+    Ok(DecodedImage { rgb: img.data, width: img.w as u32, height: img.h as u32 })
+}