edgefirst_image/cpu/

mod.rs

// SPDX-FileCopyrightText: Copyright 2025 Au-Zone Technologies
// SPDX-License-Identifier: Apache-2.0

use crate::{
    Crop, Error, Flip, FunctionTimer, ImageProcessorTrait, Rect, ResolvedCrop, Result, Rotation,
};
use edgefirst_decoder::{DetectBox, ProtoData, Segmentation};
use edgefirst_tensor::{
    DType, PixelFormat, Tensor, TensorDyn, TensorMapTrait, TensorMemory, TensorTrait,
};

mod convert;
mod masks;
mod resize;
mod simd;
mod tests;

// bilinear_dot removed — masks.rs now uses slice-native bilinear_dot_slice
// closure-based kernel, invoked through the local dtype dispatch below.

/// Resolved colorimetry parameters for a format conversion, computed once at
/// the dispatch site (`convert_impl`) from the appropriate tensor:
/// - YUV→RGB conversions use the **source** tensor's colorimetry.
/// - RGB→YUV conversions use the **destination** tensor's colorimetry.
///
/// `matrix`/`range` feed the `yuv` crate kernels; `src_full_range`/
/// `dst_full_range` gate the hand-rolled luma (limited↔full) expansion in the
/// grey/luma copy helpers.
#[derive(Debug, Clone, Copy)]
pub(crate) struct ColorParams {
    /// Matrix resolved from the YUV side (src for decode, dst for encode).
    pub matrix: yuv::YuvStandardMatrix,
    /// Range resolved from the YUV side (src for decode, dst for encode).
    pub range: yuv::YuvRange,
    /// The same matrix/range in HAL's own terms, so the hand-rolled fixed-point
    /// encoders draw their `(kr, kb)` weights and luma/chroma swings from the
    /// canonical [`edgefirst_tensor::colorimetry`] source rather than a private
    /// duplicate table.
    pub encoding: edgefirst_tensor::ColorEncoding,
    pub range_kind: edgefirst_tensor::ColorRange,
    /// True when the **source** tensor's resolved range is full (decode-side
    /// luma extraction: copy luma directly instead of limited→full expansion).
    pub src_full_range: bool,
    /// True when the **destination** tensor's resolved range is full
    /// (encode-side luma copies).
    pub dst_full_range: bool,
}

/// CPUConverter implements the ImageProcessor trait using the fallback CPU
/// implementation for image processing.
#[derive(Debug)]
pub struct CPUProcessor {
    resizer: fast_image_resize::Resizer,
    options: fast_image_resize::ResizeOptions,
    colors: [[u8; 4]; 20],
    /// Reusable scratch tensor for the U8→float widen path.
    ///
    /// Holds a U8 image in the destination's pixel format and dimensions.
    /// Reallocated only when the format or dimensions change, amortising
    /// the heap allocation across repeated same-size conversions.
    widen_scratch: Option<TensorDyn>,
    /// Reusable scratch for de-striding a padded source before resize.
    ///
    /// `fast_image_resize` needs a tightly-packed input; a padded source
    /// (64-aligned stride from codec decode / `create_image`) is copied here
    /// row-by-row first. Kept on the processor so the steady-state resize loop
    /// reuses the allocation instead of a fresh `Vec` per call (the stride
    /// alignment in this release makes that de-stride copy fire far more often).
    resize_destride_scratch: Vec<u8>,
    /// Reusable cache-resident scratch for the strip-fused NV→planar path
    /// (`convert_nv_to_planar_fused`): holds a single row strip of packed RGB
    /// (`STRIP_ROWS * width * 3` bytes). Keeping it on the processor avoids a
    /// per-frame allocation and lets each strip stay hot in L2 between the YUV
    /// decode and the deinterleave. Grown on demand; never shrunk.
    nv_strip_scratch: Vec<u8>,
    /// Reusable intermediate buffers for the multi-step convert pipeline
    /// (pre-resize format-convert and the resized-RGB scratch). Reused across
    /// frames when the dimensions/format match so the steady-state
    /// letterbox/resize loop does not reallocate — and `alloc_zeroed`-clear — a
    /// full-frame buffer per call. The region each consumer reads is always
    /// fully overwritten first, so reused (non-zeroed) contents are never read.
    convert_tmp: Option<Tensor<u8>>,
    convert_tmp2: Option<Tensor<u8>>,
}

// `CPUProcessor` was `#[derive(Clone)]` before the `widen_scratch` field was
// added; `TensorDyn` is not `Clone`, so the derive no longer applies. Restore
// the public `Clone` impl by hand to avoid a breaking API change. The scratch
// cache is a private allocation amortiser, not part of the logical value, so a
// clone starts empty rather than sharing or duplicating it.
impl Clone for CPUProcessor {
    fn clone(&self) -> Self {
        Self {
            resizer: self.resizer.clone(),
            options: self.options,
            colors: self.colors,
            widen_scratch: None,
            resize_destride_scratch: Vec::new(),
            nv_strip_scratch: Vec::new(),
            convert_tmp: None,
            convert_tmp2: None,
        }
    }
}

unsafe impl Send for CPUProcessor {}
unsafe impl Sync for CPUProcessor {}

impl Default for CPUProcessor {
    fn default() -> Self {
        Self::new_bilinear()
    }
}

/// Write the base layer of `dst` before mask rendering.
///
/// This is the terminal fallback: on CPU we have no 2D hardware, so a
/// direct buffer write is the appropriate primitive. The invariant is that
/// every call to the CPU draw_* entry points fully initialises dst — we
/// never rely on "whatever was in the buffer" from the caller.
///
/// - `background == Some(bg)` → byte-for-byte copy bg → dst (after shape /
///   format validation).
/// - `background == None` → fill dst with 0x00 (transparent black).
fn prepare_dst_base_cpu(dst: &mut TensorDyn, background: Option<&TensorDyn>) -> Result<()> {
    match background {
        Some(bg) => {
            if bg.shape() != dst.shape() {
                return Err(Error::InvalidShape(
                    "background shape does not match dst".into(),
                ));
            }
            if bg.format() != dst.format() {
                return Err(Error::InvalidShape(
                    "background pixel format does not match dst".into(),
                ));
            }
            let bg_u8 = bg.as_u8().ok_or(Error::NotAnImage)?;
            let dst_u8 = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
            let bg_map = bg_u8.map()?;
            let mut dst_map = dst_u8.map()?;
            let bg_slice = bg_map.as_slice();
            let dst_slice = dst_map.as_mut_slice();
            if bg_slice.len() != dst_slice.len() {
                return Err(Error::InvalidShape(
                    "background buffer size does not match dst".into(),
                ));
            }
            dst_slice.copy_from_slice(bg_slice);
        }
        None => {
            let dst_u8 = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
            let mut dst_map = dst_u8.map()?;
            dst_map.as_mut_slice().fill(0);
        }
    }
    Ok(())
}

/// Compute row stride for a packed-format Tensor<u8> image given its format.
fn row_stride_for(width: usize, fmt: PixelFormat) -> usize {
    use edgefirst_tensor::PixelLayout;
    match fmt.layout() {
        PixelLayout::Packed => width * fmt.channels(),
        PixelLayout::Planar | PixelLayout::SemiPlanar => width,
        _ => width, // fallback for non-exhaustive
    }
}

/// Read the effective row stride from a tensor, falling back to the computed
/// minimum stride if the tensor has no explicit stride set. This correctly
/// handles tensors with GPU pitch-alignment padding (e.g., from
/// `ImageProcessor::create_image()` or codec strided decode).
fn tensor_row_stride(tensor: &Tensor<u8>) -> usize {
    tensor.effective_row_stride().unwrap_or_else(|| {
        let w = tensor.width().unwrap_or(0);
        let fmt = tensor.format().unwrap_or(PixelFormat::Rgb);
        row_stride_for(w, fmt)
    })
}

/// Split a contiguous semi-planar (NV12/NV16/NV24) tensor's mapped bytes into
/// `(luma, chroma)` planes at the `stride * src_h` boundary, validating the map
/// holds the full combined plane first. A bare `split_at` would panic if an
/// imported tensor's caller-supplied dimensions/stride exceed its actual buffer
/// (untrusted input), so this returns `Error::InvalidShape` instead.
fn split_semi_planar(
    bytes: &[u8],
    stride: usize,
    src_h: usize,
    fmt: PixelFormat,
) -> Result<(&[u8], &[u8])> {
    let total_h = fmt.combined_plane_height(src_h).unwrap_or(src_h);
    let need = stride.checked_mul(total_h).ok_or_else(|| {
        Error::InvalidShape(format!(
            "{fmt:?} plane size overflow (stride={stride}, h={src_h})"
        ))
    })?;
    if bytes.len() < need {
        return Err(Error::InvalidShape(format!(
            "{fmt:?} source has {} bytes but needs {need} (stride={stride}, h={src_h})",
            bytes.len()
        )));
    }
    Ok(bytes.split_at(stride * src_h))
}

/// Mutable mirror of [`split_semi_planar`]: split a contiguous semi-planar
/// destination's mapped bytes into `(luma, chroma)` planes at the
/// `stride * dst_h` boundary, validating the map holds the full combined plane
/// first. A bare `split_at_mut` would panic if a caller-supplied (untrusted)
/// destination's declared dimensions/stride exceed its actual buffer, so this
/// returns `Error::InvalidShape` instead.
fn split_semi_planar_mut(
    bytes: &mut [u8],
    stride: usize,
    dst_h: usize,
    fmt: PixelFormat,
) -> Result<(&mut [u8], &mut [u8])> {
    let total_h = fmt.combined_plane_height(dst_h).unwrap_or(dst_h);
    let need = stride.checked_mul(total_h).ok_or_else(|| {
        Error::InvalidShape(format!(
            "{fmt:?} plane size overflow (stride={stride}, combined_h={total_h})"
        ))
    })?;
    if bytes.len() < need {
        return Err(Error::InvalidShape(format!(
            "{fmt:?} destination has {} bytes but needs {need} (stride={stride}, combined_h={total_h})",
            bytes.len()
        )));
    }
    Ok(bytes.split_at_mut(stride * dst_h))
}

/// Validate that a mapped plane buffer of `buf_len` bytes can hold `rows` rows
/// of `stride` bytes each, with `row_bytes` valid bytes per row. Returns
/// `Error::InvalidShape` (instead of letting a later row slice panic) when a
/// caller-supplied stride/shape exceeds the actual allocation. `what` labels the
/// buffer in the error message.
fn guard_plane(
    buf_len: usize,
    stride: usize,
    rows: usize,
    row_bytes: usize,
    what: &str,
) -> Result<()> {
    let need = stride.checked_mul(rows).ok_or_else(|| {
        Error::InvalidShape(format!(
            "{what} plane size overflow (stride={stride}, rows={rows})"
        ))
    })?;
    if row_bytes > stride || buf_len < need {
        return Err(Error::InvalidShape(format!(
            "{what} buffer too small: {buf_len} bytes, need {need} (stride={stride}, rows={rows}, row_bytes={row_bytes})"
        )));
    }
    Ok(())
}

/// Apply XOR 0x80 bias to color channels only, preserving alpha.
///
/// Matches GL int8 shader behavior: `vec4(int8_bias(c.rgb), c.a)`.
/// For formats without alpha, XORs every byte (fast path).
pub(crate) fn apply_int8_xor_bias(data: &mut [u8], fmt: PixelFormat) {
    use edgefirst_tensor::PixelLayout;
    if !fmt.has_alpha() {
        for b in data.iter_mut() {
            *b ^= 0x80;
        }
    } else if fmt.layout() == PixelLayout::Planar {
        // Planar with alpha (e.g. PlanarRgba): XOR color planes, skip alpha plane.
        let channels = fmt.channels();
        let plane_size = data.len() / channels;
        for b in data[..plane_size * (channels - 1)].iter_mut() {
            *b ^= 0x80;
        }
    } else {
        // Packed with alpha (Rgba, Bgra): XOR color bytes, skip alpha byte.
        let channels = fmt.channels();
        for pixel in data.chunks_exact_mut(channels) {
            for b in &mut pixel[..channels - 1] {
                *b ^= 0x80;
            }
        }
    }
}

impl CPUProcessor {
    /// Creates a new CPUConverter with bilinear resizing.
    pub fn new() -> Self {
        Self::new_bilinear()
    }

    /// Creates a new CPUConverter with bilinear resizing.
    fn new_bilinear() -> Self {
        let resizer = fast_image_resize::Resizer::new();
        let options = fast_image_resize::ResizeOptions::new()
            .resize_alg(fast_image_resize::ResizeAlg::Convolution(
                fast_image_resize::FilterType::Bilinear,
            ))
            .use_alpha(false);

        log::debug!("CPUConverter created");
        Self {
            resizer,
            options,
            colors: crate::DEFAULT_COLORS_U8,
            widen_scratch: None,
            resize_destride_scratch: Vec::new(),
            nv_strip_scratch: Vec::new(),
            convert_tmp: None,
            convert_tmp2: None,
        }
    }

    /// Creates a new CPUConverter with nearest neighbor resizing.
    pub fn new_nearest() -> Self {
        let resizer = fast_image_resize::Resizer::new();
        let options = fast_image_resize::ResizeOptions::new()
            .resize_alg(fast_image_resize::ResizeAlg::Nearest)
            .use_alpha(false);
        log::debug!("CPUConverter created");
        Self {
            resizer,
            options,
            colors: crate::DEFAULT_COLORS_U8,
            widen_scratch: None,
            resize_destride_scratch: Vec::new(),
            nv_strip_scratch: Vec::new(),
            convert_tmp: None,
            convert_tmp2: None,
        }
    }

    pub(crate) fn support_conversion_pf(src: PixelFormat, dst: PixelFormat) -> bool {
        use PixelFormat::*;
        matches!(
            (src, dst),
            (Nv12, Rgb)
                | (Nv12, Rgba)
                | (Nv12, Grey)
                | (Nv16, Rgb)
                | (Nv16, Rgba)
                | (Nv16, Bgra)
                | (Nv24, Rgb)
                | (Nv24, Rgba)
                | (Nv24, Grey)
                | (Nv24, Bgra)
                | (Yuyv, Rgb)
                | (Yuyv, Rgba)
                | (Yuyv, Grey)
                | (Yuyv, Yuyv)
                | (Yuyv, PlanarRgb)
                | (Yuyv, PlanarRgba)
                | (Yuyv, Nv16)
                | (Vyuy, Rgb)
                | (Vyuy, Rgba)
                | (Vyuy, Grey)
                | (Vyuy, Vyuy)
                | (Vyuy, PlanarRgb)
                | (Vyuy, PlanarRgba)
                | (Vyuy, Nv16)
                | (Rgba, Rgb)
                | (Rgba, Rgba)
                | (Rgba, Grey)
                | (Rgba, Yuyv)
                | (Rgba, PlanarRgb)
                | (Rgba, PlanarRgba)
                | (Rgba, Nv16)
                | (Rgb, Rgb)
                | (Rgb, Rgba)
                | (Rgb, Grey)
                | (Rgb, Yuyv)
                | (Rgb, PlanarRgb)
                | (Rgb, PlanarRgba)
                | (Rgb, Nv16)
                | (Grey, Rgb)
                | (Grey, Rgba)
                | (Grey, Grey)
                | (Grey, Yuyv)
                | (Grey, PlanarRgb)
                | (Grey, PlanarRgba)
                | (Grey, Nv16)
                | (Nv12, Bgra)
                | (Yuyv, Bgra)
                | (Vyuy, Bgra)
                | (Rgba, Bgra)
                | (Rgb, Bgra)
                | (Grey, Bgra)
                | (Bgra, Bgra)
                | (PlanarRgb, Rgb)
                | (PlanarRgb, Rgba)
                | (PlanarRgba, Rgb)
                | (PlanarRgba, Rgba)
                | (PlanarRgb, Bgra)
                | (PlanarRgba, Bgra)
        )
    }

    /// Format conversion dispatch for Tensor<u8> with PixelFormat metadata.
    pub(crate) fn convert_format_pf(
        src: &Tensor<u8>,
        dst: &mut Tensor<u8>,
        src_fmt: PixelFormat,
        dst_fmt: PixelFormat,
        cp: ColorParams,
    ) -> Result<()> {
        let _timer = FunctionTimer::new(format!(
            "ImageProcessor::convert_format {} to {}",
            src_fmt, dst_fmt,
        ));

        use PixelFormat::*;
        match (src_fmt, dst_fmt) {
            (Nv12, Rgb) => Self::convert_nv12_to_rgb(src, dst, cp),
            (Nv12, Rgba) => Self::convert_nv12_to_rgba(src, dst, cp),
            (Nv12, Grey) => Self::convert_nv12_to_grey(src, dst, cp),
            (Yuyv, Rgb) => Self::convert_yuyv_to_rgb(src, dst, cp),
            (Yuyv, Rgba) => Self::convert_yuyv_to_rgba(src, dst, cp),
            (Yuyv, Grey) => Self::convert_yuyv_to_grey(src, dst, cp),
            (Yuyv, Yuyv) => Self::copy_image(src, dst),
            (Yuyv, PlanarRgb) => Self::convert_yuyv_to_8bps(src, dst, cp),
            (Yuyv, PlanarRgba) => Self::convert_yuyv_to_prgba(src, dst, cp),
            (Yuyv, Nv16) => Self::convert_yuyv_to_nv16(src, dst),
            (Vyuy, Rgb) => Self::convert_vyuy_to_rgb(src, dst, cp),
            (Vyuy, Rgba) => Self::convert_vyuy_to_rgba(src, dst, cp),
            (Vyuy, Grey) => Self::convert_vyuy_to_grey(src, dst, cp),
            (Vyuy, Vyuy) => Self::copy_image(src, dst),
            (Vyuy, PlanarRgb) => Self::convert_vyuy_to_8bps(src, dst, cp),
            (Vyuy, PlanarRgba) => Self::convert_vyuy_to_prgba(src, dst, cp),
            (Vyuy, Nv16) => Self::convert_vyuy_to_nv16(src, dst),
            (Rgba, Rgb) => Self::convert_rgba_to_rgb(src, dst),
            (Rgba, Rgba) => Self::copy_image(src, dst),
            (Rgba, Grey) => Self::convert_rgba_to_grey(src, dst),
            (Rgba, Yuyv) => Self::convert_rgba_to_yuyv(src, dst, cp),
            (Rgba, PlanarRgb) => Self::convert_rgba_to_8bps(src, dst),
            (Rgba, PlanarRgba) => Self::convert_rgba_to_prgba(src, dst),
            (Rgba, Nv16) => Self::convert_rgba_to_nv16(src, dst, cp),
            (Rgb, Rgb) => Self::copy_image(src, dst),
            (Rgb, Rgba) => Self::convert_rgb_to_rgba(src, dst),
            (Rgb, Grey) => Self::convert_rgb_to_grey(src, dst),
            (Rgb, Yuyv) => Self::convert_rgb_to_yuyv(src, dst, cp),
            (Rgb, PlanarRgb) => Self::convert_rgb_to_8bps(src, dst),
            (Rgb, PlanarRgba) => Self::convert_rgb_to_prgba(src, dst),
            (Rgb, Nv16) => Self::convert_rgb_to_nv16(src, dst, cp),
            (Grey, Rgb) => Self::convert_grey_to_rgb(src, dst),
            (Grey, Rgba) => Self::convert_grey_to_rgba(src, dst),
            (Grey, Grey) => Self::copy_image(src, dst),
            (Grey, Yuyv) => Self::convert_grey_to_yuyv(src, dst, cp),
            (Grey, PlanarRgb) => Self::convert_grey_to_8bps(src, dst),
            (Grey, PlanarRgba) => Self::convert_grey_to_prgba(src, dst),
            (Grey, Nv16) => Self::convert_grey_to_nv16(src, dst, cp),

            // the following converts are added for use in testing
            (Nv16, Rgb) => Self::convert_nv16_to_rgb(src, dst, cp),
            (Nv16, Rgba) => Self::convert_nv16_to_rgba(src, dst, cp),
            (Nv24, Rgb) => Self::convert_nv24_to_rgb(src, dst, cp),
            (Nv24, Rgba) => Self::convert_nv24_to_rgba(src, dst, cp),
            (Nv24, Grey) => Self::convert_nv24_to_grey(src, dst, cp),
            (PlanarRgb, Rgb) => Self::convert_8bps_to_rgb(src, dst),
            (PlanarRgb, Rgba) => Self::convert_8bps_to_rgba(src, dst),
            (PlanarRgba, Rgb) => Self::convert_prgba_to_rgb(src, dst),
            (PlanarRgba, Rgba) => Self::convert_prgba_to_rgba(src, dst),

            // BGRA destination: convert to RGBA layout, then swap R and B
            (Bgra, Bgra) => Self::copy_image(src, dst),
            (Nv12, Bgra) => {
                Self::convert_nv12_to_rgba(src, dst, cp)?;
                Self::swizzle_rb_4chan(dst)
            }
            (Nv16, Bgra) => {
                Self::convert_nv16_to_rgba(src, dst, cp)?;
                Self::swizzle_rb_4chan(dst)
            }
            (Nv24, Bgra) => {
                Self::convert_nv24_to_rgba(src, dst, cp)?;
                Self::swizzle_rb_4chan(dst)
            }
            (Yuyv, Bgra) => {
                Self::convert_yuyv_to_rgba(src, dst, cp)?;
                Self::swizzle_rb_4chan(dst)
            }
            (Vyuy, Bgra) => {
                Self::convert_vyuy_to_rgba(src, dst, cp)?;
                Self::swizzle_rb_4chan(dst)
            }
            (Rgba, Bgra) => {
                dst.map()?.copy_from_slice(&src.map()?);
                Self::swizzle_rb_4chan(dst)
            }
            (Rgb, Bgra) => {
                Self::convert_rgb_to_rgba(src, dst)?;
                Self::swizzle_rb_4chan(dst)
            }
            (Grey, Bgra) => {
                Self::convert_grey_to_rgba(src, dst)?;
                Self::swizzle_rb_4chan(dst)
            }
            (PlanarRgb, Bgra) => {
                Self::convert_8bps_to_rgba(src, dst)?;
                Self::swizzle_rb_4chan(dst)
            }
            (PlanarRgba, Bgra) => {
                Self::convert_prgba_to_rgba(src, dst)?;
                Self::swizzle_rb_4chan(dst)
            }

            (s, d) => Err(Error::NotSupported(format!("Conversion from {s} to {d}",))),
        }
    }

    /// Tensor<u8>-based fill_image_outside_crop.
    pub(crate) fn fill_image_outside_crop_u8(
        dst: &mut Tensor<u8>,
        rgba: [u8; 4],
        crop: Rect,
    ) -> Result<()> {
        let dst_fmt = dst.format().unwrap();
        let dst_w = dst.width().unwrap();
        let dst_h = dst.height().unwrap();
        // Resolve the YUV fill encoding from the destination tensor so the
        // border color matches the same matrix/range as a later YUV→RGB
        // decode of this image. RGB/Grey fills ignore these params.
        let cm = crate::colorimetry::resolve_colorimetry(dst.colorimetry(), dst.height());
        let cp = ColorParams {
            matrix: crate::colorimetry::yuv_matrix(cm.encoding.unwrap()),
            range: crate::colorimetry::yuv_range(cm.range.unwrap()),
            encoding: cm.encoding.unwrap(),
            range_kind: cm.range.unwrap(),
            src_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
            dst_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
        };
        let mut dst_map = dst.map()?;
        let dst_tup = (dst_map.as_mut_slice(), dst_w, dst_h);
        Self::fill_outside_crop_dispatch(dst_tup, dst_fmt, rgba, crop, cp)
    }

    /// Common fill dispatch by format.
    fn fill_outside_crop_dispatch(
        dst: (&mut [u8], usize, usize),
        fmt: PixelFormat,
        rgba: [u8; 4],
        crop: Rect,
        cp: ColorParams,
    ) -> Result<()> {
        use PixelFormat::*;
        match fmt {
            Rgba | Bgra => Self::fill_image_outside_crop_(dst, rgba, crop),
            Rgb => Self::fill_image_outside_crop_(dst, Self::rgba_to_rgb(rgba), crop),
            Grey => Self::fill_image_outside_crop_(dst, Self::rgba_to_grey(rgba), crop),
            Yuyv => Self::fill_image_outside_crop_(
                (dst.0, dst.1 / 2, dst.2),
                Self::rgba_to_yuyv(rgba, cp),
                Rect::new(crop.left / 2, crop.top, crop.width.div_ceil(2), crop.height),
            ),
            PlanarRgb => Self::fill_image_outside_crop_planar(dst, Self::rgba_to_rgb(rgba), crop),
            PlanarRgba => Self::fill_image_outside_crop_planar(dst, rgba, crop),
            Nv16 => {
                let yuyv = Self::rgba_to_yuyv(rgba, cp);
                Self::fill_image_outside_crop_yuv_semiplanar(dst, yuyv[0], [yuyv[1], yuyv[3]], crop)
            }
            _ => Err(Error::Internal(format!(
                "Found unexpected destination {fmt}",
            ))),
        }
    }
}

impl ImageProcessorTrait for CPUProcessor {
    fn convert(
        &mut self,
        src: &TensorDyn,
        dst: &mut TensorDyn,
        rotation: Rotation,
        flip: Flip,
        crop: Crop,
    ) -> Result<()> {
        let crop = crop.resolve(
            src.width().unwrap_or(0),
            src.height().unwrap_or(0),
            dst.width().unwrap_or(0),
            dst.height().unwrap_or(0),
        )?;
        self.convert_impl(src, dst, rotation, flip, crop)
    }

    fn draw_decoded_masks(
        &mut self,
        dst: &mut TensorDyn,
        detect: &[DetectBox],
        segmentation: &[Segmentation],
        overlay: crate::MaskOverlay<'_>,
    ) -> Result<()> {
        // CPU is the terminal fallback — it must always produce the full
        // output, never assume the caller cleared dst. Every call writes
        // the base layer first (bg copy or zero fill) and then the masks.
        prepare_dst_base_cpu(dst, overlay.background)?;
        let dst = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
        self.draw_decoded_masks_impl(
            dst,
            detect,
            segmentation,
            overlay.opacity,
            overlay.color_mode,
        )
    }

    fn draw_proto_masks(
        &mut self,
        dst: &mut TensorDyn,
        detect: &[DetectBox],
        proto_data: &ProtoData,
        overlay: crate::MaskOverlay<'_>,
    ) -> Result<()> {
        prepare_dst_base_cpu(dst, overlay.background)?;
        let dst = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
        self.draw_proto_masks_impl(
            dst,
            detect,
            proto_data,
            overlay.opacity,
            overlay.letterbox,
            overlay.color_mode,
        )
    }

    fn set_class_colors(&mut self, colors: &[[u8; 4]]) -> Result<()> {
        for (c, new_c) in self.colors.iter_mut().zip(colors.iter()) {
            *c = *new_c;
        }
        Ok(())
    }
}

// Internal methods — dtype-aware dispatch layer.
impl CPUProcessor {
    /// Top-level conversion dispatcher: handles dtype combinations.
    pub(crate) fn convert_impl(
        &mut self,
        src: &TensorDyn,
        dst: &mut TensorDyn,
        rotation: Rotation,
        flip: Flip,
        crop: ResolvedCrop,
    ) -> Result<()> {
        let src_fmt = src.format().ok_or(Error::NotAnImage)?;
        let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;

        // Resolve per-tensor colorimetry once, at the use site, without
        // mutating either tensor. YUV→RGB conversions take their matrix/range
        // from the source; RGB→YUV conversions from the destination. The grey/
        // luma expansion is gated on the resolved range of the relevant side.
        let src_cm = crate::colorimetry::effective_colorimetry(src);
        let dst_cm = crate::colorimetry::effective_colorimetry(dst);
        let src_full = src_cm.range == Some(edgefirst_tensor::ColorRange::Full);
        let dst_full = dst_cm.range == Some(edgefirst_tensor::ColorRange::Full);
        let src_params = ColorParams {
            matrix: crate::colorimetry::yuv_matrix(src_cm.encoding.unwrap()),
            range: crate::colorimetry::yuv_range(src_cm.range.unwrap()),
            encoding: src_cm.encoding.unwrap(),
            range_kind: src_cm.range.unwrap(),
            src_full_range: src_full,
            dst_full_range: dst_full,
        };
        let dst_params = ColorParams {
            matrix: crate::colorimetry::yuv_matrix(dst_cm.encoding.unwrap()),
            range: crate::colorimetry::yuv_range(dst_cm.range.unwrap()),
            encoding: dst_cm.encoding.unwrap(),
            range_kind: dst_cm.range.unwrap(),
            src_full_range: src_full,
            dst_full_range: dst_full,
        };
        match (src.dtype(), dst.dtype()) {
            (DType::U8, DType::U8) => {
                let src = src.as_u8().unwrap();
                let dst = dst.as_u8_mut().unwrap();
                self.convert_u8(
                    src, dst, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
                )
            }
            (DType::U8, DType::I8) => {
                // Int8 output: reinterpret the i8 destination as u8 (layout-
                // identical), convert directly into it, then XOR 0x80 in-place.
                let src_u8 = src.as_u8().unwrap();
                let dst_i8 = dst.as_i8_mut().unwrap();
                // SAFETY: Tensor<i8> and Tensor<u8> are layout-identical
                // (same element size, no T-dependent drop glue). Same
                // rationale as gl::processor::tensor_i8_as_u8_mut.
                let dst_u8 = unsafe { &mut *(dst_i8 as *mut Tensor<i8> as *mut Tensor<u8>) };
                self.convert_u8(
                    src_u8, dst_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
                )?;
                // Apply XOR 0x80 bias in-place (u8 → i8 conversion)
                let mut map = dst_u8.map()?;
                apply_int8_xor_bias(map.as_mut_slice(), dst_fmt);
                Ok(())
            }
            (DType::U8, d @ (DType::F32 | DType::F16)) => {
                let src_u8 = src.as_u8().unwrap();
                let dw = dst.width().ok_or(Error::NotAnImage)?;
                let dh = dst.height().ok_or(Error::NotAnImage)?;
                // Reuse the scratch tensor when format and dimensions match;
                // otherwise reallocate and cache the new scratch.  Take the
                // scratch out of `self` so that `convert_u8` can borrow `self`
                // exclusively, then restore it afterwards.
                let scratch_matches = self.widen_scratch.as_ref().is_some_and(|t| {
                    t.width() == Some(dw) && t.height() == Some(dh) && t.format() == Some(dst_fmt)
                });
                let mut tmp = if scratch_matches {
                    self.widen_scratch.take().unwrap()
                } else {
                    TensorDyn::image(dw, dh, dst_fmt, DType::U8, Some(TensorMemory::Mem))?
                };
                {
                    let tmp_u8 = tmp.as_u8_mut().unwrap();
                    self.convert_u8(
                        src_u8, tmp_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params,
                        dst_params,
                    )?;
                }
                // Widen the u8 scratch into the float destination, then restore
                // the scratch for reuse on the next call.
                {
                    let tmp_u8 = tmp.as_u8().unwrap();
                    let src_map = tmp_u8.map()?;
                    match d {
                        DType::F32 => {
                            let dst_t = dst.as_f32_mut().unwrap();
                            let mut dst_map = dst_t.map()?;
                            debug_assert_eq!(src_map.as_slice().len(), dst_map.as_slice().len());
                            // NEON-accelerated u8→f32 `/255` widen (bit-identical
                            // to the scalar `b as f32 / 255.0`); the scalar
                            // iterator form did not vectorise. See cpu::simd.
                            simd::widen_u8_to_f32_norm(src_map.as_slice(), dst_map.as_mut_slice());
                        }
                        DType::F16 => {
                            let dst_t = dst.as_f16_mut().unwrap();
                            let mut dst_map = dst_t.map()?;
                            debug_assert_eq!(src_map.as_slice().len(), dst_map.as_slice().len());
                            // u8→f16 `/255` widen; uses native FP16
                            // (`ucvtf`+`fdiv`) at runtime on FEAT_FP16 CPUs
                            // (Orin), scalar `half::f16::from_f32` elsewhere.
                            // See cpu::simd.
                            simd::widen_u8_to_f16_norm(src_map.as_slice(), dst_map.as_mut_slice());
                        }
                        _ => unreachable!(),
                    }
                }
                self.widen_scratch = Some(tmp);
                Ok(())
            }
            (s, d) => Err(Error::NotSupported(format!("dtype {s} -> {d}",))),
        }
    }

    /// Reuse `cached` if it already has dimensions `(w, h)` and pixel format
    /// `fmt`, otherwise allocate a fresh `Mem` image. The returned buffer's
    /// contents are **not** zeroed on reuse — callers must fully overwrite the
    /// region they later read (see the note at the convert pipeline's tmp/tmp2
    /// site). This amortises the per-frame `Tensor::image` `alloc_zeroed`.
    fn reuse_or_alloc_image(
        cached: Option<Tensor<u8>>,
        w: usize,
        h: usize,
        fmt: PixelFormat,
    ) -> Result<Tensor<u8>> {
        if let Some(t) = cached {
            if t.width() == Some(w) && t.height() == Some(h) && t.format() == Some(fmt) {
                return Ok(t);
            }
        }
        Ok(Tensor::<u8>::image(w, h, fmt, Some(TensorMemory::Mem))?)
    }

    /// U8-to-U8 conversion: the full format conversion + resize pipeline.
    #[allow(clippy::too_many_arguments)]
    fn convert_u8(
        &mut self,
        src: &Tensor<u8>,
        dst: &mut Tensor<u8>,
        src_fmt: PixelFormat,
        dst_fmt: PixelFormat,
        rotation: Rotation,
        flip: Flip,
        crop: ResolvedCrop,
        src_params: ColorParams,
        dst_params: ColorParams,
    ) -> Result<()> {
        use PixelFormat::*;

        let src_w = src.width().unwrap();
        let src_h = src.height().unwrap();
        let dst_w = dst.width().unwrap();
        let dst_h = dst.height().unwrap();

        crop.check_crop_dims(src_w, src_h, dst_w, dst_h)?;

        // Determine intermediate format for the resize step
        let intermediate = match (src_fmt, dst_fmt) {
            (Nv12, Rgb) => Rgb,
            (Nv12, Rgba) => Rgba,
            (Nv12, Grey) => Grey,
            (Nv12, Yuyv) => Rgba,
            (Nv12, Nv16) => Rgba,
            (Nv12, PlanarRgb) => Rgb,
            (Nv12, PlanarRgba) => Rgba,
            (Nv16, PlanarRgb) => Rgb,
            (Nv16, PlanarRgba) => Rgba,
            (Nv24, PlanarRgb) => Rgb,
            (Nv24, PlanarRgba) => Rgba,
            (Yuyv, Rgb) => Rgb,
            (Yuyv, Rgba) => Rgba,
            (Yuyv, Grey) => Grey,
            (Yuyv, Yuyv) => Rgba,
            (Yuyv, PlanarRgb) => Rgb,
            (Yuyv, PlanarRgba) => Rgba,
            (Yuyv, Nv16) => Rgba,
            (Vyuy, Rgb) => Rgb,
            (Vyuy, Rgba) => Rgba,
            (Vyuy, Grey) => Grey,
            (Vyuy, Vyuy) => Rgba,
            (Vyuy, PlanarRgb) => Rgb,
            (Vyuy, PlanarRgba) => Rgba,
            (Vyuy, Nv16) => Rgba,
            (Rgba, Rgb) => Rgba,
            (Rgba, Rgba) => Rgba,
            (Rgba, Grey) => Grey,
            (Rgba, Yuyv) => Rgba,
            (Rgba, PlanarRgb) => Rgba,
            (Rgba, PlanarRgba) => Rgba,
            (Rgba, Nv16) => Rgba,
            (Rgb, Rgb) => Rgb,
            (Rgb, Rgba) => Rgb,
            (Rgb, Grey) => Grey,
            (Rgb, Yuyv) => Rgb,
            (Rgb, PlanarRgb) => Rgb,
            (Rgb, PlanarRgba) => Rgb,
            (Rgb, Nv16) => Rgb,
            (Grey, Rgb) => Rgb,
            (Grey, Rgba) => Rgba,
            (Grey, Grey) => Grey,
            (Grey, Yuyv) => Grey,
            (Grey, PlanarRgb) => Grey,
            (Grey, PlanarRgba) => Grey,
            (Grey, Nv16) => Grey,
            (Nv12, Bgra) => Rgba,
            (Yuyv, Bgra) => Rgba,
            (Vyuy, Bgra) => Rgba,
            (Rgba, Bgra) => Rgba,
            (Rgb, Bgra) => Rgb,
            (Grey, Bgra) => Grey,
            (Bgra, Bgra) => Bgra,
            (Nv16, Rgb) => Rgb,
            (Nv16, Rgba) => Rgba,
            (Nv16, Bgra) => Rgba,
            (Nv24, Rgb) => Rgb,
            (Nv24, Rgba) => Rgba,
            (Nv24, Grey) => Grey,
            (Nv24, Bgra) => Rgba,
            (PlanarRgb, Rgb) => Rgb,
            (PlanarRgb, Rgba) => Rgb,
            (PlanarRgb, Bgra) => Rgb,
            (PlanarRgba, Rgb) => Rgba,
            (PlanarRgba, Rgba) => Rgba,
            (PlanarRgba, Bgra) => Rgba,
            (s, d) => {
                return Err(Error::NotSupported(format!("Conversion from {s} to {d}",)));
            }
        };

        let need_resize_flip_rotation = rotation != Rotation::None
            || flip != Flip::None
            || src_w != dst_w
            || src_h != dst_h
            || crop.src_rect.is_some_and(|c| {
                c != Rect {
                    left: 0,
                    top: 0,
                    width: src_w,
                    height: src_h,
                }
            })
            || crop.dst_rect.is_some_and(|c| {
                c != Rect {
                    left: 0,
                    top: 0,
                    width: dst_w,
                    height: dst_h,
                }
            });

        // Pick the resolved colorimetry for a single conversion by its YUV
        // side: YUV→RGB decodes use the source side, RGB→YUV encodes the dest.
        let direct_is_yuv_src = matches!(src_fmt, Nv12 | Nv16 | Nv24 | Yuyv | Vyuy);
        let direct_params = if direct_is_yuv_src {
            src_params
        } else {
            dst_params
        };

        // Fused NV→planar (no resize/flip/rotation/crop): decode the YUV source
        // into packed RGB one cache-resident row strip at a time and
        // NEON-deinterleave each strip straight into the destination planes, so
        // the full-size packed-RGB intermediate never round-trips through DRAM
        // and is not reallocated per frame. JPEG decodes to the NV family and the
        // model wants planar RGB, so this is the hot Orin CPU-preprocess path.
        if !need_resize_flip_rotation
            && matches!(src_fmt, Nv12 | Nv16 | Nv24)
            && matches!(dst_fmt, PlanarRgb | PlanarRgba)
        {
            return self.convert_nv_to_planar_fused(src, dst, src_fmt, dst_fmt, direct_params);
        }

        // check if a direct conversion can be done
        if !need_resize_flip_rotation && Self::support_conversion_pf(src_fmt, dst_fmt) {
            return Self::convert_format_pf(src, dst, src_fmt, dst_fmt, direct_params);
        }

        // any extra checks
        if dst_fmt == Yuyv && !dst_w.is_multiple_of(2) {
            return Err(Error::NotSupported(format!(
                "{} destination must have width divisible by 2",
                dst_fmt,
            )));
        }

        // Take the cached intermediates out of `self` so the resize step can
        // borrow `self` exclusively; they are restored before returning on the
        // success path. Reused buffers are not re-zeroed — every consumer below
        // fully overwrites the region it later reads (the pre-resize convert
        // writes all of `tmp`; the resize writes the scaled rect of `tmp2` and
        // its letterbox border is either pre-filled from `dst` or overwritten in
        // `dst` by the final `fill_image_outside_crop_u8`).
        let mut cached_tmp = self.convert_tmp.take();
        let mut cached_tmp2 = self.convert_tmp2.take();

        // create tmp buffer (reusing the cached one when its geometry matches)
        let tmp_holder: Option<Tensor<u8>> = if intermediate != src_fmt {
            let _s = tracing::trace_span!(
                "image.convert.cpu.format_convert",
                from = ?src_fmt,
                to = ?intermediate,
                pass = "pre_resize",
            )
            .entered();
            let mut t = Self::reuse_or_alloc_image(cached_tmp.take(), src_w, src_h, intermediate)?;
            Self::convert_format_pf(src, &mut t, src_fmt, intermediate, src_params)?;
            Some(t)
        } else {
            None
        };
        let (tmp, tmp_fmt): (&Tensor<u8>, PixelFormat) = match &tmp_holder {
            Some(t) => (t, intermediate),
            None => (src, src_fmt),
        };

        // format must be RGB/RGBA/GREY
        debug_assert!(matches!(tmp_fmt, Rgb | Rgba | Grey));
        if tmp_fmt == dst_fmt {
            let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
            self.resize_flip_rotate_pf(tmp, dst, dst_fmt, rotation, flip, crop)?;
        } else if !need_resize_flip_rotation {
            let _s = tracing::trace_span!(
                "image.convert.cpu.format_convert",
                from = ?tmp_fmt,
                to = ?dst_fmt,
                pass = "direct",
            )
            .entered();
            Self::convert_format_pf(tmp, dst, tmp_fmt, dst_fmt, dst_params)?;
        } else {
            let mut tmp2 = Self::reuse_or_alloc_image(cached_tmp2.take(), dst_w, dst_h, tmp_fmt)?;
            if crop.dst_rect.is_some_and(|c| {
                c != Rect {
                    left: 0,
                    top: 0,
                    width: dst_w,
                    height: dst_h,
                }
            }) && crop.dst_color.is_none()
            {
                Self::convert_format_pf(dst, &mut tmp2, dst_fmt, tmp_fmt, dst_params)?;
            }
            {
                let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
                self.resize_flip_rotate_pf(tmp, &mut tmp2, tmp_fmt, rotation, flip, crop)?;
            }
            {
                let _s = tracing::trace_span!(
                    "image.convert.cpu.format_convert",
                    from = ?tmp_fmt,
                    to = ?dst_fmt,
                    pass = "post_resize",
                )
                .entered();
                Self::convert_format_pf(&tmp2, dst, tmp_fmt, dst_fmt, dst_params)?;
            }
            cached_tmp2 = Some(tmp2);
        }
        // Restore the intermediates to the cache for the next call (`tmp` — a
        // borrow of `tmp_holder` — is no longer used past this point).
        if let Some(t) = tmp_holder {
            cached_tmp = Some(t);
        }
        self.convert_tmp = cached_tmp;
        self.convert_tmp2 = cached_tmp2;

        if let (Some(dst_rect), Some(dst_color)) = (crop.dst_rect, crop.dst_color) {
            let full_rect = Rect {
                left: 0,
                top: 0,
                width: dst_w,
                height: dst_h,
            };
            if dst_rect != full_rect {
                Self::fill_image_outside_crop_u8(dst, dst_color, dst_rect)?;
            }
        }

        Ok(())
    }

    fn draw_decoded_masks_impl(
        &mut self,
        dst: &mut Tensor<u8>,
        detect: &[DetectBox],
        segmentation: &[Segmentation],
        opacity: f32,
        color_mode: crate::ColorMode,
    ) -> Result<()> {
        let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
        if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
            return Err(crate::Error::NotSupported(
                "CPU image rendering only supports RGBA or RGB images".to_string(),
            ));
        }

        let _timer = FunctionTimer::new("CPUProcessor::draw_decoded_masks");

        let dst_w = dst.width().unwrap();
        let dst_h = dst.height().unwrap();
        let dst_rs = tensor_row_stride(dst);
        let dst_c = dst_fmt.channels();

        let mut map = dst.map()?;
        let dst_slice = map.as_mut_slice();

        self.render_box(dst_w, dst_h, dst_rs, dst_c, dst_slice, detect, color_mode)?;

        if segmentation.is_empty() {
            return Ok(());
        }

        // Semantic segmentation (e.g. ModelPack) has C > 1 (multi-class),
        // instance segmentation (e.g. YOLO) has C = 1 (binary per-instance).
        let is_semantic = segmentation[0].segmentation.shape()[2] > 1;

        if is_semantic {
            self.render_modelpack_segmentation(
                dst_w,
                dst_h,
                dst_rs,
                dst_c,
                dst_slice,
                &segmentation[0],
                opacity,
            )?;
        } else {
            for (idx, (seg, det)) in segmentation.iter().zip(detect).enumerate() {
                let color_index = color_mode.index(idx, det.label);
                self.render_yolo_segmentation(
                    dst_w,
                    dst_h,
                    dst_rs,
                    dst_c,
                    dst_slice,
                    seg,
                    color_index,
                    opacity,
                )?;
            }
        }

        Ok(())
    }

    fn draw_proto_masks_impl(
        &mut self,
        dst: &mut Tensor<u8>,
        detect: &[DetectBox],
        proto_data: &ProtoData,
        opacity: f32,
        letterbox: Option<[f32; 4]>,
        color_mode: crate::ColorMode,
    ) -> Result<()> {
        let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
        if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
            return Err(crate::Error::NotSupported(
                "CPU image rendering only supports RGBA or RGB images".to_string(),
            ));
        }

        let _timer = FunctionTimer::new("CPUProcessor::draw_proto_masks");

        let dst_w = dst.width().unwrap();
        let dst_h = dst.height().unwrap();
        let dst_rs = tensor_row_stride(dst);
        let channels = dst_fmt.channels();

        let mut map = dst.map()?;
        let dst_slice = map.as_mut_slice();

        self.render_box(
            dst_w, dst_h, dst_rs, channels, dst_slice, detect, color_mode,
        )?;

        if detect.is_empty() {
            return Ok(());
        }
        let proto_shape = proto_data.protos.shape();
        if proto_shape.len() != 3 {
            return Err(Error::InvalidShape(format!(
                "protos tensor must be rank-3, got {proto_shape:?}"
            )));
        }
        let proto_h = proto_shape[0];
        let proto_w = proto_shape[1];
        let num_protos = proto_shape[2];
        let coeff_shape = proto_data.mask_coefficients.shape();
        if coeff_shape.len() != 2 {
            return Err(Error::InvalidShape(format!(
                "mask_coefficients tensor must be rank-2, got {coeff_shape:?}"
            )));
        }
        // Genuine "no detections this frame" → nothing to render.
        if coeff_shape[0] == 0 {
            return Ok(());
        }
        if coeff_shape[1] != num_protos {
            return Err(Error::InvalidShape(format!(
                "mask_coefficients second dimension must match num_protos \
                 ({num_protos}), got {coeff_shape:?}"
            )));
        }

        // Widen coefficients to f32 once; shape [N, num_protos].
        let coeff_f32: Vec<f32> = match proto_data.mask_coefficients.dtype() {
            DType::F32 => {
                let t = proto_data.mask_coefficients.as_f32().expect("F32");
                let m = t.map()?;
                m.as_slice().to_vec()
            }
            DType::F16 => {
                let t = proto_data.mask_coefficients.as_f16().expect("F16");
                let m = t.map()?;
                m.as_slice().iter().map(|v| v.to_f32()).collect()
            }
            DType::I8 => {
                let t = proto_data.mask_coefficients.as_i8().expect("I8");
                let m = t.map()?;
                if let Some(q) = t.quantization() {
                    use edgefirst_tensor::QuantMode;
                    let (scale, zp) = match q.mode() {
                        QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
                        QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
                        other => {
                            return Err(Error::NotSupported(format!(
                                "I8 mask_coefficients quantization mode {other:?} not supported"
                            )));
                        }
                    };
                    m.as_slice()
                        .iter()
                        .map(|&v| (v as f32 - zp) * scale)
                        .collect()
                } else {
                    m.as_slice().iter().map(|&v| v as f32).collect()
                }
            }
            DType::I16 => {
                let t = proto_data.mask_coefficients.as_i16().expect("I16");
                let m = t.map()?;
                if let Some(q) = t.quantization() {
                    use edgefirst_tensor::QuantMode;
                    let (scale, zp) = match q.mode() {
                        QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
                        QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
                        other => {
                            return Err(Error::NotSupported(format!(
                                "I16 mask_coefficients quantization mode {other:?} not supported"
                            )));
                        }
                    };
                    m.as_slice()
                        .iter()
                        .map(|&v| (v as f32 - zp) * scale)
                        .collect()
                } else {
                    m.as_slice().iter().map(|&v| v as f32).collect()
                }
            }
            other => {
                return Err(Error::InvalidShape(format!(
                    "mask_coefficients dtype {other:?} not supported"
                )));
            }
        };

        // Precompute letterbox scale/offset for output-pixel → proto-pixel mapping.
        let (lx0, lx_range, ly0, ly_range) = match letterbox {
            Some([lx0, ly0, lx1, ly1]) => (lx0, lx1 - lx0, ly0, ly1 - ly0),
            None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
        };

        // Per-dtype dispatch. Map protos once, call the inner draw loop
        // with a dtype-specialized loader closure.
        match proto_data.protos.dtype() {
            DType::F32 => {
                let t = proto_data.protos.as_f32().expect("F32");
                let m = t.map()?;
                self.draw_proto_masks_inner(
                    dst_slice,
                    dst_w,
                    dst_h,
                    dst_rs,
                    channels,
                    detect,
                    m.as_slice(),
                    &coeff_f32,
                    proto_h,
                    proto_w,
                    num_protos,
                    opacity,
                    (lx0, lx_range, ly0, ly_range),
                    color_mode,
                    0.0_f32,
                    |p: &f32, _| *p,
                );
            }
            DType::F16 => {
                let t = proto_data.protos.as_f16().expect("F16");
                let m = t.map()?;
                self.draw_proto_masks_inner(
                    dst_slice,
                    dst_w,
                    dst_h,
                    dst_rs,
                    channels,
                    detect,
                    m.as_slice(),
                    &coeff_f32,
                    proto_h,
                    proto_w,
                    num_protos,
                    opacity,
                    (lx0, lx_range, ly0, ly_range),
                    color_mode,
                    0.0_f32,
                    |p: &half::f16, _| p.to_f32(),
                );
            }
            DType::I8 => {
                use edgefirst_tensor::QuantMode;
                let t = proto_data.protos.as_i8().expect("I8");
                let m = t.map()?;
                let quant = t.quantization().ok_or_else(|| {
                    Error::InvalidShape("I8 protos require quantization metadata".into())
                })?;
                let (scale, zp) = match quant.mode() {
                    QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
                    QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
                    QuantMode::PerChannel { axis, .. }
                    | QuantMode::PerChannelSymmetric { axis, .. } => {
                        return Err(Error::NotSupported(format!(
                            "per-channel quantization (axis={axis}) in draw_proto_masks \
                             CPU path not yet supported"
                        )));
                    }
                };
                self.draw_proto_masks_inner(
                    dst_slice,
                    dst_w,
                    dst_h,
                    dst_rs,
                    channels,
                    detect,
                    m.as_slice(),
                    &coeff_f32,
                    proto_h,
                    proto_w,
                    num_protos,
                    opacity,
                    (lx0, lx_range, ly0, ly_range),
                    color_mode,
                    scale,
                    move |p: &i8, _| (*p as f32) - zp,
                );
            }
            other => {
                return Err(Error::InvalidShape(format!(
                    "proto tensor dtype {other:?} not supported"
                )));
            }
        }

        Ok(())
    }

    #[allow(clippy::too_many_arguments)]
    fn draw_proto_masks_inner<P: Copy>(
        &self,
        dst_slice: &mut [u8],
        dst_w: usize,
        dst_h: usize,
        dst_rs: usize,
        channels: usize,
        detect: &[DetectBox],
        protos: &[P],
        coeff_all_f32: &[f32],
        proto_h: usize,
        proto_w: usize,
        num_protos: usize,
        opacity: f32,
        letterbox_xy: (f32, f32, f32, f32),
        color_mode: crate::ColorMode,
        acc_scale: f32,
        load_f32: impl Fn(&P, f32) -> f32 + Copy,
    ) {
        let (lx0, lx_range, ly0, ly_range) = letterbox_xy;
        let stride_y = proto_w * num_protos;
        for (idx, det) in detect.iter().enumerate() {
            let coeff = &coeff_all_f32[idx * num_protos..(idx + 1) * num_protos];
            let color_index = color_mode.index(idx, det.label);
            let color = self.colors[color_index % self.colors.len()];
            let alpha = if opacity == 1.0 {
                color[3] as u16
            } else {
                (color[3] as f32 * opacity).round() as u16
            };

            let start_x = (dst_w as f32 * det.bbox.xmin).round() as usize;
            let start_y = (dst_h as f32 * det.bbox.ymin).round() as usize;
            let end_x = ((dst_w as f32 * det.bbox.xmax).round() as usize).min(dst_w);
            let end_y = ((dst_h as f32 * det.bbox.ymax).round() as usize).min(dst_h);

            for y in start_y..end_y {
                for x in start_x..end_x {
                    let px = (lx0 + (x as f32 / dst_w as f32) * lx_range) * proto_w as f32 - 0.5;
                    let py = (ly0 + (y as f32 / dst_h as f32) * ly_range) * proto_h as f32 - 0.5;

                    // Bilinear interpolation with per-load widening. Inline
                    // bilinear-sample since bilinear_dot_slice takes a
                    // different closure shape (no `zp` arg).
                    let x0 = (px.floor() as isize).clamp(0, proto_w as isize - 1) as usize;
                    let y0 = (py.floor() as isize).clamp(0, proto_h as isize - 1) as usize;
                    let x1 = (x0 + 1).min(proto_w - 1);
                    let y1 = (y0 + 1).min(proto_h - 1);
                    let fx = px - px.floor();
                    let fy = py - py.floor();
                    let w00 = (1.0 - fx) * (1.0 - fy);
                    let w10 = fx * (1.0 - fy);
                    let w01 = (1.0 - fx) * fy;
                    let w11 = fx * fy;
                    let b00 = y0 * stride_y + x0 * num_protos;
                    let b10 = y0 * stride_y + x1 * num_protos;
                    let b01 = y1 * stride_y + x0 * num_protos;
                    let b11 = y1 * stride_y + x1 * num_protos;
                    let mut acc = 0.0_f32;
                    for p in 0..num_protos {
                        let v00 = load_f32(&protos[b00 + p], 0.0);
                        let v10 = load_f32(&protos[b10 + p], 0.0);
                        let v01 = load_f32(&protos[b01 + p], 0.0);
                        let v11 = load_f32(&protos[b11 + p], 0.0);
                        let val = w00 * v00 + w10 * v10 + w01 * v01 + w11 * v11;
                        acc += coeff[p] * val;
                    }
                    let final_acc = if acc_scale == 0.0 {
                        acc
                    } else {
                        acc_scale * acc
                    };
                    // Pass-through: acc_scale=0.0 means "no scaling" (f32/f16
                    // native); non-zero means "apply scale once" (i8 with
                    // per-tensor quant).
                    let mask = 1.0 / (1.0 + (-final_acc).exp());
                    if mask < 0.5 {
                        continue;
                    }
                    let dst_index = y * dst_rs + x * channels;
                    for c in 0..3 {
                        dst_slice[dst_index + c] = ((color[c] as u16 * alpha
                            + dst_slice[dst_index + c] as u16 * (255 - alpha))
                            / 255) as u8;
                    }
                }
            }
        }
    }
}