nanocodec-cli 0.1.0

//! `nanocodec stat` (analysis builds: `--features analysis`): decode an
//! H.264, HEVC, AV1, VP9, VVC or AV2 stream with coding-tool
//! statistics enabled ([`nc_common::stream_stats`]) and write them as JSON,
//! to show which block sizes, modes, motion, transforms, QPs and filters an
//! encoder actually used. The input is Annex-B or an IVF file of Annex-B
//! access units (H.264, HEVC), Annex-B (VVC) or IVF (AV1, VP9, AV2); `eval/nceval/streamstats.py` summarises the output.

use std::io::{Cursor, Write};
use std::path::PathBuf;

use nc_common::stream_stats::StreamStats;
use nc_encoder::Codec;

use crate::framing::{is_ivf, read_ivf, read_ivf_annex_b};
use crate::stdio;

#[derive(Debug)]
struct StatArgs {
    codec: Codec,
    input: PathBuf,
    json: Option<PathBuf>,
}

fn parse(args: &[String]) -> Result<StatArgs, String> {
    let get = |flag: &str| {
        let i = args.iter().position(|a| a == flag)?;
        args.get(i + 1).cloned()
    };
    let codec = match get("--codec").as_deref() {
        Some("h264") => Codec::H264,
        Some("hevc") => Codec::Hevc,
        Some("av1") => Codec::Av1,
        Some("vp9") => Codec::Vp9,
        Some("vvc") => Codec::Vvc,
        Some("av2") => Codec::Av2,
        _ => return Err("stat needs --codec h264|hevc|av1|vp9|vvc|av2".into()),
    };
    let input = get("--input").ok_or("missing --input")?.into();
    Ok(StatArgs {
        codec,
        input,
        json: get("--json").map(PathBuf::from),
    })
}

/// Run `nanocodec stat`.
pub(super) fn run(args: &[String]) -> Result<(), String> {
    let args = parse(args)?;
    let data = stdio::read_all(&args.input)?;
    let stats = match args.codec {
        Codec::Hevc => hevc(&data)?,
        Codec::H264 => h264(&data)?,
        Codec::Av1 => av1(&data)?,
        Codec::Vp9 => vp9(&data)?,
        Codec::Vvc => vvc(&data)?,
        Codec::Av2 => av2(&data)?,
    };
    let json = stats.to_json();
    match &args.json {
        Some(path) => stdio::write_all(path, json.as_bytes()),
        None => std::io::stdout()
            .write_all(json.as_bytes())
            .map_err(|e| e.to_string()),
    }
}

/// Call `visit` with the stream's Annex-B data: each IVF frame, or the
/// whole input.
fn for_each_packet(
    data: &[u8],
    mut visit: impl FnMut(&[u8]) -> Result<(), String>,
) -> Result<(), String> {
    if is_ivf(data) {
        read_ivf_annex_b(&mut Cursor::new(data), |_, frame| visit(frame.packet))?;
        Ok(())
    } else {
        visit(data)
    }
}

fn hevc(data: &[u8]) -> Result<StreamStats, String> {
    let mut decoder = nc_hevc::Decoder::new();
    decoder.collect_stats(StreamStats::default());
    for_each_packet(data, |packet| {
        decoder.push(packet).map_err(|e| e.to_string())?;
        while decoder.next_output().is_some() {}
        Ok(())
    })?;
    decoder.flush().map_err(|e| e.to_string())?;
    Ok(decoder.take_stats().unwrap_or_default())
}

/// AV1 in IVF: every temporal unit, shown or not.
fn av1(data: &[u8]) -> Result<StreamStats, String> {
    let mut decoder = nc_av1::Decoder::new();
    decoder.collect_stats(StreamStats::default());
    read_ivf(&mut Cursor::new(data), |_, frame| {
        decoder
            .decode_with_layer(frame.packet)
            .map(|_| ())
            .map_err(|e| e.to_string())
    })?;
    Ok(decoder.take_stats().unwrap_or_default())
}

/// VP9 in IVF: every frame, shown or not.
fn vp9(data: &[u8]) -> Result<StreamStats, String> {
    let mut decoder = nc_vp9::stream::Decoder::new();
    decoder.collect_stats(StreamStats::default());
    read_ivf(&mut Cursor::new(data), |_, frame| {
        decoder
            .decode_all_layers(frame.packet)
            .map(|_| ())
            .map_err(|e| e.to_string())
    })?;
    Ok(decoder.take_stats().unwrap_or_default())
}

/// VVC: Annex-B NAL units.
fn vvc(data: &[u8]) -> Result<StreamStats, String> {
    let mut decoder = nc_vvc::Decoder::new();
    decoder.collect_stats(StreamStats::default());
    for_each_packet(data, |packet| {
        for nal in nc_common::decode::annex_b_units(packet).map_err(|e| e.to_string())? {
            decoder.push_nal(nal).map_err(|e| e.to_string())?;
            while decoder.next_output().is_some() {}
        }
        Ok(())
    })?;
    decoder.flush().map_err(|e| e.to_string())?;
    Ok(decoder.take_stats().unwrap_or_default())
}

/// AV2 in IVF: every temporal unit.
fn av2(data: &[u8]) -> Result<StreamStats, String> {
    let mut decoder = nc_av2::stream::Decoder::default();
    decoder.collect_stats(StreamStats::default());
    read_ivf(&mut Cursor::new(data), |_, frame| {
        decoder
            .decode_shown(frame.packet)
            .map_err(|e| e.to_string())
    })?;
    Ok(decoder.take_stats().unwrap_or_default())
}

fn h264(data: &[u8]) -> Result<StreamStats, String> {
    let mut decoder = nc_h264::Decoder::new();
    decoder.collect_stats(StreamStats::default());
    for_each_packet(data, |packet| {
        for nal in nc_common::decode::annex_b_units(packet).map_err(|e| e.to_string())? {
            decoder.decode_nal(nal).map_err(|e| e.to_string())?;
        }
        Ok(())
    })?;
    Ok(decoder.take_stats().unwrap_or_default())
}