zoomvtools 2.0.0

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

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

use smallvec::SmallVec;

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

use super::{Degrain, DegrainOptions};
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 super_info(height: usize, hpad: usize, vpad: usize) -> SuperClipInfo {
    SuperClipInfo {
        height: NonZeroUsize::new(height).unwrap(),
        hpad,
        vpad,
        pel: Subpel::Full,
        mode_yuv: MVPlaneSet::YPLANE,
        levels: 1,
    }
}

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

fn analysis_data(
    width: usize,
    height: usize,
    blk_size_x: usize,
    blk_size_y: usize,
    overlap_x: usize,
    overlap_y: usize,
    blk_x: usize,
    blk_y: usize,
    is_backward: bool,
) -> 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,
        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 degrain_options(limit: u16) -> DegrainOptions {
    DegrainOptions {
        thsad: 400,
        thsad_chroma: 400,
        plane: PlaneSelection::All,
        limit,
        limit_chroma: limit,
        thscd1: 400,
        thscd2: 130,
    }
}

fn frame_view<'a>(plane: &'a [u8], stride: usize) -> FrameView<'a, u8> {
    FrameView::new(
        FramePlanes::new([Some(PlaneRef::new(plane)), None, None]),
        (NonZeroUsize::new(stride).unwrap(), None, None),
    )
}

fn fake_vectors(data: &MVAnalysisData, sad: i64) -> FakeGroupOfPlanes {
    let mut fake_gop = FakeGroupOfPlanes::new(data);
    fake_gop.set_validity(true);
    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 degrain_render_frame_limits_changes_after_filtering_no_overlap() {
    let backward = analysis_data(2, 2, 2, 2, 0, 0, 1, 1, true);
    let forward = analysis_data(2, 2, 2, 2, 0, 0, 1, 1, false);
    let degrain = Degrain::<1>::new(
        gray8_info(2, 2),
        actual_super_width(2, 1),
        super_info(2, 1, 1),
        SmallVec::from_vec(vec![backward, forward]),
        degrain_options(5),
    )
    .unwrap();
    let source = [40u8; 4];
    let backward_super = [100u8; 16];
    let forward_super = [100u8; 16];
    let source_view = frame_view(&source, 2);
    let backward_view = frame_view(&backward_super, 4);
    let forward_view = frame_view(&forward_super, 4);
    let vectors = vec![fake_vectors(&backward, 0), fake_vectors(&forward, 0)];
    let references = [Some(&backward_view), Some(&forward_view)];
    let mut output = [0u8; 4];
    let mut output_planes = FramePlanesMut::new([Some(&mut output), None, None]);

    degrain
        .render_frame(
            &source_view,
            &mut output_planes,
            (NonZeroUsize::new(2).unwrap(), None, None),
            &vectors,
            &references,
        )
        .unwrap();

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

#[test]
fn degrain_render_frame_limits_changes_after_filtering_overlap() {
    let backward = analysis_data(8, 8, 4, 4, 2, 2, 3, 3, true);
    let forward = analysis_data(8, 8, 4, 4, 2, 2, 3, 3, false);
    let degrain = Degrain::<1>::new(
        gray8_info(8, 8),
        actual_super_width(8, 1),
        super_info(8, 1, 1),
        SmallVec::from_vec(vec![backward, forward]),
        degrain_options(5),
    )
    .unwrap();
    let source = [40u8; 64];
    let backward_super = [100u8; 100];
    let forward_super = [100u8; 100];
    let source_view = frame_view(&source, 8);
    let backward_view = frame_view(&backward_super, 10);
    let forward_view = frame_view(&forward_super, 10);
    let vectors = vec![fake_vectors(&backward, 0), fake_vectors(&forward, 0)];
    let references = [Some(&backward_view), Some(&forward_view)];
    let mut output = [0u8; 64];
    let mut output_planes = FramePlanesMut::new([Some(&mut output), None, None]);

    degrain
        .render_frame(
            &source_view,
            &mut output_planes,
            (NonZeroUsize::new(8).unwrap(), None, None),
            &vectors,
            &references,
        )
        .unwrap();

    assert_eq!(output, [45u8; 64]);
}

#[test]
fn degrain_new_rejects_wrong_actual_super_width() {
    let backward = analysis_data(2, 2, 2, 2, 0, 0, 1, 1, true);
    let forward = analysis_data(2, 2, 2, 2, 0, 0, 1, 1, false);

    let Err(err) = Degrain::<1>::new(
        gray8_info(2, 2),
        3,
        super_info(2, 1, 1),
        SmallVec::from_vec(vec![backward, forward]),
        degrain_options(5),
    ) else {
        panic!("expected malformed super width to be rejected");
    };

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