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::{MVPlaneSet, MotionFlags, SceneChangeBehavior, Subpel},
    video::{ColorFamily, Resolution, SampleType, VideoFormat, VideoInfo},
};

use super::{Compensate, CompensateOptions};
use crate::filters::analyse::SuperClipInfo;

fn gray8_info(width: usize, height: usize) -> VideoInfo {
    VideoInfo {
        format: VideoFormat {
            color_family: ColorFamily::Gray,
            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, height },
        num_frames: 10,
    }
}

fn yuv444_8_info(width: usize, height: usize) -> 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, height },
        num_frames: 10,
    }
}

fn super_info() -> SuperClipInfo {
    SuperClipInfo {
        height: NonZeroUsize::new(2).unwrap(),
        hpad: 1,
        vpad: 1,
        pel: Subpel::Full,
        mode_yuv: MVPlaneSet::YPLANE,
        levels: 1,
    }
}

fn super_info_with_geometry(width: usize, height: usize, mode_yuv: MVPlaneSet) -> SuperClipInfo {
    let _ = width;
    SuperClipInfo {
        height: NonZeroUsize::new(height).unwrap(),
        hpad: 1,
        vpad: 1,
        pel: Subpel::Full,
        mode_yuv,
        levels: 1,
    }
}

fn actual_super_width(width: usize, hpad: usize) -> usize {
    width + hpad * 2
}

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

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

fn compensate_options(scene_change_behavior: SceneChangeBehavior) -> CompensateOptions {
    CompensateOptions {
        scene_change_behavior,
        thsad: 400,
        fields: false,
        time_percent: 100.0,
        thscd1: 400,
        thscd2: 130,
    }
}

fn super_view<'a>(plane: &'a [u8; 16]) -> FrameView<'a, u8> {
    FrameView::new(
        FramePlanes::new([Some(PlaneRef::new(plane)), None, None]),
        (NonZeroUsize::new(4).unwrap(), None, None),
    )
}

fn yuv_super_view<'a>(y: &'a [u8], u: &'a [u8], v: &'a [u8], stride: usize) -> FrameView<'a, u8> {
    let stride = NonZeroUsize::new(stride).unwrap();
    FrameView::new(
        FramePlanes::new([
            Some(PlaneRef::new(y)),
            Some(PlaneRef::new(u)),
            Some(PlaneRef::new(v)),
        ]),
        (stride, Some(stride), Some(stride)),
    )
}

fn fake_vectors(validity: bool, sad: i64) -> FakeGroupOfPlanes {
    let mut fake_gop = FakeGroupOfPlanes::new(&analysis_data());
    fake_gop.set_validity(validity);
    let plane = fake_gop.plane_mut(0);
    let block = plane.blocks.first_mut().expect("expected first block");
    block.vector = MotionVector { x: 0, y: 0, sad };
    fake_gop
}

fn fake_vectors_for(data: &MVAnalysisData, validity: bool, sad: i64) -> FakeGroupOfPlanes {
    let mut fake_gop = FakeGroupOfPlanes::new(data);
    fake_gop.set_validity(validity);
    let plane = fake_gop.plane_mut(0);
    for block in &mut plane.blocks {
        block.vector = MotionVector { x: 0, y: 0, sad };
    }
    fake_gop
}

#[test]
fn compensate_render_frame_copies_reference_when_vectors_are_usable_no_overlap() {
    let compensate = Compensate::new(
        gray8_info(2, 2),
        actual_super_width(2, 1),
        super_info(),
        analysis_data(),
        compensate_options(SceneChangeBehavior::ReferenceFrame),
    )
    .unwrap();
    let current_super = [3u8; 16];
    let reference_super = [7u8; 16];
    let current_view = super_view(&current_super);
    let reference_view = super_view(&reference_super);
    let mut output = [0u8; 4];
    let mut output_planes = FramePlanesMut::new([Some(&mut output), None, None]);

    compensate
        .render_frame(
            &current_view,
            Some(&reference_view),
            &mut output_planes,
            (NonZeroUsize::new(2).unwrap(), None, None),
            &fake_vectors(true, 0),
            0,
        )
        .unwrap();

    assert_eq!(output, [7, 7, 7, 7]);
}

#[test]
fn compensate_render_frame_uses_reference_frame_when_scene_change_requests_reference_fallback() {
    let compensate = Compensate::new(
        gray8_info(2, 2),
        actual_super_width(2, 1),
        super_info(),
        analysis_data(),
        compensate_options(SceneChangeBehavior::ReferenceFrame),
    )
    .unwrap();
    let current_super = [3u8; 16];
    let reference_super = [7u8; 16];
    let current_view = super_view(&current_super);
    let reference_view = super_view(&reference_super);
    let mut output = [0u8; 4];
    let mut output_planes = FramePlanesMut::new([Some(&mut output), None, None]);

    compensate
        .render_frame(
            &current_view,
            Some(&reference_view),
            &mut output_planes,
            (NonZeroUsize::new(2).unwrap(), None, None),
            &fake_vectors(false, 0),
            0,
        )
        .unwrap();

    assert_eq!(output, [7, 7, 7, 7]);
}

#[test]
fn compensate_render_frame_preserves_current_chroma_when_super_clip_is_y_only() {
    let info = yuv444_8_info(2, 2);
    let vectors_data = analysis_data();
    let compensate = Compensate::new(
        info,
        actual_super_width(2, 1),
        super_info_with_geometry(2, 2, MVPlaneSet::YPLANE),
        vectors_data,
        compensate_options(SceneChangeBehavior::ReferenceFrame),
    )
    .unwrap();
    let current_y = [3u8; 16];
    let reference_y = [7u8; 16];
    let current_u = [11u8; 16];
    let reference_u = [13u8; 16];
    let current_v = [21u8; 16];
    let reference_v = [23u8; 16];
    let current_view = yuv_super_view(&current_y, &current_u, &current_v, 4);
    let reference_view = yuv_super_view(&reference_y, &reference_u, &reference_v, 4);
    let mut y = [0u8; 4];
    let mut u = [0u8; 4];
    let mut v = [0u8; 4];
    let mut output_planes = FramePlanesMut::new([Some(&mut y), Some(&mut u), Some(&mut v)]);

    compensate
        .render_frame(
            &current_view,
            Some(&reference_view),
            &mut output_planes,
            (
                NonZeroUsize::new(2).unwrap(),
                Some(NonZeroUsize::new(2).unwrap()),
                Some(NonZeroUsize::new(2).unwrap()),
            ),
            &fake_vectors_for(&vectors_data, true, 0),
            0,
        )
        .unwrap();

    assert_eq!(y, [7, 7, 7, 7]);
    assert_eq!(u, [11, 11, 11, 11]);
    assert_eq!(v, [21, 21, 21, 21]);
}

#[test]
fn compensate_render_frame_overlap_preserves_current_chroma_when_super_clip_is_y_only() {
    let info = yuv444_8_info(8, 8);
    let vectors_data = analysis_data_with_geometry(8, 8, 4, 4, 2, 2, 3, 3);
    let compensate = Compensate::new(
        info,
        actual_super_width(8, 1),
        super_info_with_geometry(8, 8, MVPlaneSet::YPLANE),
        vectors_data,
        compensate_options(SceneChangeBehavior::ReferenceFrame),
    )
    .unwrap();
    let current_y = [3u8; 100];
    let reference_y = [7u8; 100];
    let current_u = [11u8; 100];
    let reference_u = [13u8; 100];
    let current_v = [21u8; 100];
    let reference_v = [23u8; 100];
    let current_view = yuv_super_view(&current_y, &current_u, &current_v, 10);
    let reference_view = yuv_super_view(&reference_y, &reference_u, &reference_v, 10);
    let mut y = [0u8; 64];
    let mut u = [0u8; 64];
    let mut v = [0u8; 64];
    let mut output_planes = FramePlanesMut::new([Some(&mut y), Some(&mut u), Some(&mut v)]);

    compensate
        .render_frame(
            &current_view,
            Some(&reference_view),
            &mut output_planes,
            (
                NonZeroUsize::new(8).unwrap(),
                Some(NonZeroUsize::new(8).unwrap()),
                Some(NonZeroUsize::new(8).unwrap()),
            ),
            &fake_vectors_for(&vectors_data, true, 0),
            0,
        )
        .unwrap();

    assert_eq!(y, [7u8; 64]);
    assert_eq!(u, [11u8; 64]);
    assert_eq!(v, [21u8; 64]);
}

#[test]
fn compensate_new_rejects_wrong_actual_super_width() {
    let Err(err) = Compensate::new(
        gray8_info(2, 2),
        3,
        super_info(),
        analysis_data(),
        compensate_options(SceneChangeBehavior::ReferenceFrame),
    ) else {
        panic!("expected malformed super width to be rejected");
    };

    assert_eq!(
        err.to_string(),
        "Compensate: wrong source or super clip frame size"
    );
}