box_image_pyramid/

pyramid.rs

//! Image pyramid construction using fixed 2x box-filter downsampling.
//!
//! The API is allocation-friendly: construct a [`PyramidBuffers`] once, then
//! reuse it to build pyramids for successive frames without re-allocating
//! intermediate levels. When both the `par_pyramid` and `simd` features are
//! enabled, the 2x box downsample uses portable SIMD for higher throughput.

use crate::imageview::{ImageBuffer, ImageView};
#[cfg(feature = "tracing")]
use tracing::instrument;

/// Reusable backing storage for pyramid construction.
///
/// Typically you construct a [`PyramidBuffers`] once (for example with
/// [`PyramidBuffers::with_capacity`]) and reuse it across frames by
/// passing a mutable reference into [`build_pyramid`]. The internal level
/// buffers are resized on demand to match the requested pyramid shape.
pub struct PyramidBuffers {
    levels: Vec<ImageBuffer>,
}

impl Default for PyramidBuffers {
    fn default() -> Self {
        Self::new()
    }
}

impl PyramidBuffers {
    /// Create an empty buffer set.
    pub fn new() -> Self {
        Self { levels: Vec::new() }
    }

    /// Create a buffer set with capacity reserved for `num_levels`.
    pub fn with_capacity(num_levels: u8) -> Self {
        Self {
            levels: Vec::with_capacity(num_levels.saturating_sub(1) as usize),
        }
    }

    fn ensure_level_shape(&mut self, idx: usize, w: usize, h: usize) {
        if idx >= self.levels.len() {
            self.levels.resize_with(idx + 1, || ImageBuffer::new(w, h));
        }

        let level = &mut self.levels[idx];
        if level.width != w || level.height != h {
            *level = ImageBuffer::new(w, h);
        }
    }
}

/// A single level of a [`Pyramid`].
#[non_exhaustive]
pub struct PyramidLevel<'a> {
    /// Image data for this level. Level 0 borrows the caller-supplied
    /// base; subsequent levels are views into [`PyramidBuffers`].
    pub img: ImageView<'a>,
    /// Scale of this level relative to the base image. Level 0 is
    /// always `1.0`; each subsequent level is half the previous
    /// (i.e. `0.5`, `0.25`, …).
    pub scale: f32,
}

/// A top-down image pyramid produced by [`build_pyramid`].
///
/// `levels[0]` is always the base (full resolution, `scale = 1.0`).
/// Each subsequent level is a 2× box-filter downsample of the previous.
/// The number of levels is determined by [`PyramidParams::num_levels`]
/// and [`PyramidParams::min_size`].
#[non_exhaustive]
pub struct Pyramid<'a> {
    /// Ordered pyramid levels from base (index 0) to coarsest.
    pub levels: Vec<PyramidLevel<'a>>,
}

/// Parameters controlling pyramid generation.
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct PyramidParams {
    /// Maximum number of levels (including the base).
    pub num_levels: u8,
    /// Stop building when either dimension falls below this value.
    pub min_size: usize,
}

impl Default for PyramidParams {
    fn default() -> Self {
        Self {
            num_levels: 1,
            min_size: 128,
        }
    }
}

/// Build a top-down image pyramid using fixed 2x downsampling.
///
/// The base image is always included as level 0. Each subsequent level is a
/// 2x downsampled copy (box filter) written into `buffers`. Construction stops
/// when:
/// - either dimension would fall below `min_size`, or
/// - `num_levels` is reached.
#[cfg_attr(
    feature = "tracing",
    instrument(
        level = "info",
        skip(base, params, buffers),
        fields(levels = params.num_levels, min_size = params.min_size)
    )
)]
pub fn build_pyramid<'a>(
    base: ImageView<'a>,
    params: &PyramidParams,
    buffers: &'a mut PyramidBuffers,
) -> Pyramid<'a> {
    if params.num_levels == 0 || base.width < params.min_size || base.height < params.min_size {
        return Pyramid { levels: Vec::new() };
    }

    #[derive(Clone, Copy)]
    enum LevelSource {
        Base,
        Buffer(usize),
    }

    let mut sources: Vec<(LevelSource, f32)> = Vec::with_capacity(params.num_levels as usize);
    sources.push((LevelSource::Base, 1.0));

    let mut current_src = LevelSource::Base;
    let mut current_w = base.width;
    let mut current_h = base.height;
    let mut scale = 1.0f32;

    for level_idx in 1..params.num_levels {
        let w2 = current_w / 2;
        let h2 = current_h / 2;

        if w2 == 0 || h2 == 0 || w2 < params.min_size || h2 < params.min_size {
            break;
        }

        let buf_idx = (level_idx - 1) as usize;
        buffers.ensure_level_shape(buf_idx, w2, h2);

        let (src_img, dst): (ImageView<'_>, &mut ImageBuffer) = match current_src {
            LevelSource::Base => (base, &mut buffers.levels[buf_idx]),
            LevelSource::Buffer(src_idx) => {
                debug_assert!(src_idx < buf_idx);
                let (head, tail) = buffers.levels.split_at_mut(buf_idx);
                (head[src_idx].as_view(), &mut tail[0])
            }
        };

        downsample_2x_box(src_img, dst);

        scale *= 0.5;
        current_src = LevelSource::Buffer(buf_idx);
        current_w = w2;
        current_h = h2;
        sources.push((current_src, scale));
    }

    let mut levels = Vec::with_capacity(sources.len());
    for (source, lvl_scale) in sources {
        let img = match source {
            LevelSource::Base => base,
            LevelSource::Buffer(idx) => buffers.levels[idx].as_view(),
        };
        levels.push(PyramidLevel {
            img,
            scale: lvl_scale,
        });
    }

    Pyramid { levels }
}

/// Fast 2x downsample with a 2x2 box filter into a pre-allocated destination.
///
/// Uses SIMD and/or `rayon` specializations when the `par_pyramid`
/// feature is enabled alongside the relevant flags.
#[inline]
fn downsample_2x_box(src: ImageView<'_>, dst: &mut ImageBuffer) {
    #[cfg(all(feature = "par_pyramid", feature = "rayon", feature = "simd"))]
    return downsample_2x_box_parallel_simd(src, dst);

    #[cfg(all(feature = "par_pyramid", feature = "rayon", not(feature = "simd")))]
    return downsample_2x_box_parallel_scalar(src, dst);

    #[cfg(all(feature = "par_pyramid", not(feature = "rayon"), feature = "simd"))]
    return downsample_2x_box_simd(src, dst);

    #[cfg(all(feature = "par_pyramid", not(feature = "rayon"), not(feature = "simd")))]
    return downsample_2x_box_scalar(src, dst);

    #[cfg(not(feature = "par_pyramid"))]
    return downsample_2x_box_scalar(src, dst);
}

#[inline]
#[cfg_attr(
    all(feature = "par_pyramid", any(feature = "rayon", feature = "simd")),
    allow(dead_code)
)]
fn downsample_2x_box_scalar(src: ImageView<'_>, dst: &mut ImageBuffer) {
    debug_assert_eq!(src.width / 2, dst.width);
    debug_assert_eq!(src.height / 2, dst.height);

    let src_w = src.width;
    let dst_w = dst.width;
    let dst_h = dst.height;

    for y in 0..dst_h {
        let row0 = (y * 2) * src_w;
        let row1 = row0 + src_w;

        downsample_row_scalar(
            &src.data[row0..row0 + src_w],
            &src.data[row1..row1 + src_w],
            &mut dst.data[y * dst_w..(y + 1) * dst_w],
        );
    }
}

#[cfg(all(feature = "par_pyramid", not(feature = "rayon"), feature = "simd"))]
fn downsample_2x_box_simd(src: ImageView<'_>, dst: &mut ImageBuffer) {
    debug_assert_eq!(src.width / 2, dst.width);
    debug_assert_eq!(src.height / 2, dst.height);

    let src_w = src.width;
    let dst_w = dst.width;
    let dst_h = dst.height;

    for y_out in 0..dst_h {
        let y0 = 2 * y_out;
        let y1 = y0 + 1;

        let row0 = &src.data[y0 * src_w..(y0 + 1) * src_w];
        let row1 = &src.data[y1 * src_w..(y1 + 1) * src_w];

        let dst_row = &mut dst.data[y_out * dst_w..(y_out + 1) * dst_w];

        downsample_row_simd(row0, row1, dst_row);
    }
}

#[cfg(all(feature = "par_pyramid", feature = "rayon", not(feature = "simd")))]
fn downsample_2x_box_parallel_scalar(src: ImageView<'_>, dst: &mut ImageBuffer) {
    use rayon::prelude::*;

    debug_assert_eq!(src.width / 2, dst.width);
    debug_assert_eq!(src.height / 2, dst.height);

    let src_w = src.width;
    let dst_w = dst.width;

    dst.data
        .par_chunks_mut(dst_w)
        .enumerate()
        .for_each(|(y_out, dst_row)| {
            let y0 = 2 * y_out;
            let y1 = y0 + 1;

            let row0 = &src.data[y0 * src_w..(y0 + 1) * src_w];
            let row1 = &src.data[y1 * src_w..(y1 + 1) * src_w];

            downsample_row_scalar(row0, row1, dst_row);
        });
}

#[cfg(all(feature = "par_pyramid", feature = "rayon", feature = "simd"))]
fn downsample_2x_box_parallel_simd(src: ImageView<'_>, dst: &mut ImageBuffer) {
    use rayon::prelude::*;

    debug_assert_eq!(src.width / 2, dst.width);
    debug_assert_eq!(src.height / 2, dst.height);

    let src_w = src.width;
    let dst_w = dst.width;

    dst.data
        .par_chunks_mut(dst_w)
        .enumerate()
        .for_each(|(y_out, dst_row)| {
            let y0 = 2 * y_out;
            let y1 = y0 + 1;

            let row0 = &src.data[y0 * src_w..(y0 + 1) * src_w];
            let row1 = &src.data[y1 * src_w..(y1 + 1) * src_w];

            downsample_row_simd(row0, row1, dst_row);
        });
}

#[inline]
fn downsample_row_scalar(row0: &[u8], row1: &[u8], dst_row: &mut [u8]) {
    let dst_w = dst_row.len();

    for (x, item) in dst_row.iter_mut().enumerate().take(dst_w) {
        let sx = x * 2;
        let p00 = row0[sx] as u16;
        let p01 = row0[sx + 1] as u16;
        let p10 = row1[sx] as u16;
        let p11 = row1[sx + 1] as u16;
        let sum = p00 + p01 + p10 + p11;
        *item = ((sum + 2) >> 2) as u8;
    }
}

#[cfg(all(feature = "par_pyramid", feature = "simd"))]
fn downsample_row_simd(row0: &[u8], row1: &[u8], dst_row: &mut [u8]) {
    use std::ops::Shr;
    use std::simd::num::SimdUint;
    use std::simd::{u16x16, u8x16};

    const LANES: usize = 16;
    let mut x_out = 0usize;

    while x_out + LANES <= dst_row.len() {
        let mut p00 = [0u8; LANES];
        let mut p01 = [0u8; LANES];
        let mut p10 = [0u8; LANES];
        let mut p11 = [0u8; LANES];

        for lane in 0..LANES {
            let x = x_out + lane;
            let sx = 2 * x;
            p00[lane] = row0[sx];
            p01[lane] = row0[sx + 1];
            p10[lane] = row1[sx];
            p11[lane] = row1[sx + 1];
        }

        let p00v = u8x16::from_array(p00).cast::<u16>();
        let p01v = u8x16::from_array(p01).cast::<u16>();
        let p10v = u8x16::from_array(p10).cast::<u16>();
        let p11v = u8x16::from_array(p11).cast::<u16>();

        let sum = p00v + p01v + p10v + p11v;
        let avg = (sum + u16x16::splat(2)).shr(2);
        let out = avg.cast::<u8>();

        dst_row[x_out..x_out + LANES].copy_from_slice(out.as_array());
        x_out += LANES;
    }

    // Tail
    if x_out < dst_row.len() {
        downsample_row_scalar(row0, row1, &mut dst_row[x_out..]);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn reference_downsample(src: &ImageBuffer) -> ImageBuffer {
        let mut dst = ImageBuffer::new(src.width / 2, src.height / 2);
        downsample_2x_box_scalar(src.as_view(), &mut dst);
        dst
    }

    fn gray_to_buffer(w: u32, h: u32, data: Vec<u8>) -> ImageBuffer {
        ImageBuffer {
            width: w as usize,
            height: h as usize,
            data,
        }
    }

    fn make_checker(w: u32, h: u32, a: u8, b: u8) -> ImageBuffer {
        let mut data = vec![0u8; (w * h) as usize];
        for y in 0..h {
            for x in 0..w {
                data[(y * w + x) as usize] = if (x + y) % 2 == 0 { a } else { b };
            }
        }
        gray_to_buffer(w, h, data)
    }

    #[test]
    fn downsample_matches_reference() {
        let mut src = ImageBuffer::new(8, 8);
        for (i, p) in src.data.iter_mut().enumerate() {
            *p = (i % 251) as u8;
        }
        let mut dst = ImageBuffer::new(4, 4);
        downsample_2x_box(src.as_view(), &mut dst);
        let expected = reference_downsample(&src);
        assert_eq!(dst.data, expected.data);
    }

    #[test]
    fn downsample_matches_reference_on_checker() {
        let src = make_checker(16, 14, 0, 255);
        let mut dst = ImageBuffer::new(src.width / 2, src.height / 2);
        downsample_2x_box(src.as_view(), &mut dst);
        let expected = reference_downsample(&src);
        assert_eq!(dst.data, expected.data);
    }

    #[test]
    fn build_pyramid_single_level() {
        let img = ImageBuffer::new(64, 64);
        let params = PyramidParams {
            num_levels: 1,
            min_size: 16,
        };
        let mut buffers = PyramidBuffers::new();
        let pyramid = build_pyramid(img.as_view(), &params, &mut buffers);
        assert_eq!(pyramid.levels.len(), 1);
        assert_eq!(pyramid.levels[0].scale, 1.0);
    }

    #[test]
    fn build_pyramid_multiple_levels() {
        let img = ImageBuffer::new(128, 128);
        let params = PyramidParams {
            num_levels: 4,
            min_size: 16,
        };
        let mut buffers = PyramidBuffers::new();
        let pyramid = build_pyramid(img.as_view(), &params, &mut buffers);
        assert_eq!(pyramid.levels.len(), 4);
        assert_eq!(pyramid.levels[0].img.width, 128);
        assert_eq!(pyramid.levels[1].img.width, 64);
        assert_eq!(pyramid.levels[2].img.width, 32);
        assert_eq!(pyramid.levels[3].img.width, 16);
        assert!((pyramid.levels[3].scale - 0.125).abs() < 1e-6);
    }

    #[test]
    fn build_pyramid_stops_at_min_size() {
        let img = ImageBuffer::new(64, 64);
        let params = PyramidParams {
            num_levels: 10,
            min_size: 32,
        };
        let mut buffers = PyramidBuffers::new();
        let pyramid = build_pyramid(img.as_view(), &params, &mut buffers);
        assert_eq!(pyramid.levels.len(), 2); // 64 -> 32, stops before 16
    }
}