zoomvtools 2.0.0

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

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

use crate::{
    analysis::MVAnalysisData,
    fake::group_of_planes::FakeGroupOfPlanes,
    frame::{FramePlanes, FramePlanesMut, FrameView, PlaneRef},
    mv::MotionVector,
    params::{MaskKind, MotionFlags, Subpel},
    video::{ColorFamily, Resolution, SampleType, VideoFormat, VideoInfo},
};

use super::{Mask, MaskOptions};

fn yuv444_info() -> VideoInfo {
    VideoInfo {
        format: VideoFormat {
            color_family: ColorFamily::Yuv,
            sample_type: SampleType::Integer,
            bits_per_sample: NonZeroU8::new(8).unwrap(),
            bytes_per_sample: NonZeroU8::new(1).unwrap(),
            sub_sampling_w: 0,
            sub_sampling_h: 0,
        },
        resolution: Resolution {
            width: 2,
            height: 2,
        },
        num_frames: 1,
    }
}

fn analysis_data() -> MVAnalysisData {
    MVAnalysisData {
        blk_size_x: NonZeroUsize::new(1).unwrap(),
        blk_size_y: NonZeroUsize::new(1).unwrap(),
        pel: Subpel::Full,
        level_count: 1,
        delta_frame: 1,
        is_backward: false,
        motion_flags: MotionFlags::empty(),
        width: NonZeroUsize::new(2).unwrap(),
        height: NonZeroUsize::new(2).unwrap(),
        overlap_x: 0,
        overlap_y: 0,
        blk_x: NonZeroUsize::new(2).unwrap(),
        blk_y: NonZeroUsize::new(2).unwrap(),
        bits_per_sample: NonZeroU8::new(8).unwrap(),
        y_ratio_uv: NonZeroU8::new(1).unwrap(),
        x_ratio_uv: NonZeroU8::new(1).unwrap(),
        h_padding: 0,
        v_padding: 0,
    }
}

fn fake_vectors(validity: bool) -> FakeGroupOfPlanes {
    let mut fake_gop = FakeGroupOfPlanes::new(&analysis_data());
    let plane = fake_gop.plane_mut(0);
    let [block0, block1, block2, block3] = plane.blocks.as_mut_slice() else {
        panic!("expected 4 luma blocks");
    };
    block0.vector = MotionVector {
        x: 10,
        y: -20,
        sad: 0,
    };
    block1.vector = MotionVector {
        x: -10,
        y: 20,
        sad: 0,
    };
    block2.vector = MotionVector { x: 5, y: 0, sad: 0 };
    block3.vector = MotionVector {
        x: 0,
        y: -5,
        sad: 0,
    };
    fake_gop.set_validity(validity);
    fake_gop
}

fn source_view() -> FrameView<'static, u8> {
    static Y: [u8; 4] = [9, 11, 13, 15];
    static U: [u8; 4] = [0, 0, 0, 0];
    static V: [u8; 4] = [0, 0, 0, 0];

    FrameView::new(
        FramePlanes::new([
            Some(PlaneRef::new(&Y)),
            Some(PlaneRef::new(&U)),
            Some(PlaneRef::new(&V)),
        ]),
        plane_sizes(),
    )
}

fn plane_sizes() -> (NonZeroUsize, Option<NonZeroUsize>, Option<NonZeroUsize>) {
    let stride = NonZeroUsize::new(2).unwrap();
    (stride, Some(stride), Some(stride))
}

fn source_view_with_mixed_luma_pitch() -> FrameView<'static, u8> {
    static Y: [u8; 8] = [9, 11, 201, 202, 13, 15, 203, 204];
    static U: [u8; 4] = [0, 0, 0, 0];
    static V: [u8; 4] = [0, 0, 0, 0];

    FrameView::new(
        FramePlanes::new([
            Some(PlaneRef::new(&Y)),
            Some(PlaneRef::new(&U)),
            Some(PlaneRef::new(&V)),
        ]),
        (
            NonZeroUsize::new(4).unwrap(),
            Some(NonZeroUsize::new(2).unwrap()),
            Some(NonZeroUsize::new(2).unwrap()),
        ),
    )
}

#[test]
fn mask_new_rejects_non_positive_ml() {
    for ml in [0.0, -1.0] {
        let Err(err) = Mask::new(
            yuv444_info(),
            analysis_data(),
            MaskOptions {
                ml,
                ..MaskOptions::default()
            },
            None,
            None,
        ) else {
            panic!("Mask::new should reject non-positive ml");
        };

        assert_eq!(err.to_string(), "Mask: ml must be greater than 0.");
    }
}

#[test]
fn mask_render_frame_builds_motion_colormap_without_vapoursynth_types() {
    let mask = Mask::new(
        yuv444_info(),
        analysis_data(),
        MaskOptions {
            ml: 100.0,
            kind: MaskKind::MotionColormap,
            ..MaskOptions::default()
        },
        None,
        None,
    )
    .unwrap();
    let mut y = [0u8; 4];
    let mut u = [0u8; 4];
    let mut v = [0u8; 4];
    let mut output = FramePlanesMut::new([Some(&mut y), Some(&mut u), Some(&mut v)]);

    mask.render_frame(
        &source_view(),
        &mut output,
        plane_sizes(),
        &fake_vectors(true),
    )
    .unwrap();

    assert_eq!(y, [9, 11, 13, 15]);
    assert_eq!(u, [138, 118, 133, 128]);
    assert_eq!(v, [108, 148, 128, 123]);
}

#[test]
fn mask_render_frame_motion_colormap_copies_luma_rows_with_independent_pitches() {
    let mask = Mask::new(
        yuv444_info(),
        analysis_data(),
        MaskOptions {
            kind: MaskKind::MotionColormap,
            ..MaskOptions::default()
        },
        None,
        None,
    )
    .unwrap();
    let mut y = [0u8; 6];
    let mut u = [0u8; 4];
    let mut v = [0u8; 4];
    let mut output = FramePlanesMut::new([Some(&mut y), Some(&mut u), Some(&mut v)]);

    mask.render_frame(
        &source_view_with_mixed_luma_pitch(),
        &mut output,
        (
            NonZeroUsize::new(3).unwrap(),
            Some(NonZeroUsize::new(2).unwrap()),
            Some(NonZeroUsize::new(2).unwrap()),
        ),
        &fake_vectors(true),
    )
    .unwrap();

    assert_eq!(y, [9, 11, 0, 13, 15, 0]);
}

#[test]
fn mask_render_frame_fills_scene_change_value_when_vectors_are_unusable() {
    let mask = Mask::new(
        yuv444_info(),
        analysis_data(),
        MaskOptions {
            scene_change_value: 77,
            ..MaskOptions::default()
        },
        None,
        None,
    )
    .unwrap();
    let mut y = [0u8; 4];
    let mut u = [0u8; 4];
    let mut v = [0u8; 4];
    let mut output = FramePlanesMut::new([Some(&mut y), Some(&mut u), Some(&mut v)]);

    mask.render_frame(
        &source_view(),
        &mut output,
        plane_sizes(),
        &fake_vectors(false),
    )
    .unwrap();

    assert_eq!(y, [77, 77, 77, 77]);
    assert_eq!(u, [77, 77, 77, 77]);
    assert_eq!(v, [77, 77, 77, 77]);
}