48becebcab
8a7492b23 moved ort to 2.0.0-rc.15 and gave ort_session.rs the environment rc.15 asks every session builder for; five test files build sessions of their own (jpegai_e2e and the four calibration tests) and still called `Session::builder()` with none. They build only with `neural` (and `legacy_int`), which no push-triggered job builds, and a checked list cut at ten lines hid them. rc.15 allows one environment per process, and these tests call into wai in the same process, so they take wai's: `ort_session::environment` is public as `wai::ort_environment` (a new path, compatible). All five build and run under neural,semantic,legacy_int; without their JPEG AI fixtures they return early, as before. jpegai_decode_path builds and passes under its features. clippy reports nothing in the files this touches. Changelog: none Co-Authored-By: Claude Opus 5.5 <[email protected]>
85 lines
4 KiB
Rust
85 lines
4 KiB
Rust
//! Calibration: integer MCM stage decoders (`mcm_stage0` / `mcm_stage_n`) vs the
|
|
//! float reference — run the SAME inputs through each `MCM/stage{i}.onnx` (ort) and
|
|
//! through the Rust stage, compare the fixed-point context features `y[i]` (SNR,
|
|
//! since the output is params not pixels). Each stage validated independently.
|
|
//! Needs `neural` (ort).
|
|
//!
|
|
//! env: WAI_JPEGAI_MCM_DIR = dir with stage0..3.onnx,
|
|
//! WAI_JPEGAI_MCM_WEIGHTS_DIR = dir with mcm_stage0..3_y.json
|
|
#![cfg(all(feature = "neural", feature = "legacy_int"))]
|
|
|
|
use ort::session::{builder::GraphOptimizationLevel, Session};
|
|
use ort::value::Value;
|
|
use wai::jpegai_synthesis::{mcm_stage0, mcm_stage_n, Weights};
|
|
|
|
fn arr(c: usize, h: usize, w: usize, seed: i64) -> (Vec<f32>, Vec<i64>) {
|
|
let n = c * h * w;
|
|
let f: Vec<f32> = (0..n).map(|i| ((((i as i64 + seed) * 31) % 197 - 98) as f32) / 64.0).collect();
|
|
let q: Vec<i64> = f.iter().map(|&x| (x * 65536.0).round() as i64).collect();
|
|
(f, q)
|
|
}
|
|
|
|
fn run_stage(path: &str, named: Vec<(&str, Vec<f32>, usize)>, h: usize, w: usize) -> Vec<f32> {
|
|
let mut s = Session::builder(wai::ort_environment().unwrap()).unwrap()
|
|
.with_intra_threads(1).unwrap().with_inter_threads(1).unwrap()
|
|
.with_optimization_level(GraphOptimizationLevel::Level3).unwrap()
|
|
.commit_from_file(path).unwrap();
|
|
let inputs: Vec<(String, Value)> = named.into_iter().map(|(name, data, c)| {
|
|
let a = ndarray::Array4::from_shape_vec((1, c, h, w), data).unwrap();
|
|
(name.to_string(), Value::from_array(a).unwrap().into_dyn())
|
|
}).collect();
|
|
let out = s.run(inputs).unwrap();
|
|
out[0].try_extract_array::<f32>().unwrap().iter().copied().collect()
|
|
}
|
|
|
|
fn snr(int_q16: &[i64], flt: &[f32]) -> (f64, f64) {
|
|
let (mut maxd, mut se, mut sig) = (0f64, 0f64, 0f64);
|
|
for i in 0..int_q16.len() {
|
|
let q = int_q16[i] as f64 / 65536.0;
|
|
let f = flt[i] as f64;
|
|
let d = (q - f).abs();
|
|
maxd = maxd.max(d); se += d * d; sig += f * f;
|
|
}
|
|
(maxd, if se > 0.0 { 10.0 * (sig / se).log10() } else { f64::INFINITY })
|
|
}
|
|
|
|
#[test]
|
|
fn mcm_stages_int_vs_float_snr() {
|
|
let dir = match std::env::var("WAI_JPEGAI_MCM_DIR") {
|
|
Ok(p) if std::path::Path::new(&p).exists() => p,
|
|
_ => { eprintln!("skip: set WAI_JPEGAI_MCM_DIR to the MCM onnx dir"); return; }
|
|
};
|
|
let wdir = match std::env::var("WAI_JPEGAI_MCM_WEIGHTS_DIR") {
|
|
Ok(p) if std::path::Path::new(&p).exists() => p,
|
|
_ => { eprintln!("skip: set WAI_JPEGAI_MCM_WEIGHTS_DIR to the weights dir"); return; }
|
|
};
|
|
let (h, w) = (16usize, 16usize);
|
|
let load = |i: usize| -> Weights {
|
|
let v: serde_json::Value = serde_json::from_slice(&std::fs::read(format!("{wdir}/mcm_stage{i}_y.json")).unwrap()).unwrap();
|
|
Weights::from_json(&v).unwrap()
|
|
};
|
|
let mut worst = f64::INFINITY;
|
|
|
|
// stage 0: single input p0
|
|
{
|
|
let (pf, pq) = arr(160, h, w, 1);
|
|
let rf = run_stage(&format!("{dir}/stage0.onnx"), vec![("p0", pf, 160)], h, w);
|
|
let mine = mcm_stage0(&pq, h, w, &load(0));
|
|
let (md, sn) = snr(&mine, &rf);
|
|
eprintln!("MCM stage0: max|Δ| = {md:.4} SNR = {sn:.2} dB (len {})", mine.len());
|
|
worst = worst.min(sn);
|
|
}
|
|
// stages 1..3: inputs (p_i, context of width 160/320/480)
|
|
for (i, (pn, cn, cc)) in [("p1", "y[0]", 160usize), ("p2", "y[1-0]", 320), ("p3", "y[2-0]", 480)].into_iter().enumerate() {
|
|
let stage = i + 1;
|
|
let (pf, pq) = arr(160, h, w, 7 * stage as i64 + 1);
|
|
let (cf, cq) = arr(cc, h, w, 13 * stage as i64 + 2);
|
|
let rf = run_stage(&format!("{dir}/stage{stage}.onnx"), vec![(pn, pf, 160), (cn, cf, cc)], h, w);
|
|
let mine = mcm_stage_n(&pq, &cq, cc, h, w, &load(stage));
|
|
let (md, sn) = snr(&mine, &rf);
|
|
eprintln!("MCM stage{stage}: max|Δ| = {md:.4} SNR = {sn:.2} dB (len {})", mine.len());
|
|
worst = worst.min(sn);
|
|
}
|
|
eprintln!("MCM worst-stage SNR = {worst:.2} dB");
|
|
assert!(worst > 45.0, "an MCM stage regressed (worst SNR {worst:.2} dB)");
|
|
}
|