zoomvtools 1.1.1

#[cfg(test)]
mod tests;

use std::num::NonZeroUsize;

use anyhow::{Result, anyhow, bail, ensure};
use semisafe::slice::get as semisafe_get;
use semisafe::slice::get_mut as semisafe_get_mut;

use crate::{
    analysis::MVAnalysisData,
    fake::group_of_planes::FakeGroupOfPlanes,
    frame::{FramePlanesMut, FrameView, PlaneSizeTuple},
    params::Subpel,
    resize::SimpleResize,
    util::Pixel,
    video::{ColorFamily, SampleType, VideoInfo},
};

/// Options used to construct [`FlowBlur`].
#[derive(Debug, Clone, Copy)]
pub struct FlowBlurOptions {
    /// Blur strength in the `0.0..=200.0` percent range.
    pub blur: f64,
    /// Number of temporal samples used during the blur.
    pub prec: i32,
}

/// Motion-compensated temporal blur filter.
pub struct FlowBlur {
    backward_data: MVAnalysisData,
    forward_data: MVAnalysisData,
    blur256: i32,
    prec: i32,
    width_uv: Option<NonZeroUsize>,
    height_uv: Option<NonZeroUsize>,
    h_padding_uv: usize,
    v_padding_uv: usize,
    vector_pitch_y: NonZeroUsize,
    vector_pitch_uv: Option<NonZeroUsize>,
    upsizer: SimpleResize,
    upsizer_uv: Option<SimpleResize>,
}

#[derive(Debug)]
struct ResizedVectorPair {
    vx_full: Vec<i16>,
    vy_full: Vec<i16>,
}

impl FlowBlur {
    /// Builds a flow-blur filter from validated clip and vector metadata.
    ///
    /// # Errors
    /// Returns an error if the clip format, vector metadata, or options are invalid.
    #[inline]
    pub fn new(
        info: VideoInfo,
        backward_data: MVAnalysisData,
        forward_data: MVAnalysisData,
        options: FlowBlurOptions,
    ) -> Result<Self> {
        ensure!(
            (0.0..=200.0).contains(&options.blur),
            "FlowBlur: blur must be between 0 and 200 % (inclusive)."
        );
        ensure!(options.prec >= 1, "FlowBlur: prec must be at least 1.");

        let format = info.format;
        if format.bits_per_sample.get() > 16 {
            bail!("FlowBlur: input clip must be 8-16 bits");
        }
        if format.sample_type != SampleType::Integer {
            bail!("FlowBlur: input clip must be integer super_format");
        }
        if ![ColorFamily::Yuv, ColorFamily::Gray].contains(&format.color_family)
            || format.sub_sampling_w > 1
            || format.sub_sampling_h > 1
        {
            bail!("FlowBlur: input clip must be GRAY, 420, 422, 440, or 444");
        }

        ensure!(
            backward_data.width == forward_data.width
                && backward_data.height == forward_data.height
                && backward_data.blk_x == forward_data.blk_x
                && backward_data.blk_y == forward_data.blk_y
                && backward_data.pel == forward_data.pel
                && backward_data.h_padding == forward_data.h_padding
                && backward_data.v_padding == forward_data.v_padding
                && backward_data.x_ratio_uv == forward_data.x_ratio_uv
                && backward_data.y_ratio_uv == forward_data.y_ratio_uv,
            "FlowBlur: mvbw and mvfw must use matching geometry."
        );
        ensure!(
            info.resolution.width == backward_data.width.get()
                && info.resolution.height == backward_data.height.get(),
            "FlowBlur: wrong source or super clip frame size."
        );
        let expected_x_ratio = 1usize << usize::from(format.sub_sampling_w);
        let expected_y_ratio = 1usize << usize::from(format.sub_sampling_h);
        ensure!(
            backward_data.x_ratio_uv.get() as usize == expected_x_ratio
                && backward_data.y_ratio_uv.get() as usize == expected_y_ratio,
            "FlowBlur: input clip subsampling does not match vector metadata."
        );

        let include_chroma = format.color_family == ColorFamily::Yuv;
        let x_ratio_uv = NonZeroUsize::from(backward_data.x_ratio_uv);
        let y_ratio_uv = NonZeroUsize::from(backward_data.y_ratio_uv);
        let (width_uv, height_uv, vector_pitch_uv, upsizer_uv) = if include_chroma {
            let width_uv = NonZeroUsize::new(backward_data.width.get() / x_ratio_uv.get())
                .ok_or_else(|| anyhow!("FlowBlur: wrong source or super clip frame size."))?;
            let height_uv = NonZeroUsize::new(backward_data.height.get() / y_ratio_uv.get())
                .ok_or_else(|| anyhow!("FlowBlur: wrong source or super clip frame size."))?;
            let upsizer_uv = SimpleResize::new(
                width_uv,
                height_uv,
                backward_data.blk_x,
                backward_data.blk_y,
                width_uv,
                height_uv,
                backward_data.pel,
            );
            (
                Some(width_uv),
                Some(height_uv),
                Some(width_uv),
                Some(upsizer_uv),
            )
        } else {
            (None, None, None, None)
        };

        let upsizer = SimpleResize::new(
            backward_data.width,
            backward_data.height,
            backward_data.blk_x,
            backward_data.blk_y,
            backward_data.width,
            backward_data.height,
            backward_data.pel,
        );
        Ok(Self {
            backward_data,
            forward_data,
            blur256: (options.blur * 256.0 / 200.0) as i32,
            prec: options.prec,
            width_uv,
            height_uv,
            h_padding_uv: backward_data.h_padding / x_ratio_uv.get(),
            v_padding_uv: backward_data.v_padding / y_ratio_uv.get(),
            vector_pitch_y: backward_data.width,
            vector_pitch_uv,
            upsizer,
            upsizer_uv,
        })
    }

    /// Returns whether both vector fields pass the supplied scene-change thresholds.
    #[must_use]
    #[inline]
    pub fn vectors_are_usable(
        &self,
        backward: &FakeGroupOfPlanes,
        forward: &FakeGroupOfPlanes,
        thscd1: u64,
        thscd2: u64,
    ) -> bool {
        backward.is_usable(thscd1, thscd2) && forward.is_usable(thscd1, thscd2)
    }

    /// Renders one blurred frame into `output`.
    ///
    /// # Errors
    /// Returns an error if the frame, vectors, or output buffers do not match the filter
    /// configuration.
    #[inline]
    pub fn render_frame<T: Pixel>(
        &self,
        reference: &FrameView<'_, T>,
        output: &mut FramePlanesMut<'_, T>,
        output_pitch: PlaneSizeTuple,
        backward: &FakeGroupOfPlanes,
        forward: &FakeGroupOfPlanes,
    ) -> Result<()> {
        let (vx_small_y_b, vy_small_y_b) = make_vector_small_masks(
            backward,
            self.backward_data.blk_x.get(),
            self.backward_data.blk_y.get(),
        );
        let (vx_small_y_f, vy_small_y_f) = make_vector_small_masks(
            forward,
            self.forward_data.blk_x.get(),
            self.forward_data.blk_y.get(),
        );
        let vectors_y_b = resize_vector_pair(
            &self.upsizer,
            self.backward_data.blk_x,
            self.vector_pitch_y,
            &vx_small_y_b,
            &vy_small_y_b,
        );
        let vectors_y_f = resize_vector_pair(
            &self.upsizer,
            self.forward_data.blk_x,
            self.vector_pitch_y,
            &vx_small_y_f,
            &vy_small_y_f,
        );

        {
            let ref_stride = reference.pitch_for_plane(0)?;
            let dst_stride = pitch_for_plane(output_pitch, 0)?;
            let ref_origin_offset = reference_origin_offset(
                ref_stride,
                self.backward_data.h_padding,
                self.backward_data.v_padding,
                self.backward_data.pel,
            );
            let ref_plane = reference.plane(0)?;
            let dst_plane = output.plane_mut(0)?;

            flow_blur_plane(
                dst_plane,
                dst_stride.get(),
                ref_plane,
                ref_stride.get(),
                ref_origin_offset,
                &vectors_y_b.vx_full,
                &vectors_y_f.vx_full,
                &vectors_y_b.vy_full,
                &vectors_y_f.vy_full,
                self.vector_pitch_y.get(),
                self.backward_data.width.get(),
                self.backward_data.height.get(),
                self.blur256,
                self.prec,
                self.backward_data.pel,
            );
        }

        if let Some(upsizer_uv) = &self.upsizer_uv {
            let vx_small_uv_b = vector_small_mask_y_to_uv(
                &vx_small_y_b,
                self.backward_data.x_ratio_uv.get() as usize,
            );
            let vy_small_uv_b = vector_small_mask_y_to_uv(
                &vy_small_y_b,
                self.backward_data.y_ratio_uv.get() as usize,
            );
            let vx_small_uv_f = vector_small_mask_y_to_uv(
                &vx_small_y_f,
                self.forward_data.x_ratio_uv.get() as usize,
            );
            let vy_small_uv_f = vector_small_mask_y_to_uv(
                &vy_small_y_f,
                self.forward_data.y_ratio_uv.get() as usize,
            );

            let vectors_uv_b = resize_vector_pair(
                upsizer_uv,
                self.backward_data.blk_x,
                self.vector_pitch_uv.expect("chroma pitch must exist"),
                &vx_small_uv_b,
                &vy_small_uv_b,
            );
            let vectors_uv_f = resize_vector_pair(
                upsizer_uv,
                self.forward_data.blk_x,
                self.vector_pitch_uv.expect("chroma pitch must exist"),
                &vx_small_uv_f,
                &vy_small_uv_f,
            );

            for plane in [1_usize, 2] {
                let ref_stride = reference.pitch_for_plane(plane)?;
                let dst_stride = pitch_for_plane(output_pitch, plane)?;
                let ref_origin_offset = reference_origin_offset(
                    ref_stride,
                    self.h_padding_uv,
                    self.v_padding_uv,
                    self.backward_data.pel,
                );
                let ref_plane = reference.plane(plane)?;
                let dst_plane = output.plane_mut(plane)?;

                flow_blur_plane(
                    dst_plane,
                    dst_stride.get(),
                    ref_plane,
                    ref_stride.get(),
                    ref_origin_offset,
                    &vectors_uv_b.vx_full,
                    &vectors_uv_f.vx_full,
                    &vectors_uv_b.vy_full,
                    &vectors_uv_f.vy_full,
                    self.vector_pitch_uv.expect("chroma pitch must exist").get(),
                    self.width_uv.expect("chroma width must exist").get(),
                    self.height_uv.expect("chroma height must exist").get(),
                    self.blur256,
                    self.prec,
                    self.backward_data.pel,
                );
            }
        }

        Ok(())
    }
}

fn flow_blur_plane<T: Pixel>(
    dst: &mut [T],
    dst_stride: usize,
    reference: &[T],
    ref_stride: usize,
    ref_origin_offset: isize,
    vx_full_b: &[i16],
    vx_full_f: &[i16],
    vy_full_b: &[i16],
    vy_full_f: &[i16],
    vector_pitch: usize,
    width: usize,
    height: usize,
    blur256: i32,
    prec: i32,
    pel: Subpel,
) {
    let pel_log = pel.log();

    for y in 0..height {
        for x in 0..width {
            let ref_index = ref_origin_offset
                + ((y * ref_stride) << pel_log) as isize
                + (x << pel_log) as isize;
            let center_index =
                usize::try_from(ref_index).expect("flow blur reference index must be in range");
            let mut blurred_sum: u64 =
                <T as num_traits::AsPrimitive<u64>>::as_(*semisafe_get(reference, center_index));

            let vx_f0 = *semisafe_get(vx_full_f, y * vector_pitch + x) as i32 * blur256;
            let vy_f0 = *semisafe_get(vy_full_f, y * vector_pitch + x) as i32 * blur256;
            let m_f = (vx_f0.abs().max(vy_f0.abs()) / prec) >> 8;
            if m_f > 0 {
                let step_x = vx_f0 / m_f;
                let step_y = vy_f0 / m_f;
                let mut vx = step_x;
                let mut vy = step_y;
                for _ in 0..m_f {
                    let sample_index =
                        ref_index + (vy >> 8) as isize * ref_stride as isize + (vx >> 8) as isize;
                    let sample_index = usize::try_from(sample_index)
                        .expect("flow blur forward sample index must be in range");
                    blurred_sum += <T as num_traits::AsPrimitive<u64>>::as_(*semisafe_get(
                        reference,
                        sample_index,
                    ));
                    vx += step_x;
                    vy += step_y;
                }
            }

            let vx_b0 = *semisafe_get(vx_full_b, y * vector_pitch + x) as i32 * blur256;
            let vy_b0 = *semisafe_get(vy_full_b, y * vector_pitch + x) as i32 * blur256;
            let m_b = (vx_b0.abs().max(vy_b0.abs()) / prec) >> 8;
            if m_b > 0 {
                let step_x = vx_b0 / m_b;
                let step_y = vy_b0 / m_b;
                let mut vx = step_x;
                let mut vy = step_y;
                for _ in 0..m_b {
                    let sample_index =
                        ref_index + (vy >> 8) as isize * ref_stride as isize + (vx >> 8) as isize;
                    let sample_index = usize::try_from(sample_index)
                        .expect("flow blur backward sample index must be in range");
                    blurred_sum += <T as num_traits::AsPrimitive<u64>>::as_(*semisafe_get(
                        reference,
                        sample_index,
                    ));
                    vx += step_x;
                    vy += step_y;
                }
            }

            let denominator = (m_f + m_b + 1) as u64;
            *semisafe_get_mut(dst, y * dst_stride + x) =
                T::from_u32_or_max_value((blurred_sum / denominator) as u32);
        }
    }
}

fn make_vector_small_masks(
    fake_gop: &FakeGroupOfPlanes,
    blk_x: usize,
    blk_y: usize,
) -> (Vec<i16>, Vec<i16>) {
    let mut vx = vec![0; blk_x * blk_y];
    let mut vy = vec![0; blk_x * blk_y];

    for by in 0..blk_y {
        for bx in 0..blk_x {
            let block = fake_gop.get_block(0, bx + by * blk_x);
            *semisafe_get_mut(&mut vx, bx + by * blk_x) = block.vector.x as i16;
            *semisafe_get_mut(&mut vy, bx + by * blk_x) = block.vector.y as i16;
        }
    }

    (vx, vy)
}

fn vector_small_mask_y_to_uv(v_small_y: &[i16], ratio_uv: usize) -> Vec<i16> {
    v_small_y
        .iter()
        .map(|value| if ratio_uv == 2 { *value >> 1 } else { *value })
        .collect()
}

fn resize_vector_pair(
    upsizer: &SimpleResize,
    blk_stride: NonZeroUsize,
    vec_pitch: NonZeroUsize,
    vx_small: &[i16],
    vy_small: &[i16],
) -> ResizedVectorPair {
    let vx_full = upsizer.resize_i16_to_vec(vec_pitch, vx_small, blk_stride, true);
    let vy_full = upsizer.resize_i16_to_vec(vec_pitch, vy_small, blk_stride, false);
    ResizedVectorPair { vx_full, vy_full }
}

fn pitch_for_plane(pitch: PlaneSizeTuple, plane: usize) -> Result<NonZeroUsize> {
    match plane {
        0 => Ok(pitch.0),
        1 => pitch
            .1
            .ok_or_else(|| anyhow!("requested plane 1 is not available")),
        2 => pitch
            .2
            .ok_or_else(|| anyhow!("requested plane 2 is not available")),
        _ => bail!("requested plane {plane} is not available"),
    }
}

#[inline]
const fn reference_origin_offset(
    stride: NonZeroUsize,
    h_padding: usize,
    v_padding: usize,
    pel: Subpel,
) -> isize {
    (stride.get() * v_padding * pel as usize + h_padding * pel as usize) as isize
}