open-standards/wai/wai-rs/tests/jpegai_scope_ort.rs

//! The float ONNX path (`wai::neural::jpegai`) accepts exactly the codestreams the
//! exact decoder accepts: both hard-code one configuration (tools_0, β = 0, the
//! high-operation-point synthesis, one eICCI model pair, 256×256 4:4:4), and a
//! codestream outside it is refused with `Unsupported` rather than decoded to the
//! wrong image.
//!
//! env: WAI_JPEGAI_ASSETS = the jpegai-vm-assets root (the ONNX model bundle).
//! Skips without it; no CI job provides it.
#![cfg(feature = "neural")]

use wai::neural::jpegai::{decode, JpegAiError, JpegAiModel};

const TEST256: &[u8] = include_bytes!("vectors/jpegai/test256.bin");

/// Overwrite `n` bits (MSB-first) at absolute bit `pos` with `value`.
fn set_bits(data: &mut [u8], pos: usize, n: usize, value: u64) {
    for i in 0..n {
        let bit = (value >> (n - 1 - i)) & 1;
        let (byte, sh) = ((pos + i) / 8, 7 - (pos + i) % 8);
        data[byte] = (data[byte] & !(1 << sh)) | ((bit as u8) << sh);
    }
}

#[test]
fn ort_path_refuses_what_the_exact_decoder_refuses() {
    let Ok(dir) = std::env::var("WAI_JPEGAI_ASSETS") else {
        eprintln!("skip: set WAI_JPEGAI_ASSETS to the jpegai-vm-assets root");
        return;
    };
    let model = JpegAiModel::load(&dir).expect("load the ONNX model bundle");

    use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, Marker};
    let pih = find_segment(&parse_container(TEST256), Marker::Pih).unwrap().start_bit();
    // (field offset inside the PIH payload, width, new value, word the refusal names)
    for (at, width, value, word) in [(90, 12, 2048 + 7, "beta"), (85, 4, 1, "model_id")] {
        let mut d = TEST256.to_vec();
        set_bits(&mut d, pih + at, width, value);
        let ph = parse_picture_header(&d, pih);
        assert!(ph.beta_displacement_log != [0, 0] || ph.model_id != 0, "the edit hit a field");
        match decode(&d, &model) {
            Err(JpegAiError::Unsupported(why)) => assert!(why.contains(word), "{why}"),
            Err(e) => panic!("expected Unsupported naming {word}, got {e}"),
            Ok(_) => panic!("an out-of-scope codestream ({word}) was decoded"),
        }
    }

    let img = decode(TEST256, &model).expect("the in-scope codestream decodes");
    assert_eq!((img.width, img.height), (256, 256));
}