zoomvtools 1.1.1

Video motion vector analysis utilities in pure Rust
Documentation
#[cfg(test)]
mod tests;

use std::cmp::{max, min};

use crate::{fake::group_of_planes::FakeGroupOfPlanes, params::Subpel};
use semisafe::slice::get as semisafe_get;
use semisafe::slice::get_mut as semisafe_get_mut;

#[inline]
/// Extends a smaller block mask to the padded block grid by repeating the last
/// valid row and column.
pub fn check_and_pad_mask_small(
    mask_small: &mut [u8],
    blk_x_padded: usize,
    blk_y_padded: usize,
    blk_x: usize,
    blk_y: usize,
) {
    if blk_x_padded > blk_x {
        for row in 0..blk_y {
            let right = *semisafe_get(mask_small, row * blk_x_padded + blk_x - 1);
            for col in blk_x..blk_x_padded {
                *semisafe_get_mut(mask_small, row * blk_x_padded + col) = right;
            }
        }
    }

    if blk_y_padded > blk_y {
        for col in 0..blk_x_padded {
            let bottom = *semisafe_get(mask_small, blk_x_padded * (blk_y - 1) + col);
            for row in blk_y..blk_y_padded {
                *semisafe_get_mut(mask_small, blk_x_padded * row + col) = bottom;
            }
        }
    }
}

fn byte_occ_mask(mask_ptr: &mut u8, occlusion: i32, occ_norm: f64, gamma: f64) {
    *mask_ptr = max(
        *mask_ptr,
        if gamma == 1.0 {
            ((255.0 * occlusion as f64 * occ_norm) as i32).clamp(0, 255) as u8
        } else {
            ((255.0 * (occlusion as f64 * occ_norm).powf(gamma)) as i32).clamp(0, 255) as u8
        },
    );
}

fn byte_norm(sad: i64, sad_norm_factor: f64, gamma: f64) -> u8 {
    let l = 255.0 * (sad as f64 * sad_norm_factor).powf(gamma);
    if l > 255.0 { 255 } else { l as u8 }
}

#[expect(
    clippy::too_many_arguments,
    reason = "mirror C MakeVectorOcclusionMaskTime inputs"
)]
#[must_use]
#[inline]
/// Builds the per-block occlusion mask used by flow-based interpolation.
pub fn make_vector_occlusion_mask_time(
    fake_gop: &FakeGroupOfPlanes,
    is_backward: bool,
    blk_x: usize,
    blk_y: usize,
    mask_norm_divider: f64,
    gamma: f64,
    pel: Subpel,
    mask_pitch: usize,
    mask_height: usize,
    time256: i32,
    blk_step_x: usize,
    blk_step_y: usize,
) -> Vec<u8> {
    let mut mask = vec![0; mask_height * mask_pitch];

    let pel = pel as i32;
    let time4096_x = time256 * 16 / (blk_step_x as i32 * pel);
    let time4096_y = time256 * 16 / (blk_step_y as i32 * pel);
    let occ_norm_x = 80.0 / (mask_norm_divider * blk_step_x as f64 * pel as f64);
    let occ_norm_y = 80.0 / (mask_norm_divider * blk_step_y as f64 * pel as f64);

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

            if bx + 1 < blk_x {
                let right = fake_gop.get_block(0, bx + 1 + by * blk_x);
                if right.vector.x < vx {
                    let occlusion = vx - right.vector.x;
                    let min_b = if is_backward {
                        max(0, bx as i32 + 1 - occlusion * time4096_x / 4096) as usize
                    } else {
                        bx
                    };
                    let max_b = if is_backward {
                        bx + 1
                    } else {
                        min(
                            bx as i32 + 1 - occlusion * time4096_x / 4096,
                            blk_x as i32 - 1,
                        ) as usize
                    };
                    for bxi in min_b..=max_b {
                        byte_occ_mask(
                            semisafe_get_mut(&mut mask, bxi + by * mask_pitch),
                            occlusion,
                            occ_norm_x,
                            gamma,
                        );
                    }
                }
            }

            if by + 1 < blk_y {
                let bottom = fake_gop.get_block(0, bx + (by + 1) * blk_x);
                if bottom.vector.y < vy {
                    let occlusion = vy - bottom.vector.y;
                    let min_b = if is_backward {
                        max(0, by as i32 + 1 - occlusion * time4096_y / 4096) as usize
                    } else {
                        by
                    };
                    let max_b = if is_backward {
                        by + 1
                    } else {
                        min(
                            by as i32 + 1 - occlusion * time4096_y / 4096,
                            blk_y as i32 - 1,
                        ) as usize
                    };
                    for byi in min_b..=max_b {
                        byte_occ_mask(
                            semisafe_get_mut(&mut mask, bx + byi * mask_pitch),
                            occlusion,
                            occ_norm_y,
                            gamma,
                        );
                    }
                }
            }
        }
    }

    mask
}

#[expect(clippy::too_many_arguments, reason = "mirror C MakeSADMaskTime inputs")]
#[must_use]
#[inline]
/// Builds the per-block SAD mask used by flow-based interpolation.
pub fn make_sad_mask_time(
    fake_gop: &FakeGroupOfPlanes,
    blk_x: usize,
    blk_y: usize,
    sad_norm_factor: f64,
    gamma: f64,
    pel: Subpel,
    mask_pitch: usize,
    mask_height: usize,
    time256: i32,
    blk_step_x: usize,
    blk_step_y: usize,
    bits_per_sample: u8,
) -> Vec<u8> {
    let mut mask = vec![0; mask_height * mask_pitch];

    let pel = pel as i32;
    let time4096_x = (256 - time256) * 16 / (blk_step_x as i32 * pel);
    let time4096_y = (256 - time256) * 16 / (blk_step_y as i32 * pel);

    for by in 0..blk_y {
        for bx in 0..blk_x {
            let i = bx + by * blk_x;
            let block = fake_gop.get_block(0, i);
            let vx = block.vector.x;
            let vy = block.vector.y;
            let mut bxi = bx as i32 - vx * time4096_x / 4096;
            let mut byi = by as i32 - vy * time4096_y / 4096;
            if bxi < 0 || bxi >= blk_x as i32 || byi < 0 || byi >= blk_y as i32 {
                bxi = bx as i32;
                byi = by as i32;
            }

            let i1 = bxi as usize + byi as usize * blk_x;
            let sad = fake_gop.get_block(0, i1).vector.sad >> (bits_per_sample - 8);
            *semisafe_get_mut(&mut mask, bx + by * mask_pitch) =
                byte_norm(sad, sad_norm_factor, gamma);
        }
    }

    mask
}