zoomvtools 1.1.1

Video motion vector analysis utilities in pure Rust
Documentation
use std::{
    cmp::{max, min},
    num::{NonZeroU8, NonZeroUsize},
};

use crate::util::Pixel;
use semisafe::slice::get as semisafe_get;
use semisafe::slice::get_mut as semisafe_get_mut;

pub(super) fn refine_horizontal_wiener<T: Pixel>(
    dest: &mut [T],
    src: &[T],
    pitch: NonZeroUsize,
    width: NonZeroUsize,
    height: NonZeroUsize,
    bits_per_sample: NonZeroU8,
) {
    let pixel_max = (1i32 << bits_per_sample.get()) - 1;
    let mut offset = 0;

    for _j in 0..height.get() {
        let src_row = semisafe_get(semisafe_get(src, offset..), ..width.get());
        let dest_row = semisafe_get_mut(semisafe_get_mut(dest, offset..), ..width.get());

        // Handle first two pixels with bilinear interpolation (if width >= 2)
        if width.get() >= 2 {
            let a: u32 = semisafe_get(src_row, 0).as_();
            let b: u32 = semisafe_get(src_row, 1).as_();
            *semisafe_get_mut(dest_row, 0) = T::from_u32_or_max_value((a + b + 1) / 2);

            if width.get() >= 3 {
                let c: u32 = semisafe_get(src_row, 2).as_();
                *semisafe_get_mut(dest_row, 1) = T::from_u32_or_max_value((b + c + 1) / 2);
            }
        }

        // Process middle pixels with Wiener filter
        let wiener_start = 2;
        let wiener_end = if width.get() >= 4 {
            width.get() - 4
        } else {
            wiener_start
        };

        for i in wiener_start..wiener_end {
            let mut m0: i32 = semisafe_get(src_row, i - 2).as_();
            let m1: i32 = semisafe_get(src_row, i - 1).as_();
            let mut m2: i32 = semisafe_get(src_row, i).as_();
            let m3: i32 = semisafe_get(src_row, i + 1).as_();
            let m4: i32 = semisafe_get(src_row, i + 2).as_();
            let m5: i32 = semisafe_get(src_row, i + 3).as_();

            m2 = (m2 + m3) * 4;

            m2 -= m1 + m4;
            m2 *= 5;

            m0 += m5 + m2 + 16;
            m0 >>= 5;

            *semisafe_get_mut(dest_row, i) =
                T::from_u32_or_max_value(max(0, min(m0, pixel_max)) as u32);
        }

        // Handle last few pixels with bilinear interpolation
        for i in wiener_end..(width.get() - 1).min(width.get()) {
            let a: u32 = semisafe_get(src_row, i).as_();
            let b: u32 = semisafe_get(src_row, i + 1).as_();
            *semisafe_get_mut(dest_row, i) = T::from_u32_or_max_value((a + b + 1) / 2);
        }

        // Copy last pixel
        if width.get() > 0 {
            *semisafe_get_mut(dest_row, width.get() - 1) = *semisafe_get(src_row, width.get() - 1);
        }
        offset += pitch.get();
    }
}

pub(super) fn refine_vertical_wiener<T: Pixel>(
    dest: &mut [T],
    src: &[T],
    pitch: NonZeroUsize,
    width: NonZeroUsize,
    height: NonZeroUsize,
    bits_per_sample: NonZeroU8,
) {
    let pixel_max = (1i32 << bits_per_sample.get()) - 1;
    let mut offset = 0;

    for _j in 0..2 {
        for i in 0..width.get() {
            let a: u32 = semisafe_get(src, offset + i).as_();
            let b: u32 = semisafe_get(src, offset + i + pitch.get()).as_();
            *semisafe_get_mut(dest, offset + i) = T::from_u32_or_max_value((a + b + 1) / 2);
        }
        offset += pitch.get();
    }

    for _j in 2..(height.get() - 4) {
        for i in 0..width.get() {
            let mut m0: i32 = semisafe_get(src, offset + i - pitch.get() * 2).as_();
            let m1: i32 = semisafe_get(src, offset + i - pitch.get()).as_();
            let mut m2: i32 = semisafe_get(src, offset + i).as_();
            let m3: i32 = semisafe_get(src, offset + i + pitch.get()).as_();
            let m4: i32 = semisafe_get(src, offset + i + pitch.get() * 2).as_();
            let m5: i32 = semisafe_get(src, offset + i + pitch.get() * 3).as_();

            m2 = (m2 + m3) * 4;

            m2 -= m1 + m4;
            m2 *= 5;

            m0 += m5 + m2 + 16;
            m0 >>= 5;

            *semisafe_get_mut(dest, offset + i) =
                T::from_u32_or_max_value(max(0, min(m0, pixel_max)) as u32);
        }
        offset += pitch.get();
    }

    for _j in (height.get() - 4)..(height.get() - 1) {
        for i in 0..width.get() {
            let a: u32 = semisafe_get(src, offset + i).as_();
            let b: u32 = semisafe_get(src, offset + i + pitch.get()).as_();
            *semisafe_get_mut(dest, offset + i) = T::from_u32_or_max_value((a + b + 1) / 2);
        }

        offset += pitch.get();
    }

    // last row
    let dest_row = semisafe_get_mut(semisafe_get_mut(dest, offset..), ..width.get());
    let src_row = semisafe_get(semisafe_get(src, offset..), ..width.get());
    dest_row.copy_from_slice(src_row);
}