zoomvtools 1.1.0

Video motion vector analysis utilities in pure Rust
Documentation
use anyhow::{Result, bail, ensure};
use std::num::{NonZeroU8, NonZeroUsize};

use crate::params::{MotionFlags, Subpel};

pub const PROP_MVANALYSISDATA: &str = "MVTools_MVAnalysisData";
pub const PROP_VECTORS: &str = "MVTools_vectors";

#[derive(Debug, Clone, Copy)]
pub struct MVAnalysisData {
    pub blk_size_x: NonZeroUsize,
    pub blk_size_y: NonZeroUsize,
    pub pel: Subpel,
    pub level_count: usize,
    pub delta_frame: isize,
    pub is_backward: bool,
    pub motion_flags: MotionFlags,
    pub width: NonZeroUsize,
    pub height: NonZeroUsize,
    pub overlap_x: usize,
    pub overlap_y: usize,
    pub blk_x: NonZeroUsize,
    pub blk_y: NonZeroUsize,
    pub bits_per_sample: NonZeroU8,
    pub y_ratio_uv: NonZeroU8,
    pub x_ratio_uv: NonZeroU8,
    pub h_padding: usize,
    pub v_padding: usize,
}

impl MVAnalysisData {
    #[must_use]
    #[inline]
    pub fn bytes(&self) -> Vec<u8> {
        let prop_data = MVAnalysisPropData::from(*self);
        // SAFETY: `MVAnalysisPropData` is `repr(C)` and copied as raw bytes.
        unsafe {
            std::slice::from_raw_parts(
                (&prop_data as *const MVAnalysisPropData).cast::<u8>(),
                size_of::<MVAnalysisPropData>(),
            )
            .to_vec()
        }
    }

    #[inline]
    pub fn from_bytes(data: &[u8], filter_name: &str, vector_name: &str) -> Result<Self> {
        if data.len() != size_of::<MVAnalysisPropData>() {
            bail!(
                "{filter_name}: Property {PROP_MVANALYSISDATA} in first frame of {vector_name} has wrong size"
            );
        }

        // SAFETY: Size has been verified above and the representation is `repr(C)`.
        let data = unsafe { std::ptr::read_unaligned(data.as_ptr().cast::<MVAnalysisPropData>()) };
        Self::try_from(data)
    }

    #[inline]
    pub fn scale_thscd(&self, thscd1: &mut u64, thscd2: &mut u64, filter_name: &str) -> Result<()> {
        const MAX_SAD: u64 = 8 * 8 * 255;

        if *thscd1 > MAX_SAD {
            bail!("{filter_name}: thscd1 can be at most {MAX_SAD}");
        }

        const REFERENCE_BLOCK_SIZE: u64 = 8 * 8;
        *thscd1 =
            *thscd1 * (self.blk_size_x.get() * self.blk_size_y.get()) as u64 / REFERENCE_BLOCK_SIZE;
        if self.motion_flags.contains(MotionFlags::USE_CHROMA_MOTION) {
            *thscd1 += *thscd1 / (self.x_ratio_uv.get() * self.y_ratio_uv.get()) as u64 * 2;
        }

        let pixel_max = (1u64 << self.bits_per_sample.get()) - 1;
        *thscd1 = ((*thscd1 as f64) * pixel_max as f64 / 255.0 + 0.5) as u64;
        *thscd2 = *thscd2 * self.blk_x.get() as u64 * self.blk_y.get() as u64 / 256;

        Ok(())
    }

    #[inline]
    pub fn check_similarity(
        &self,
        other: &MVAnalysisData,
        filter_name: &str,
        name1: &str,
        name2: &str,
    ) -> Result<()> {
        ensure!(
            self.width == other.width,
            "{filter_name}: {name1} and {name2} have different widths.",
        );
        ensure!(
            self.height == other.height,
            "{filter_name}: {name1} and {name2} have different heights."
        );
        ensure!(
            self.blk_size_x == other.blk_size_x && self.blk_size_y == other.blk_size_y,
            "{filter_name}: {name1} and {name2} have different block sizes."
        );
        ensure!(
            self.pel == other.pel,
            "{filter_name}: {name1} and {name2} have different pel precision."
        );
        ensure!(
            self.overlap_x == other.overlap_x && self.overlap_y == other.overlap_y,
            "{filter_name}: {name1} and {name2} have different overlap."
        );
        ensure!(
            self.x_ratio_uv == other.x_ratio_uv,
            "{filter_name}: {name1} and {name2} have different horizontal subsampling."
        );
        ensure!(
            self.y_ratio_uv == other.y_ratio_uv,
            "{filter_name}: {name1} and {name2} have different vertical subsampling."
        );
        ensure!(
            self.bits_per_sample == other.bits_per_sample,
            "{filter_name}: {name1} and {name2} have different bit depth."
        );
        ensure!(
            self.h_padding == other.h_padding && self.v_padding == other.v_padding,
            "{filter_name}: {name1} and {name2} have different padding."
        );

        Ok(())
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
struct MVAnalysisPropData {
    magic_key: u32,
    version: u32,
    blk_size_x: u32,
    blk_size_y: u32,
    pel: u32,
    level_count: u32,
    delta_frame: i32,
    is_backward: u32,
    cpu_flags: u32,
    motion_flags: u32,
    width: u32,
    height: u32,
    overlap_x: u32,
    overlap_y: u32,
    blk_x: u32,
    blk_y: u32,
    bits_per_sample: u32,
    y_ratio_uv: u32,
    x_ratio_uv: u32,
    h_padding: u32,
    v_padding: u32,
}

#[allow(clippy::non_zero_suggestions)]
impl From<MVAnalysisData> for MVAnalysisPropData {
    fn from(value: MVAnalysisData) -> Self {
        Self {
            magic_key: 0,
            version: 0,
            blk_size_x: u32::try_from(value.blk_size_x.get()).expect("blk_size_x fits u32"),
            blk_size_y: u32::try_from(value.blk_size_y.get()).expect("blk_size_y fits u32"),
            pel: u32::from(u8::from(value.pel)),
            level_count: u32::try_from(value.level_count).expect("level_count fits u32"),
            delta_frame: i32::try_from(value.delta_frame).expect("delta_frame fits i32"),
            is_backward: u32::from(value.is_backward),
            cpu_flags: 0,
            motion_flags: u32::from(value.motion_flags.bits()),
            width: u32::try_from(value.width.get()).expect("width fits u32"),
            height: u32::try_from(value.height.get()).expect("height fits u32"),
            overlap_x: u32::try_from(value.overlap_x).expect("overlap_x fits u32"),
            overlap_y: u32::try_from(value.overlap_y).expect("overlap_y fits u32"),
            blk_x: u32::try_from(value.blk_x.get()).expect("blk_x fits u32"),
            blk_y: u32::try_from(value.blk_y.get()).expect("blk_y fits u32"),
            bits_per_sample: u32::from(value.bits_per_sample.get()),
            y_ratio_uv: u32::from(value.y_ratio_uv.get()),
            x_ratio_uv: u32::from(value.x_ratio_uv.get()),
            h_padding: u32::try_from(value.h_padding).expect("h_padding fits u32"),
            v_padding: u32::try_from(value.v_padding).expect("v_padding fits u32"),
        }
    }
}

impl TryFrom<MVAnalysisPropData> for MVAnalysisData {
    type Error = anyhow::Error;

    #[inline]
    fn try_from(value: MVAnalysisPropData) -> Result<Self> {
        Ok(Self {
            blk_size_x: NonZeroUsize::new(usize::try_from(value.blk_size_x)?)
                .ok_or_else(|| anyhow::anyhow!("invalid zero blk_size_x"))?,
            blk_size_y: NonZeroUsize::new(usize::try_from(value.blk_size_y)?)
                .ok_or_else(|| anyhow::anyhow!("invalid zero blk_size_y"))?,
            pel: Subpel::try_from(i64::from(value.pel))?,
            level_count: usize::try_from(value.level_count)?,
            delta_frame: isize::try_from(value.delta_frame)?,
            is_backward: value.is_backward != 0,
            motion_flags: MotionFlags::from_bits(u8::try_from(value.motion_flags)?)
                .ok_or_else(|| anyhow::anyhow!("invalid motion flags"))?,
            width: NonZeroUsize::new(usize::try_from(value.width)?)
                .ok_or_else(|| anyhow::anyhow!("invalid zero width"))?,
            height: NonZeroUsize::new(usize::try_from(value.height)?)
                .ok_or_else(|| anyhow::anyhow!("invalid zero height"))?,
            overlap_x: usize::try_from(value.overlap_x)?,
            overlap_y: usize::try_from(value.overlap_y)?,
            blk_x: NonZeroUsize::new(usize::try_from(value.blk_x)?)
                .ok_or_else(|| anyhow::anyhow!("invalid zero blk_x"))?,
            blk_y: NonZeroUsize::new(usize::try_from(value.blk_y)?)
                .ok_or_else(|| anyhow::anyhow!("invalid zero blk_y"))?,
            bits_per_sample: NonZeroU8::new(u8::try_from(value.bits_per_sample)?)
                .ok_or_else(|| anyhow::anyhow!("invalid zero bits_per_sample"))?,
            y_ratio_uv: NonZeroU8::new(u8::try_from(value.y_ratio_uv)?)
                .ok_or_else(|| anyhow::anyhow!("invalid zero y_ratio_uv"))?,
            x_ratio_uv: NonZeroU8::new(u8::try_from(value.x_ratio_uv)?)
                .ok_or_else(|| anyhow::anyhow!("invalid zero x_ratio_uv"))?,
            h_padding: usize::try_from(value.h_padding)?,
            v_padding: usize::try_from(value.v_padding)?,
        })
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::unwrap_used, reason = "allow in test files")]

    use std::num::{NonZeroU8, NonZeroUsize};

    use crate::params::{MotionFlags, Subpel};

    use super::MVAnalysisData;

    fn sample_analysis_data() -> MVAnalysisData {
        MVAnalysisData {
            blk_size_x: NonZeroUsize::new(8).unwrap(),
            blk_size_y: NonZeroUsize::new(8).unwrap(),
            pel: Subpel::Half,
            level_count: 3,
            delta_frame: 1,
            is_backward: true,
            motion_flags: MotionFlags::USE_CHROMA_MOTION,
            width: NonZeroUsize::new(640).unwrap(),
            height: NonZeroUsize::new(480).unwrap(),
            overlap_x: 2,
            overlap_y: 2,
            blk_x: NonZeroUsize::new(80).unwrap(),
            blk_y: NonZeroUsize::new(60).unwrap(),
            bits_per_sample: NonZeroU8::new(10).unwrap(),
            y_ratio_uv: NonZeroU8::new(2).unwrap(),
            x_ratio_uv: NonZeroU8::new(2).unwrap(),
            h_padding: 16,
            v_padding: 16,
        }
    }

    #[test]
    fn bytes_round_trip_through_from_bytes() {
        let data = sample_analysis_data();

        let decoded = MVAnalysisData::from_bytes(&data.bytes(), "Analyse", "vectors").unwrap();

        assert_eq!(decoded.blk_size_x, data.blk_size_x);
        assert_eq!(decoded.blk_size_y, data.blk_size_y);
        assert_eq!(decoded.pel, data.pel);
        assert_eq!(decoded.level_count, data.level_count);
        assert_eq!(decoded.delta_frame, data.delta_frame);
        assert_eq!(decoded.is_backward, data.is_backward);
        assert_eq!(decoded.motion_flags, data.motion_flags);
        assert_eq!(decoded.width, data.width);
        assert_eq!(decoded.height, data.height);
        assert_eq!(decoded.overlap_x, data.overlap_x);
        assert_eq!(decoded.overlap_y, data.overlap_y);
        assert_eq!(decoded.blk_x, data.blk_x);
        assert_eq!(decoded.blk_y, data.blk_y);
        assert_eq!(decoded.bits_per_sample, data.bits_per_sample);
        assert_eq!(decoded.y_ratio_uv, data.y_ratio_uv);
        assert_eq!(decoded.x_ratio_uv, data.x_ratio_uv);
        assert_eq!(decoded.h_padding, data.h_padding);
        assert_eq!(decoded.v_padding, data.v_padding);
    }

    #[test]
    fn from_bytes_rejects_wrong_size() {
        let err = MVAnalysisData::from_bytes(&[0u8; 4], "Analyse", "vectors").unwrap_err();

        assert!(err.to_string().contains("has wrong size"));
    }
}