open-standards/wai/wai-pack/src/encode.rs

//! The base encode's choice of step ([packaging] §5.6): the function the packer's base
//! encode of a `wai.video.int_motion` rung calls for each piece.
//!
//! `q_candidates` is read as a set. The piece is encoded at each candidate from the
//! largest down, and the first whose `floor(SSE · 1000 / (3·w·h·n)) ≤
//! max_sse_px_milli` is chosen, `SSE` being its reconstruction's against the source and
//! `n` the piece's frames: the largest step that meets the bound, so candidates that meet
//! it equally are decided by their step. Where none meets it, the smallest candidate is
//! chosen.
//!
//! [packaging]: ../../../extensions/packaging.md

use crate::policy::Rung;
use wai::int_payload::{IntPayloadError, Refusal};
use wai::int_video_enc::{encode, EncParams, Encoded};

/// The parameters of a `wai.video.int_motion` rung that fix its encode.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MotionParams {
    pub block: u32,
    pub search: u32,
    pub mv_precision: u8,
    pub q_candidates: Vec<u32>,
    pub max_sse_px_milli: u64,
}

impl MotionParams {
    /// A rung's parameters, every default written in by the policy (§3.3); `None` for a
    /// rung of another capability.
    pub fn of(rung: &Rung) -> Option<MotionParams> {
        if rung.capability != "wai.video.int_motion" {
            return None;
        }
        Some(MotionParams {
            block: rung.param_u64("block")? as u32,
            search: rung.param_u64("search")? as u32,
            mv_precision: rung.param_u64("mv_precision")? as u8,
            q_candidates: rung.param_u64_list("q_candidates")?.into_iter().map(|q| q as u32).collect(),
            max_sse_px_milli: rung.param_u64("max_sse_px_milli")?,
        })
    }
}

/// A piece's step and its encode.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct QChoice {
    pub q: u32,
    pub encoded: Encoded,
}

/// `floor(sse · 1000 / (3·w·h·n))`: the SSE per sample, in thousandths.
pub fn sse_px_milli(sse: u64, h: u32, w: u32, n: usize) -> u64 {
    let samples = 3 * h as u128 * w as u128 * n as u128;
    (sse as u128 * 1000 / samples.max(1)) as u64
}

/// Choose the step for a piece of `frames` (each `h × w` RGB8) under `p`, and encode it.
pub fn choose_q(frames: &[Vec<u8>], h: u32, w: u32, p: &MotionParams) -> Result<QChoice, IntPayloadError> {
    let mut qs: Vec<u32> = p.q_candidates.clone();
    qs.sort_unstable_by(|a, b| b.cmp(a));
    qs.dedup();
    let smallest = *qs.last().ok_or_else(|| IntPayloadError::new(Refusal::FieldRange, "no step to choose from"))?;
    for &q in &qs {
        let e = encode(frames, h, w, &EncParams { block: p.block, search: p.search, mv_precision: p.mv_precision, q })?;
        if sse_px_milli(e.sse, h, w, frames.len()) <= p.max_sse_px_milli || q == smallest {
            return Ok(QChoice { q, encoded: e });
        }
    }
    unreachable!("the smallest step is always chosen when it is reached")
}

/// The sum, over every sample, of the squared difference between two clips of RGB8
/// frames (§6.6).
pub fn sse(a: &[Vec<u8>], b: &[Vec<u8>]) -> u64 {
    a.iter().zip(b).flat_map(|(x, y)| x.iter().zip(y)).map(|(&x, &y)| (x as i64 - y as i64).unsigned_abs().pow(2)).sum()
}

/// One refinement layer of a piece (§5.4): a `WIR1` of unit kind RGB8, one unit per frame,
/// the identity predictor, one table, at most 4095 table symbols, that brings each frame
/// of `below` (the stage the layer refines) toward its source frame in `src`, at step
/// `q`. This is `wai-rs`'s refinement encoder (`staged::wir1::encode`), the one `wai stage
/// refine` runs: the residual `rdiv(source − below, q)`, its table, its streams. Returns
/// the payload and the frames a sink reconstructs, which the encoder has decoded back.
pub fn encode_refinement(below: &[Vec<u8>], src: &[Vec<u8>], h: u32, w: u32, q: u32) -> Result<(Vec<u8>, Vec<Vec<u8>>), String> {
    use wai::staged::wir1::{encode, Below, Next, Predictor, RefineParams, UnitKind};
    let belows: Vec<Below> = below.iter().map(|f| Below::U8(f)).collect();
    let targets: Vec<Below> = src.iter().map(|f| Below::U8(f)).collect();
    let p = RefineParams { q, predictor: Predictor::Identity, rate: (0, 0), table_per_unit: false, max_symbols: 4095 };
    let (payload, next) = encode(UnitKind::Rgb8, &belows, (h, w, 3), &targets, (h, w, 3), &p).map_err(|e| e.0)?;
    let frames = next
        .into_iter()
        .map(|n| match n {
            Next::U8(v) => Ok(v),
            Next::I16(_) => Err("an RGB8 layer decoded to PCM".to_string()),
        })
        .collect::<Result<Vec<_>, _>>()?;
    Ok((payload, frames))
}

#[cfg(test)]
mod tests {
    use super::*;

    /// A layer brings every frame closer to its source, and a finer step closer still.
    #[test]
    fn a_refinement_reduces_the_error() {
        let src = clip();
        let base = encode(&src, 10, 12, &EncParams { block: 4, search: 1, mv_precision: 1, q: 24 }).unwrap();
        let (p1, r1) = encode_refinement(&base.recon, &src, 10, 12, 8).unwrap();
        let (_, r2) = encode_refinement(&r1, &src, 10, 12, 2).unwrap();
        assert_eq!(&p1[..4], b"WIR1");
        assert_eq!(sse(&base.recon, &src), base.sse);
        assert!(sse(&r1, &src) < base.sse && sse(&r2, &src) < sse(&r1, &src), "{} {} {}", base.sse, sse(&r1, &src), sse(&r2, &src));
        // Step 1 reconstructs the source exactly.
        assert_eq!(encode_refinement(&r1, &src, 10, 12, 1).unwrap().1, src);
    }

    fn clip() -> Vec<Vec<u8>> {
        (0..4).map(|t| (0..12 * 10 * 3).map(|i| ((i * 37 + t * 11) % 251) as u8).collect()).collect()
    }

    fn params(q_candidates: Vec<u32>, max_sse_px_milli: u64) -> MotionParams {
        MotionParams { block: 4, search: 1, mv_precision: 1, q_candidates, max_sse_px_milli }
    }

    fn metric(q: u32) -> u64 {
        let e = encode(&clip(), 10, 12, &EncParams { block: 4, search: 1, mv_precision: 1, q }).unwrap();
        sse_px_milli(e.sse, 10, 12, 4)
    }

    /// The largest step that meets the bound; the smallest when none does; candidates
    /// read as a set, in any order.
    #[test]
    fn the_largest_step_meeting_the_bound() {
        let (m16, m8, m2) = (metric(16), metric(8), metric(2));
        assert!(m16 > m8 && m8 > m2, "{m16} {m8} {m2}");
        let pick = |c: Vec<u32>, bound| choose_q(&clip(), 10, 12, &params(c, bound)).unwrap().q;
        assert_eq!(pick(vec![2, 16, 8], m16), 16);
        assert_eq!(pick(vec![2, 16, 8], m16 - 1), 8);
        assert_eq!(pick(vec![8, 2, 16, 8], m8), 8);
        assert_eq!(pick(vec![16, 8, 2], m2), 2);
        assert_eq!(pick(vec![16, 8, 2], 0), 2, "none meets it: the smallest");
        let c = choose_q(&clip(), 10, 12, &params(vec![16, 8], m8)).unwrap();
        assert_eq!(c.encoded, encode(&clip(), 10, 12, &EncParams { block: 4, search: 1, mv_precision: 1, q: 8 }).unwrap());
    }
}