Skip to main content

edgefirst_image/cpu/
mod.rs

1// SPDX-FileCopyrightText: Copyright 2025 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::{
5    Crop, Error, Flip, FunctionTimer, ImageProcessorTrait, Rect, ResolvedCrop, Result, Rotation,
6};
7use edgefirst_tensor::{
8    DType, PixelFormat, Tensor, TensorDyn, TensorMapTrait, TensorMemory, TensorTrait,
9};
10use edgefirst_tensor::{DetectBox, ProtoData, Segmentation};
11
12mod convert;
13mod masks;
14mod resize;
15mod simd;
16mod tests;
17
18// bilinear_dot removed — masks.rs now uses slice-native bilinear_dot_slice
19// closure-based kernel, invoked through the local dtype dispatch below.
20
21/// Resolved colorimetry parameters for a format conversion, computed once at
22/// the dispatch site (`convert_impl`) from the appropriate tensor:
23/// - YUV→RGB conversions use the **source** tensor's colorimetry.
24/// - RGB→YUV conversions use the **destination** tensor's colorimetry.
25///
26/// `matrix`/`range` feed the `yuv` crate kernels; `src_full_range`/
27/// `dst_full_range` gate the hand-rolled luma (limited↔full) expansion in the
28/// grey/luma copy helpers.
29#[derive(Debug, Clone, Copy)]
30pub(crate) struct ColorParams {
31    /// Matrix resolved from the YUV side (src for decode, dst for encode).
32    pub matrix: yuv::YuvStandardMatrix,
33    /// Range resolved from the YUV side (src for decode, dst for encode).
34    pub range: yuv::YuvRange,
35    /// The same matrix/range in HAL's own terms, so the hand-rolled fixed-point
36    /// encoders draw their `(kr, kb)` weights and luma/chroma swings from the
37    /// canonical [`edgefirst_tensor::colorimetry`] source rather than a private
38    /// duplicate table.
39    pub encoding: edgefirst_tensor::ColorEncoding,
40    pub range_kind: edgefirst_tensor::ColorRange,
41    /// True when the **source** tensor's resolved range is full (decode-side
42    /// luma extraction: copy luma directly instead of limited→full expansion).
43    pub src_full_range: bool,
44    /// True when the **destination** tensor's resolved range is full
45    /// (encode-side luma copies).
46    pub dst_full_range: bool,
47}
48
49/// CPUConverter implements the ImageProcessor trait using the fallback CPU
50/// implementation for image processing.
51#[derive(Debug)]
52pub struct CPUProcessor {
53    resizer: fast_image_resize::Resizer,
54    options: fast_image_resize::ResizeOptions,
55    colors: [[u8; 4]; 20],
56    /// Reusable scratch tensor for the U8→float widen path.
57    ///
58    /// Holds a U8 image in the destination's pixel format and dimensions.
59    /// Reallocated only when the format or dimensions change, amortising
60    /// the heap allocation across repeated same-size conversions.
61    widen_scratch: Option<TensorDyn>,
62    /// Reusable scratch for de-striding a padded source before resize.
63    ///
64    /// `fast_image_resize` needs a tightly-packed input; a padded source
65    /// (64-aligned stride from codec decode / `create_image`) is copied here
66    /// row-by-row first. Kept on the processor so the steady-state resize loop
67    /// reuses the allocation instead of a fresh `Vec` per call (the stride
68    /// alignment in this release makes that de-stride copy fire far more often).
69    resize_destride_scratch: Vec<u8>,
70    /// Reusable scratch for de-striding a padded resize *destination* before
71    /// resize.
72    ///
73    /// `fast_image_resize` needs a tightly-packed output buffer; a padded
74    /// destination (a DMA pitch-aligned tensor, or a `view()` narrower than
75    /// its parent's row stride) is resized into this tight scratch first,
76    /// then copied back into the real destination row-by-row at its true
77    /// stride. Kept on the processor so the steady-state resize loop reuses
78    /// the allocation instead of a fresh `Vec` per call — mirrors
79    /// `resize_destride_scratch` on the source side.
80    resize_dst_destride_scratch: Vec<u8>,
81    /// Reusable cache-resident scratch for the strip-fused NV→planar path
82    /// (`convert_nv_to_planar_fused`): holds a single row strip of packed RGB
83    /// (`STRIP_ROWS * width * 3` bytes). Keeping it on the processor avoids a
84    /// per-frame allocation and lets each strip stay hot in L2 between the YUV
85    /// decode and the deinterleave. Grown on demand; never shrunk.
86    nv_strip_scratch: Vec<u8>,
87    /// Reusable cache-resident scratch for the fused NV→planar path's
88    /// **region** case: a tight (stride == width) copy of the current
89    /// strip's luma rows, starting at column 0. The `yuv` crate's row
90    /// iteration processes fixed-`stride`-sized chunks internally, using
91    /// only the first `width` bytes of each — feeding it a slice that starts
92    /// at a mid-row column offset (a source-crop `region` with a nonzero
93    /// `left`) reads past the *source* buffer's true end for the region's
94    /// last row when that row is also the source's last row (the unused
95    /// chunk tail has no next row to alias into there). Packing each row
96    /// into a `stride == width` buffer first sidesteps that read. Unused
97    /// (and zero-cost) for any row-aligned read (`region.left == 0`,
98    /// including the whole-frame `region: None` case): with no column
99    /// shift, every row chunk the crate sees is a real, fully-owned source
100    /// row, so those reads stay fully zero-copy. Grown on demand; never
101    /// shrunk.
102    nv_strip_y_pack: Vec<u8>,
103    /// Same as [`Self::nv_strip_y_pack`], for the chroma (UV) rows.
104    nv_strip_uv_pack: Vec<u8>,
105    /// Test-only counter of `convert` calls that took the fused NV→planar
106    /// strip path (`convert_nv_to_planar_fused`), incremented at the gate in
107    /// `convert_u8`. Exists only under `#[cfg(test)]` — a field that's
108    /// written but (outside tests) never read trips clippy's dead-code lint
109    /// on both OS lanes. Read via [`Self::fused_hits`] so path-selection
110    /// tests assert on the actual gate decision rather than on timing.
111    #[cfg(test)]
112    fused_hits: u64,
113    /// Test-only record of the dimensions of the pre-resize intermediate the
114    /// most recent `convert_u8` call allocated, or `None` when that call did
115    /// not need one. Lets the allocation-proportionality tests assert that a
116    /// cropped convert sizes its intermediate to the crop rather than to the
117    /// whole frame, without reaching into the allocator. Reset at the top of
118    /// every `convert_u8`.
119    #[cfg(test)]
120    last_tmp_dims: Option<(usize, usize)>,
121    /// Reusable intermediate buffers for the multi-step convert pipeline
122    /// (pre-resize format-convert and the resized-RGB scratch). Reused across
123    /// frames when the dimensions/format match so the steady-state
124    /// letterbox/resize loop does not reallocate — and `alloc_zeroed`-clear — a
125    /// full-frame buffer per call. The region each consumer reads is always
126    /// fully overwritten first, so reused (non-zeroed) contents are never read.
127    convert_tmp: Option<Tensor<u8>>,
128    convert_tmp2: Option<Tensor<u8>>,
129    /// Reusable crop-sized copy of the source's own pixel format, holding the
130    /// halo-grown crop rectangle extracted by [`Self::extract_nv_region`] for
131    /// the crop-sized pre-resize intermediate. Only allocated for a cropped
132    /// convert that takes that path; the whole-frame path never touches it.
133    convert_src_sub: Option<Tensor<u8>>,
134}
135
136// `CPUProcessor` was `#[derive(Clone)]` before the `widen_scratch` field was
137// added; `TensorDyn` is not `Clone`, so the derive no longer applies. Restore
138// the public `Clone` impl by hand to avoid a breaking API change. The scratch
139// cache is a private allocation amortiser, not part of the logical value, so a
140// clone starts empty rather than sharing or duplicating it.
141impl Clone for CPUProcessor {
142    fn clone(&self) -> Self {
143        Self {
144            resizer: self.resizer.clone(),
145            options: self.options,
146            colors: self.colors,
147            widen_scratch: None,
148            resize_destride_scratch: Vec::new(),
149            resize_dst_destride_scratch: Vec::new(),
150            nv_strip_scratch: Vec::new(),
151            nv_strip_y_pack: Vec::new(),
152            nv_strip_uv_pack: Vec::new(),
153            #[cfg(test)]
154            fused_hits: 0,
155            #[cfg(test)]
156            last_tmp_dims: None,
157            convert_tmp: None,
158            convert_tmp2: None,
159            convert_src_sub: None,
160        }
161    }
162}
163
164unsafe impl Send for CPUProcessor {}
165unsafe impl Sync for CPUProcessor {}
166
167impl Default for CPUProcessor {
168    fn default() -> Self {
169        Self::new_bilinear()
170    }
171}
172
173/// Write the base layer of `dst` before mask rendering.
174///
175/// This is the terminal fallback: on CPU we have no 2D hardware, so a
176/// direct buffer write is the appropriate primitive. The invariant is that
177/// every call to the CPU draw_* entry points fully initialises dst — we
178/// never rely on "whatever was in the buffer" from the caller.
179///
180/// - `background == Some(bg)` → byte-for-byte copy bg → dst (after shape /
181///   format validation).
182/// - `background == None` → fill dst with 0x00 (transparent black).
183fn prepare_dst_base_cpu(dst: &mut TensorDyn, background: Option<&TensorDyn>) -> Result<()> {
184    match background {
185        Some(bg) => {
186            if bg.shape() != dst.shape() {
187                return Err(Error::InvalidShape(
188                    "background shape does not match dst".into(),
189                ));
190            }
191            if bg.format() != dst.format() {
192                return Err(Error::InvalidShape(
193                    "background pixel format does not match dst".into(),
194                ));
195            }
196            let bg_u8 = bg.as_typed::<u8>().ok_or(Error::NotAnImage)?;
197            let dst_u8 = dst.as_typed_mut::<u8>().ok_or(Error::NotAnImage)?;
198            // Row-wise, not a flat copy_from_slice: `bg` and `dst` can each
199            // carry their own row_stride (a padded D3D11/GPU-pitch destination
200            // composited from a tight background, or vice versa), so a single
201            // equal-length check is wrong even when the logical shapes match.
202            // `CPUProcessor::copy_image` already walks both sides at their
203            // own pitch via `logical_surface`/`packed_row_pairs`.
204            CPUProcessor::copy_image(bg_u8, dst_u8)?;
205        }
206        None => {
207            let dst_u8 = dst.as_typed_mut::<u8>().ok_or(Error::NotAnImage)?;
208            let mut dst_map = dst_u8.map_mut()?;
209            dst_map.as_mut_slice().fill(0);
210        }
211    }
212    Ok(())
213}
214
215/// Whether a source-crop origin is safe for the fused NV→planar strip path
216/// (`convert_nv_to_planar_fused`'s `region` parameter). NV12 (4:2:0)
217/// subsamples chroma both directions, so an odd `top` would desync the
218/// luma/chroma row pairing; NV16 (4:2:2) subsamples only horizontally, so
219/// only `left` must be even; NV24 (4:4:4) carries full-resolution chroma and
220/// has no alignment constraint. Misaligned origins fall through to the
221/// general path rather than snapping or shifting.
222fn chroma_alignment_ok(fmt: PixelFormat, r: Rect) -> bool {
223    match fmt {
224        PixelFormat::Nv12 => r.left.is_multiple_of(2) && r.top.is_multiple_of(2),
225        PixelFormat::Nv16 => r.left.is_multiple_of(2),
226        PixelFormat::Nv24 => true,
227        _ => false,
228    }
229}
230
231/// Compute row stride for a packed-format Tensor<u8> image given its format.
232fn row_stride_for(width: usize, fmt: PixelFormat) -> usize {
233    use edgefirst_tensor::PixelLayout;
234    match fmt.layout() {
235        PixelLayout::Packed => width * fmt.channels(),
236        PixelLayout::Planar | PixelLayout::SemiPlanar => width,
237        _ => width, // fallback for non-exhaustive
238    }
239}
240
241/// Read the effective row stride from a tensor, falling back to the computed
242/// minimum stride if the tensor has no explicit stride set. This correctly
243/// handles tensors with GPU pitch-alignment padding (e.g., from
244/// `ImageProcessor::create_image()` or codec strided decode).
245fn tensor_row_stride(tensor: &Tensor<u8>) -> usize {
246    tensor.effective_row_stride().unwrap_or_else(|| {
247        let w = tensor.width().unwrap_or(0);
248        let fmt = tensor.format().unwrap_or(PixelFormat::Rgb);
249        row_stride_for(w, fmt)
250    })
251}
252
253/// The `(rows, row_bytes)` of an image tensor's **logical** pixel surface: how
254/// many rows a converter may write, and how many bytes of each row are pixels.
255///
256/// Everything between `row_bytes` and [`tensor_row_stride`] is off-limits. For
257/// an allocation-padded tensor those bytes are dead padding; for a
258/// [`Tensor::view`](edgefirst_tensor::Tensor::view) destination they are the
259/// **parent image's neighbouring columns**, so a writer that treats the mapped
260/// slice as one flat buffer both mis-places its own rows and corrupts pixels
261/// that are not its to write.
262fn logical_surface(tensor: &Tensor<u8>) -> Result<(usize, usize)> {
263    use edgefirst_tensor::PixelLayout;
264    let fmt = tensor.format().ok_or(Error::NotAnImage)?;
265    let w = tensor.width().ok_or(Error::NotAnImage)?;
266    let h = tensor.height().ok_or(Error::NotAnImage)?;
267    Ok(match fmt.layout() {
268        PixelLayout::Packed => (h, w * fmt.channels()),
269        PixelLayout::Planar => (fmt.channels() * h, w),
270        PixelLayout::SemiPlanar => (fmt.combined_plane_height(h).unwrap_or(h), w),
271        // `PixelLayout` is non-exhaustive; treat any future layout as one row
272        // per shape row at the tensor's own pitch (see `row_stride_for`).
273        _ => (h, row_stride_for(w, fmt)),
274    })
275}
276
277/// Iterate `(src_row, dst_row)` pairs over a pair of image surfaces, each row
278/// clipped to its own logical width so neither side's stride padding — nor a
279/// view destination's neighbouring parent pixels — is read or written.
280///
281/// Callers must have validated both surfaces with [`guard_plane`] first.
282fn packed_row_pairs<'s, 'd>(
283    src: &'s [u8],
284    src_stride: usize,
285    src_row_bytes: usize,
286    dst: &'d mut [u8],
287    dst_stride: usize,
288    dst_row_bytes: usize,
289    rows: usize,
290) -> impl Iterator<Item = (&'s [u8], &'d mut [u8])> {
291    src.chunks(src_stride)
292        .zip(dst.chunks_mut(dst_stride))
293        .take(rows)
294        .map(move |(s, d)| (&s[..src_row_bytes], &mut d[..dst_row_bytes]))
295}
296
297/// Split a contiguous semi-planar (NV12/NV16/NV24) tensor's mapped bytes into
298/// `(luma, chroma)` planes at the `stride * src_h` boundary, validating the map
299/// holds the full combined plane first. A bare `split_at` would panic if an
300/// imported tensor's caller-supplied dimensions/stride exceed its actual buffer
301/// (untrusted input), so this returns `Error::InvalidShape` instead.
302fn split_semi_planar(
303    bytes: &[u8],
304    stride: usize,
305    src_h: usize,
306    fmt: PixelFormat,
307) -> Result<(&[u8], &[u8])> {
308    let total_h = fmt.combined_plane_height(src_h).unwrap_or(src_h);
309    let need = stride.checked_mul(total_h).ok_or_else(|| {
310        Error::InvalidShape(format!(
311            "{fmt:?} plane size overflow (stride={stride}, h={src_h})"
312        ))
313    })?;
314    if bytes.len() < need {
315        return Err(Error::InvalidShape(format!(
316            "{fmt:?} source has {} bytes but needs {need} (stride={stride}, h={src_h})",
317            bytes.len()
318        )));
319    }
320    Ok(bytes.split_at(stride * src_h))
321}
322
323/// Mutable mirror of [`split_semi_planar`]: split a contiguous semi-planar
324/// destination's mapped bytes into `(luma, chroma)` planes at the
325/// `stride * dst_h` boundary, validating the map holds the full combined plane
326/// first. A bare `split_at_mut` would panic if a caller-supplied (untrusted)
327/// destination's declared dimensions/stride exceed its actual buffer, so this
328/// returns `Error::InvalidShape` instead.
329fn split_semi_planar_mut(
330    bytes: &mut [u8],
331    stride: usize,
332    dst_h: usize,
333    fmt: PixelFormat,
334) -> Result<(&mut [u8], &mut [u8])> {
335    let total_h = fmt.combined_plane_height(dst_h).unwrap_or(dst_h);
336    let need = stride.checked_mul(total_h).ok_or_else(|| {
337        Error::InvalidShape(format!(
338            "{fmt:?} plane size overflow (stride={stride}, combined_h={total_h})"
339        ))
340    })?;
341    if bytes.len() < need {
342        return Err(Error::InvalidShape(format!(
343            "{fmt:?} destination has {} bytes but needs {need} (stride={stride}, combined_h={total_h})",
344            bytes.len()
345        )));
346    }
347    Ok(bytes.split_at_mut(stride * dst_h))
348}
349
350/// Validate that a mapped plane buffer of `buf_len` bytes can hold `rows` rows
351/// of `stride` bytes each, with `row_bytes` valid bytes per row. Returns
352/// `Error::InvalidShape` (instead of letting a later row slice panic) when a
353/// caller-supplied stride/shape exceeds the actual allocation. `what` labels the
354/// buffer in the error message.
355fn guard_plane(
356    buf_len: usize,
357    stride: usize,
358    rows: usize,
359    row_bytes: usize,
360    what: &str,
361) -> Result<()> {
362    let need = stride.checked_mul(rows).ok_or_else(|| {
363        Error::InvalidShape(format!(
364            "{what} plane size overflow (stride={stride}, rows={rows})"
365        ))
366    })?;
367    if row_bytes > stride || buf_len < need {
368        return Err(Error::InvalidShape(format!(
369            "{what} buffer too small: {buf_len} bytes, need {need} (stride={stride}, rows={rows}, row_bytes={row_bytes})"
370        )));
371    }
372    Ok(())
373}
374
375/// Apply XOR 0x80 bias to color channels only, preserving alpha.
376///
377/// Matches GL int8 shader behavior: `vec4(int8_bias(c.rgb), c.a)`.
378/// For formats without alpha, XORs every byte (fast path).
379pub(crate) fn apply_int8_xor_bias(data: &mut [u8], fmt: PixelFormat) {
380    use edgefirst_tensor::PixelLayout;
381    if !fmt.has_alpha() {
382        for b in data.iter_mut() {
383            *b ^= 0x80;
384        }
385    } else if fmt.layout() == PixelLayout::Planar {
386        // Planar with alpha (e.g. PlanarRgba): XOR color planes, skip alpha plane.
387        let channels = fmt.channels();
388        let plane_size = data.len() / channels;
389        for b in data[..plane_size * (channels - 1)].iter_mut() {
390            *b ^= 0x80;
391        }
392    } else {
393        // Packed with alpha (Rgba, Bgra): XOR color bytes, skip alpha byte.
394        let channels = fmt.channels();
395        for pixel in data.chunks_exact_mut(channels) {
396            for b in &mut pixel[..channels - 1] {
397                *b ^= 0x80;
398            }
399        }
400    }
401}
402
403/// Row-confined [`apply_int8_xor_bias`] for a destination whose row pitch may
404/// exceed its pixel bytes — an allocation-padded tensor, or a `Tensor::view()`
405/// whose pitch is the parent's and whose trailing bytes are the parent's
406/// neighbouring pixels. The `I8` post-pass of the CPU convert and of the G2D
407/// blit alike.
408pub(crate) fn apply_int8_xor_bias_rows(tensor: &mut Tensor<u8>, fmt: PixelFormat) -> Result<()> {
409    use edgefirst_tensor::PixelLayout;
410    let (rows, row_bytes) = logical_surface(tensor)?;
411    let stride = tensor_row_stride(tensor);
412    // A planar alpha channel is a whole trailing plane, so the rows to bias are
413    // the leading `channels - 1` planes rather than a prefix of each row.
414    let color_rows = if fmt.has_alpha() && fmt.layout() == PixelLayout::Planar {
415        rows / fmt.channels() * (fmt.channels() - 1)
416    } else {
417        rows
418    };
419    // Within a colour row the packed-alpha skip still applies; a planar row is
420    // all colour, so bias it as a format without alpha.
421    let row_fmt = if fmt.layout() == PixelLayout::Planar {
422        PixelFormat::Grey
423    } else {
424        fmt
425    };
426    let mut map = tensor.map_mut()?;
427    let buf = map.as_mut_slice();
428    guard_plane(buf.len(), stride, rows, row_bytes, "int8 bias dst")?;
429    for row in buf.chunks_mut(stride).take(color_rows) {
430        apply_int8_xor_bias(&mut row[..row_bytes], row_fmt);
431    }
432    Ok(())
433}
434
435impl CPUProcessor {
436    /// Creates a new CPUConverter with bilinear resizing.
437    pub fn new() -> Self {
438        Self::new_bilinear()
439    }
440
441    /// Creates a new CPUConverter with bilinear resizing.
442    fn new_bilinear() -> Self {
443        let resizer = fast_image_resize::Resizer::new();
444        let options = fast_image_resize::ResizeOptions::new()
445            .resize_alg(fast_image_resize::ResizeAlg::Convolution(
446                fast_image_resize::FilterType::Bilinear,
447            ))
448            .use_alpha(false);
449
450        log::debug!("CPUConverter created");
451        Self {
452            resizer,
453            options,
454            colors: crate::DEFAULT_COLORS_U8,
455            widen_scratch: None,
456            resize_destride_scratch: Vec::new(),
457            resize_dst_destride_scratch: Vec::new(),
458            nv_strip_scratch: Vec::new(),
459            nv_strip_y_pack: Vec::new(),
460            nv_strip_uv_pack: Vec::new(),
461            #[cfg(test)]
462            fused_hits: 0,
463            #[cfg(test)]
464            last_tmp_dims: None,
465            convert_tmp: None,
466            convert_tmp2: None,
467            convert_src_sub: None,
468        }
469    }
470
471    /// Sets how resizing converts resample (see [`crate::Interpolation`]):
472    /// nearest-neighbour, or the bilinear filter of [`Self::new`].
473    pub fn set_interpolation(&mut self, mode: crate::Interpolation) {
474        let alg = match mode {
475            crate::Interpolation::Nearest => fast_image_resize::ResizeAlg::Nearest,
476            crate::Interpolation::Bilinear => {
477                fast_image_resize::ResizeAlg::Convolution(fast_image_resize::FilterType::Bilinear)
478            }
479        };
480        self.options = self.options.clone().resize_alg(alg);
481    }
482
483    /// Creates a new CPUConverter with nearest neighbor resizing.
484    pub fn new_nearest() -> Self {
485        let resizer = fast_image_resize::Resizer::new();
486        let options = fast_image_resize::ResizeOptions::new()
487            .resize_alg(fast_image_resize::ResizeAlg::Nearest)
488            .use_alpha(false);
489        log::debug!("CPUConverter created");
490        Self {
491            resizer,
492            options,
493            colors: crate::DEFAULT_COLORS_U8,
494            widen_scratch: None,
495            resize_destride_scratch: Vec::new(),
496            resize_dst_destride_scratch: Vec::new(),
497            nv_strip_scratch: Vec::new(),
498            nv_strip_y_pack: Vec::new(),
499            nv_strip_uv_pack: Vec::new(),
500            #[cfg(test)]
501            fused_hits: 0,
502            #[cfg(test)]
503            last_tmp_dims: None,
504            convert_tmp: None,
505            convert_tmp2: None,
506            convert_src_sub: None,
507        }
508    }
509
510    /// Test-only accessor for the `fused_hits` counter (see the field doc).
511    #[cfg(test)]
512    pub(super) fn fused_hits(&self) -> u64 {
513        self.fused_hits
514    }
515
516    /// Test-only accessor for the dimensions of the pre-resize intermediate the
517    /// last `convert_u8` allocated (see the `last_tmp_dims` field doc).
518    #[cfg(test)]
519    pub(super) fn last_tmp_dims(&self) -> Option<(usize, usize)> {
520        self.last_tmp_dims
521    }
522
523    pub(crate) fn support_conversion_pf(src: PixelFormat, dst: PixelFormat) -> bool {
524        use PixelFormat::*;
525        matches!(
526            (src, dst),
527            (Nv12, Rgb)
528                | (Nv12, Rgba)
529                | (Nv12, Grey)
530                | (Nv16, Rgb)
531                | (Nv16, Rgba)
532                | (Nv16, Bgra)
533                | (Nv24, Rgb)
534                | (Nv24, Rgba)
535                | (Nv24, Grey)
536                | (Nv24, Bgra)
537                | (Yuyv, Rgb)
538                | (Yuyv, Rgba)
539                | (Yuyv, Grey)
540                | (Yuyv, Yuyv)
541                | (Yuyv, PlanarRgb)
542                | (Yuyv, PlanarRgba)
543                | (Yuyv, Nv16)
544                | (Vyuy, Rgb)
545                | (Vyuy, Rgba)
546                | (Vyuy, Grey)
547                | (Vyuy, Vyuy)
548                | (Vyuy, PlanarRgb)
549                | (Vyuy, PlanarRgba)
550                | (Vyuy, Nv16)
551                | (Rgba, Rgb)
552                | (Rgba, Rgba)
553                | (Rgba, Grey)
554                | (Rgba, Yuyv)
555                | (Rgba, PlanarRgb)
556                | (Rgba, PlanarRgba)
557                | (Rgba, Nv16)
558                | (Rgb, Rgb)
559                | (Rgb, Rgba)
560                | (Rgb, Grey)
561                | (Rgb, Yuyv)
562                | (Rgb, PlanarRgb)
563                | (Rgb, PlanarRgba)
564                | (Rgb, Nv16)
565                | (Grey, Rgb)
566                | (Grey, Rgba)
567                | (Grey, Grey)
568                | (Grey, Yuyv)
569                | (Grey, PlanarRgb)
570                | (Grey, PlanarRgba)
571                | (Grey, Nv16)
572                | (Nv12, Bgra)
573                | (Yuyv, Bgra)
574                | (Vyuy, Bgra)
575                | (Rgba, Bgra)
576                | (Rgb, Bgra)
577                | (Grey, Bgra)
578                | (Bgra, Bgra)
579                | (PlanarRgb, Rgb)
580                | (PlanarRgb, Rgba)
581                | (PlanarRgba, Rgb)
582                | (PlanarRgba, Rgba)
583                | (PlanarRgb, Bgra)
584                | (PlanarRgba, Bgra)
585        )
586    }
587
588    /// Format conversion dispatch for Tensor<u8> with PixelFormat metadata.
589    pub(crate) fn convert_format_pf(
590        src: &Tensor<u8>,
591        dst: &mut Tensor<u8>,
592        src_fmt: PixelFormat,
593        dst_fmt: PixelFormat,
594        cp: ColorParams,
595    ) -> Result<()> {
596        let _timer = FunctionTimer::new(format!(
597            "ImageProcessor::convert_format {} to {}",
598            src_fmt, dst_fmt,
599        ));
600
601        use PixelFormat::*;
602        match (src_fmt, dst_fmt) {
603            (Nv12, Rgb) => Self::convert_nv12_to_rgb(src, dst, cp),
604            (Nv12, Rgba) => Self::convert_nv12_to_rgba(src, dst, cp),
605            (Nv12, Grey) => Self::convert_nv12_to_grey(src, dst, cp),
606            (Yuyv, Rgb) => Self::convert_yuyv_to_rgb(src, dst, cp),
607            (Yuyv, Rgba) => Self::convert_yuyv_to_rgba(src, dst, cp),
608            (Yuyv, Grey) => Self::convert_yuyv_to_grey(src, dst, cp),
609            (Yuyv, Yuyv) => Self::copy_image(src, dst),
610            (Yuyv, PlanarRgb) => Self::convert_yuyv_to_8bps(src, dst, cp),
611            (Yuyv, PlanarRgba) => Self::convert_yuyv_to_prgba(src, dst, cp),
612            (Yuyv, Nv16) => Self::convert_yuyv_to_nv16(src, dst),
613            (Vyuy, Rgb) => Self::convert_vyuy_to_rgb(src, dst, cp),
614            (Vyuy, Rgba) => Self::convert_vyuy_to_rgba(src, dst, cp),
615            (Vyuy, Grey) => Self::convert_vyuy_to_grey(src, dst, cp),
616            (Vyuy, Vyuy) => Self::copy_image(src, dst),
617            (Vyuy, PlanarRgb) => Self::convert_vyuy_to_8bps(src, dst, cp),
618            (Vyuy, PlanarRgba) => Self::convert_vyuy_to_prgba(src, dst, cp),
619            (Vyuy, Nv16) => Self::convert_vyuy_to_nv16(src, dst),
620            (Rgba, Rgb) => Self::convert_rgba_to_rgb(src, dst),
621            (Rgba, Rgba) => Self::copy_image(src, dst),
622            (Rgba, Grey) => Self::convert_rgba_to_grey(src, dst),
623            (Rgba, Yuyv) => Self::convert_rgba_to_yuyv(src, dst, cp),
624            (Rgba, PlanarRgb) => Self::convert_rgba_to_8bps(src, dst),
625            (Rgba, PlanarRgba) => Self::convert_rgba_to_prgba(src, dst),
626            (Rgba, Nv16) => Self::convert_rgba_to_nv16(src, dst, cp),
627            (Rgb, Rgb) => Self::copy_image(src, dst),
628            (Rgb, Rgba) => Self::convert_rgb_to_rgba(src, dst),
629            (Rgb, Grey) => Self::convert_rgb_to_grey(src, dst),
630            (Rgb, Yuyv) => Self::convert_rgb_to_yuyv(src, dst, cp),
631            (Rgb, PlanarRgb) => Self::convert_rgb_to_8bps(src, dst),
632            (Rgb, PlanarRgba) => Self::convert_rgb_to_prgba(src, dst),
633            (Rgb, Nv16) => Self::convert_rgb_to_nv16(src, dst, cp),
634            (Grey, Rgb) => Self::convert_grey_to_rgb(src, dst),
635            (Grey, Rgba) => Self::convert_grey_to_rgba(src, dst),
636            (Grey, Grey) => Self::copy_image(src, dst),
637            (Grey, Yuyv) => Self::convert_grey_to_yuyv(src, dst, cp),
638            (Grey, PlanarRgb) => Self::convert_grey_to_8bps(src, dst),
639            (Grey, PlanarRgba) => Self::convert_grey_to_prgba(src, dst),
640            (Grey, Nv16) => Self::convert_grey_to_nv16(src, dst, cp),
641
642            // the following converts are added for use in testing
643            (Nv16, Rgb) => Self::convert_nv16_to_rgb(src, dst, cp),
644            (Nv16, Rgba) => Self::convert_nv16_to_rgba(src, dst, cp),
645            (Nv24, Rgb) => Self::convert_nv24_to_rgb(src, dst, cp),
646            (Nv24, Rgba) => Self::convert_nv24_to_rgba(src, dst, cp),
647            (Nv24, Grey) => Self::convert_nv24_to_grey(src, dst, cp),
648            (PlanarRgb, Rgb) => Self::convert_8bps_to_rgb(src, dst),
649            (PlanarRgb, Rgba) => Self::convert_8bps_to_rgba(src, dst),
650            (PlanarRgba, Rgb) => Self::convert_prgba_to_rgb(src, dst),
651            (PlanarRgba, Rgba) => Self::convert_prgba_to_rgba(src, dst),
652
653            // BGRA destination: convert to RGBA layout, then swap R and B
654            (Bgra, Bgra) => Self::copy_image(src, dst),
655            (Nv12, Bgra) => {
656                Self::convert_nv12_to_rgba(src, dst, cp)?;
657                Self::swizzle_rb_4chan(dst)
658            }
659            (Nv16, Bgra) => {
660                Self::convert_nv16_to_rgba(src, dst, cp)?;
661                Self::swizzle_rb_4chan(dst)
662            }
663            (Nv24, Bgra) => {
664                Self::convert_nv24_to_rgba(src, dst, cp)?;
665                Self::swizzle_rb_4chan(dst)
666            }
667            (Yuyv, Bgra) => {
668                Self::convert_yuyv_to_rgba(src, dst, cp)?;
669                Self::swizzle_rb_4chan(dst)
670            }
671            (Vyuy, Bgra) => {
672                Self::convert_vyuy_to_rgba(src, dst, cp)?;
673                Self::swizzle_rb_4chan(dst)
674            }
675            (Rgba, Bgra) => {
676                Self::copy_image(src, dst)?;
677                Self::swizzle_rb_4chan(dst)
678            }
679            (Rgb, Bgra) => {
680                Self::convert_rgb_to_rgba(src, dst)?;
681                Self::swizzle_rb_4chan(dst)
682            }
683            (Grey, Bgra) => {
684                Self::convert_grey_to_rgba(src, dst)?;
685                Self::swizzle_rb_4chan(dst)
686            }
687            (PlanarRgb, Bgra) => {
688                Self::convert_8bps_to_rgba(src, dst)?;
689                Self::swizzle_rb_4chan(dst)
690            }
691            (PlanarRgba, Bgra) => {
692                Self::convert_prgba_to_rgba(src, dst)?;
693                Self::swizzle_rb_4chan(dst)
694            }
695
696            (s, d) => Err(Error::NotSupported(format!("Conversion from {s} to {d}",))),
697        }
698    }
699
700    /// Tensor<u8>-based fill_image_outside_crop.
701    pub(crate) fn fill_image_outside_crop_u8(
702        dst: &mut Tensor<u8>,
703        rgba: [u8; 4],
704        crop: Rect,
705    ) -> Result<()> {
706        let dst_fmt = dst.format().unwrap();
707        let dst_w = dst.width().unwrap();
708        let dst_h = dst.height().unwrap();
709        // Resolve the YUV fill encoding from the destination tensor so the
710        // border color matches the same matrix/range as a later YUV→RGB
711        // decode of this image. RGB/Grey fills ignore these params.
712        let cm = crate::colorimetry::resolve_colorimetry(dst.colorimetry(), dst.height());
713        let cp = ColorParams {
714            matrix: crate::colorimetry::yuv_matrix(cm.encoding.unwrap()),
715            range: crate::colorimetry::yuv_range(cm.range.unwrap()),
716            encoding: cm.encoding.unwrap(),
717            range_kind: cm.range.unwrap(),
718            src_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
719            dst_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
720        };
721        let dst_stride = tensor_row_stride(dst);
722        let mut dst_map = dst.map_mut()?;
723        let dst_tup = (dst_map.as_mut_slice(), dst_w, dst_h, dst_stride);
724        Self::fill_outside_crop_dispatch(dst_tup, dst_fmt, rgba, crop, cp)
725    }
726
727    /// Common fill dispatch by format. The tuple is
728    /// `(bytes, logical_width, logical_height, row_stride)` — the stride is the
729    /// destination's real byte pitch, which for a `view()` destination is the
730    /// parent image's, not `width × bpp`.
731    fn fill_outside_crop_dispatch(
732        dst: (&mut [u8], usize, usize, usize),
733        fmt: PixelFormat,
734        rgba: [u8; 4],
735        crop: Rect,
736        cp: ColorParams,
737    ) -> Result<()> {
738        use PixelFormat::*;
739        match fmt {
740            Rgba | Bgra => Self::fill_image_outside_crop_(dst, rgba, crop),
741            Rgb => Self::fill_image_outside_crop_(dst, Self::rgba_to_rgb(rgba), crop),
742            Grey => Self::fill_image_outside_crop_(dst, Self::rgba_to_grey(rgba), crop),
743            Yuyv => {
744                let (bytes, w, h, stride) = dst;
745                let yuyv = Self::rgba_to_yuyv(rgba, cp);
746                // Bulk fill in 4-byte [Y0,U,Y1,V] macropixels; crop bounds in
747                // macropixel units so a shared chroma pair is never split.
748                Self::fill_image_outside_crop_(
749                    (&mut *bytes, w / 2, h, stride),
750                    yuyv,
751                    Rect::new(crop.left / 2, crop.top, crop.width.div_ceil(2), crop.height),
752                )?;
753                // An odd width leaves a trailing unpaired pixel that the
754                // macropixel walk above (w / 2 units) structurally never
755                // reaches — the same trailing-pixel class the odd-width YUYV
756                // encoders handle. Its 2 bytes are [Y, U] (a row is w × 2
757                // bytes, so there is no room for its V). Fill it on every row
758                // where the last column lies outside the crop, judged in
759                // pixel units against the original crop.
760                if w % 2 == 1 {
761                    let last = w - 1;
762                    let col_outside = last < crop.left || last >= crop.left + crop.width;
763                    for y in 0..h {
764                        let row_outside = y < crop.top || y >= crop.top + crop.height;
765                        if row_outside || col_outside {
766                            let off = y * stride + last * 2;
767                            bytes[off] = yuyv[0];
768                            bytes[off + 1] = yuyv[1];
769                        }
770                    }
771                }
772                Ok(())
773            }
774            PlanarRgb => Self::fill_image_outside_crop_planar(dst, Self::rgba_to_rgb(rgba), crop),
775            PlanarRgba => Self::fill_image_outside_crop_planar(dst, rgba, crop),
776            Nv16 => {
777                let yuyv = Self::rgba_to_yuyv(rgba, cp);
778                Self::fill_image_outside_crop_yuv_semiplanar(dst, yuyv[0], [yuyv[1], yuyv[3]], crop)
779            }
780            _ => Err(Error::Internal(format!(
781                "Found unexpected destination {fmt}",
782            ))),
783        }
784    }
785}
786
787impl ImageProcessorTrait for CPUProcessor {
788    fn convert(
789        &mut self,
790        src: &TensorDyn,
791        dst: &mut TensorDyn,
792        rotation: Rotation,
793        flip: Flip,
794        crop: Crop,
795    ) -> Result<()> {
796        let crop = crop.resolve(
797            src.width().unwrap_or(0),
798            src.height().unwrap_or(0),
799            dst.width().unwrap_or(0),
800            dst.height().unwrap_or(0),
801        )?;
802        self.convert_impl(src, dst, rotation, flip, crop)
803    }
804
805    fn draw_decoded_masks(
806        &mut self,
807        dst: &mut TensorDyn,
808        detect: &[DetectBox],
809        segmentation: &[Segmentation],
810        overlay: crate::MaskOverlay<'_>,
811    ) -> Result<()> {
812        // CPU is the terminal fallback — it must always produce the full
813        // output, never assume the caller cleared dst. Every call writes
814        // the base layer first (bg copy or zero fill) and then the masks.
815        prepare_dst_base_cpu(dst, overlay.background)?;
816        let dst = dst.as_typed_mut::<u8>().ok_or(Error::NotAnImage)?;
817        self.draw_decoded_masks_impl(
818            dst,
819            detect,
820            segmentation,
821            overlay.opacity,
822            overlay.color_mode,
823        )
824    }
825
826    fn draw_proto_masks(
827        &mut self,
828        dst: &mut TensorDyn,
829        detect: &[DetectBox],
830        proto_data: &ProtoData,
831        overlay: crate::MaskOverlay<'_>,
832    ) -> Result<()> {
833        prepare_dst_base_cpu(dst, overlay.background)?;
834        let dst = dst.as_typed_mut::<u8>().ok_or(Error::NotAnImage)?;
835        self.draw_proto_masks_impl(
836            dst,
837            detect,
838            proto_data,
839            overlay.opacity,
840            overlay.letterbox,
841            overlay.color_mode,
842        )
843    }
844
845    fn set_class_colors(&mut self, colors: &[[u8; 4]]) -> Result<()> {
846        for (c, new_c) in self.colors.iter_mut().zip(colors.iter()) {
847            *c = *new_c;
848        }
849        Ok(())
850    }
851}
852
853// Internal methods — dtype-aware dispatch layer.
854impl CPUProcessor {
855    /// Top-level conversion dispatcher: handles dtype combinations.
856    pub(crate) fn convert_impl(
857        &mut self,
858        src: &TensorDyn,
859        dst: &mut TensorDyn,
860        rotation: Rotation,
861        flip: Flip,
862        crop: ResolvedCrop,
863    ) -> Result<()> {
864        let src_fmt = src.format().ok_or(Error::NotAnImage)?;
865        let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
866
867        // Resolve per-tensor colorimetry once, at the use site, without
868        // mutating either tensor. YUV→RGB conversions take their matrix/range
869        // from the source; RGB→YUV conversions from the destination. The grey/
870        // luma expansion is gated on the resolved range of the relevant side.
871        let src_cm = crate::colorimetry::effective_colorimetry(src);
872        let dst_cm = crate::colorimetry::effective_colorimetry(dst);
873        let src_full = src_cm.range == Some(edgefirst_tensor::ColorRange::Full);
874        let dst_full = dst_cm.range == Some(edgefirst_tensor::ColorRange::Full);
875        let src_params = ColorParams {
876            matrix: crate::colorimetry::yuv_matrix(src_cm.encoding.unwrap()),
877            range: crate::colorimetry::yuv_range(src_cm.range.unwrap()),
878            encoding: src_cm.encoding.unwrap(),
879            range_kind: src_cm.range.unwrap(),
880            src_full_range: src_full,
881            dst_full_range: dst_full,
882        };
883        let dst_params = ColorParams {
884            matrix: crate::colorimetry::yuv_matrix(dst_cm.encoding.unwrap()),
885            range: crate::colorimetry::yuv_range(dst_cm.range.unwrap()),
886            encoding: dst_cm.encoding.unwrap(),
887            range_kind: dst_cm.range.unwrap(),
888            src_full_range: src_full,
889            dst_full_range: dst_full,
890        };
891        match (src.dtype(), dst.dtype()) {
892            (DType::U8, DType::U8) => {
893                let src = src.as_typed::<u8>().unwrap();
894                let dst = dst.as_typed_mut::<u8>().unwrap();
895                self.convert_u8(
896                    src, dst, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
897                )
898            }
899            (DType::U8, DType::I8) => {
900                // Int8 output: reinterpret the i8 destination as u8 (layout-
901                // identical), convert directly into it, then XOR 0x80 in-place.
902                let src_u8 = src.as_typed::<u8>().unwrap();
903                let dst_i8 = dst.as_typed_mut::<i8>().unwrap();
904                // SAFETY: Tensor<i8> and Tensor<u8> are layout-identical
905                // (same element size, no T-dependent drop glue). Same
906                // rationale as gl::processor::tensor_i8_as_u8_mut.
907                let dst_u8 = unsafe { &mut *(dst_i8 as *mut Tensor<i8> as *mut Tensor<u8>) };
908                self.convert_u8(
909                    src_u8, dst_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
910                )?;
911                // Apply XOR 0x80 bias in-place (u8 → i8 conversion)
912                apply_int8_xor_bias_rows(dst_u8, dst_fmt)
913            }
914            (DType::U8, d @ (DType::F32 | DType::F16)) => {
915                let src_u8 = src.as_typed::<u8>().unwrap();
916                let dw = dst.width().ok_or(Error::NotAnImage)?;
917                let dh = dst.height().ok_or(Error::NotAnImage)?;
918                // Reuse the scratch tensor when format and dimensions match;
919                // otherwise reallocate and cache the new scratch.  Take the
920                // scratch out of `self` so that `convert_u8` can borrow `self`
921                // exclusively, then restore it afterwards.
922                let scratch_matches = self.widen_scratch.as_ref().is_some_and(|t| {
923                    t.width() == Some(dw) && t.height() == Some(dh) && t.format() == Some(dst_fmt)
924                });
925                let mut tmp = if scratch_matches {
926                    self.widen_scratch.take().unwrap()
927                } else {
928                    TensorDyn::image(
929                        dw,
930                        dh,
931                        dst_fmt,
932                        DType::U8,
933                        Some(TensorMemory::Mem),
934                        edgefirst_tensor::CpuAccess::ReadWrite,
935                    )?
936                };
937                {
938                    let tmp_u8 = tmp.as_typed_mut::<u8>().unwrap();
939                    self.convert_u8(
940                        src_u8, tmp_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params,
941                        dst_params,
942                    )?;
943                }
944                // Widen the u8 scratch into the float destination, then restore
945                // the scratch for reuse on the next call.
946                //
947                // The scratch is tight, but `dst` need not be: an allocation-
948                // padded float tensor, or a `Tensor::view()` whose pitch is the
949                // parent's, has a wider row than its pixels. Widen row by row so
950                // each output row lands at the destination's own pitch and the
951                // bytes past it stay untouched. On a tight destination the two
952                // pitches coincide and this is the previous flat widen.
953                {
954                    let tmp_u8 = tmp.as_typed::<u8>().unwrap();
955                    let (rows, row_len) = logical_surface(tmp_u8)?;
956                    let src_stride = tensor_row_stride(tmp_u8);
957                    let dst_stride_bytes = dst.effective_row_stride().ok_or(Error::NotAnImage)?;
958                    let src_map = tmp_u8.map_read()?;
959                    guard_plane(
960                        src_map.as_slice().len(),
961                        src_stride,
962                        rows,
963                        row_len,
964                        "widen src",
965                    )?;
966                    let src_rows = src_map.as_slice().chunks(src_stride).take(rows);
967                    let elem = d.size();
968                    if !dst_stride_bytes.is_multiple_of(elem) {
969                        return Err(Error::InvalidShape(format!(
970                            "{d} destination row stride {dst_stride_bytes} is not a multiple of \
971                             the element size {elem}"
972                        )));
973                    }
974                    let dst_stride = dst_stride_bytes / elem;
975                    match d {
976                        DType::F32 => {
977                            let dst_t = dst.as_typed_mut::<f32>().unwrap();
978                            let mut dst_map = dst_t.map_mut()?;
979                            // Counted in elements, not bytes — same invariant.
980                            guard_plane(
981                                dst_map.as_slice().len(),
982                                dst_stride,
983                                rows,
984                                row_len,
985                                "widen f32 dst",
986                            )?;
987                            // NEON-accelerated u8→f32 `/255` widen (bit-identical
988                            // to the scalar `b as f32 / 255.0`); the scalar
989                            // iterator form did not vectorise. See cpu::simd.
990                            for (s, dr) in
991                                src_rows.zip(dst_map.as_mut_slice().chunks_mut(dst_stride))
992                            {
993                                simd::widen_u8_to_f32_norm(s, &mut dr[..row_len]);
994                            }
995                        }
996                        DType::F16 => {
997                            let dst_t = dst.as_typed_mut::<half::f16>().unwrap();
998                            let mut dst_map = dst_t.map_mut()?;
999                            guard_plane(
1000                                dst_map.as_slice().len(),
1001                                dst_stride,
1002                                rows,
1003                                row_len,
1004                                "widen f16 dst",
1005                            )?;
1006                            // u8→f16 `/255` widen; uses native FP16
1007                            // (`ucvtf`+`fdiv`) at runtime on FEAT_FP16 CPUs
1008                            // (Orin), scalar `half::f16::from_f32` elsewhere.
1009                            // See cpu::simd.
1010                            for (s, dr) in
1011                                src_rows.zip(dst_map.as_mut_slice().chunks_mut(dst_stride))
1012                            {
1013                                simd::widen_u8_to_f16_norm(s, &mut dr[..row_len]);
1014                            }
1015                        }
1016                        _ => unreachable!(),
1017                    }
1018                }
1019                self.widen_scratch = Some(tmp);
1020                Ok(())
1021            }
1022            (s, d) => Err(Error::NotSupported(format!("dtype {s} -> {d}",))),
1023        }
1024    }
1025
1026    /// Reuse `cached` if it already has dimensions `(w, h)` and pixel format
1027    /// `fmt`, otherwise allocate a fresh `Mem` image. The returned buffer's
1028    /// contents are **not** zeroed on reuse — callers must fully overwrite the
1029    /// region they later read (see the note at the convert pipeline's tmp/tmp2
1030    /// site). This amortises the per-frame `Tensor::image` `alloc_zeroed`.
1031    fn reuse_or_alloc_image(
1032        cached: Option<Tensor<u8>>,
1033        w: usize,
1034        h: usize,
1035        fmt: PixelFormat,
1036    ) -> Result<Tensor<u8>> {
1037        if let Some(t) = cached {
1038            if t.width() == Some(w) && t.height() == Some(h) && t.format() == Some(fmt) {
1039                return Ok(t);
1040            }
1041        }
1042        Ok(Tensor::<u8>::image(
1043            w,
1044            h,
1045            fmt,
1046            Some(TensorMemory::Mem),
1047            edgefirst_tensor::CpuAccess::ReadWrite,
1048        )?)
1049    }
1050
1051    /// The source sub-rectangle to convert into the pre-resize intermediate for
1052    /// a cropped convert, or `None` to keep converting the whole frame.
1053    ///
1054    /// A cropped convert only ever reads the crop rect plus the resize filter's
1055    /// halo (see [`Self::filter_halo`]), so converting the whole frame into the
1056    /// intermediate wastes work proportional to the frame — for a 4K frame
1057    /// tiled into 640x640 crops, ~22x the pixels actually needed. Returning the
1058    /// halo-grown crop lets the pre-resize convert produce a crop-sized
1059    /// intermediate instead.
1060    ///
1061    /// The returned rect only ever *grows* the crop, and is clamped to the
1062    /// frame, so the resize reads exactly the same real source pixels — and
1063    /// clamps at exactly the same frame edges — as it did with a full-frame
1064    /// intermediate. Output is byte-identical; only the buffer size changes.
1065    ///
1066    /// `None` (unchanged full-frame behaviour) for: an uncropped convert, a
1067    /// crop covering the whole frame, a source format that cannot be extracted
1068    /// (only the semi-planar NV family can — `Tensor::view` rejects non-packed
1069    /// layouts, and NV is what the 4K tiling path feeds in), and a resize
1070    /// algorithm whose filter reach is not modelled.
1071    fn pre_resize_region(
1072        &self,
1073        src_fmt: PixelFormat,
1074        (src_w, src_h): (usize, usize),
1075        (dst_w, dst_h): (usize, usize),
1076        rotation: Rotation,
1077        crop: ResolvedCrop,
1078    ) -> Option<Rect> {
1079        use PixelFormat::{Nv12, Nv16, Nv24};
1080
1081        if !matches!(src_fmt, Nv12 | Nv16 | Nv24) {
1082            return None;
1083        }
1084        let r = crop.src_rect?;
1085        let full_src = Rect {
1086            left: 0,
1087            top: 0,
1088            width: src_w,
1089            height: src_h,
1090        };
1091        if r == full_src {
1092            return None;
1093        }
1094
1095        // The resize maps the source rect onto the destination rect; a quarter
1096        // turn swaps which destination extent each source axis is scaled onto
1097        // (see `adjust_dest_rect_for_rotate_flip_dims`), and the filter halo
1098        // depends on that scale factor.
1099        let d = crop.dst_rect.unwrap_or(Rect {
1100            left: 0,
1101            top: 0,
1102            width: dst_w,
1103            height: dst_h,
1104        });
1105        let (dst_x, dst_y) = match rotation {
1106            Rotation::None | Rotation::Rotate180 => (d.width, d.height),
1107            Rotation::Clockwise90 | Rotation::CounterClockwise90 => (d.height, d.width),
1108        };
1109        let halo_x = self.filter_halo(r.width, dst_x)?;
1110        let halo_y = self.filter_halo(r.height, dst_y)?;
1111
1112        let mut left = r.left.saturating_sub(halo_x);
1113        let mut top = r.top.saturating_sub(halo_y);
1114        let mut right = (r.left + r.width + halo_x).min(src_w);
1115        let mut bottom = (r.top + r.height + halo_y).min(src_h);
1116
1117        // NV12 subsamples chroma on both axes, NV16 on the horizontal one, so
1118        // the extracted origin must land on a chroma sample boundary for the
1119        // sub-image's chroma to line up with the frame's. Snap the origin DOWN
1120        // and the far edge UP — growing only, so every pixel inside `r` keeps
1121        // exactly the neighbours it had, whatever the crop's own parity.
1122        let (align_x, align_y) = match src_fmt {
1123            Nv12 => (2, 2),
1124            Nv16 => (2, 1),
1125            _ => (1, 1),
1126        };
1127        left -= left % align_x;
1128        top -= top % align_y;
1129        right = right.next_multiple_of(align_x).min(src_w);
1130        bottom = bottom.next_multiple_of(align_y).min(src_h);
1131
1132        let grown = Rect {
1133            left,
1134            top,
1135            width: right - left,
1136            height: bottom - top,
1137        };
1138        // The halo already reaches the whole frame — extracting would be a pure
1139        // copy with no saving.
1140        (grown != full_src).then_some(grown)
1141    }
1142
1143    /// U8-to-U8 conversion: the full format conversion + resize pipeline.
1144    #[allow(clippy::too_many_arguments)]
1145    fn convert_u8(
1146        &mut self,
1147        src: &Tensor<u8>,
1148        dst: &mut Tensor<u8>,
1149        src_fmt: PixelFormat,
1150        dst_fmt: PixelFormat,
1151        rotation: Rotation,
1152        flip: Flip,
1153        crop: ResolvedCrop,
1154        src_params: ColorParams,
1155        dst_params: ColorParams,
1156    ) -> Result<()> {
1157        use PixelFormat::*;
1158
1159        #[cfg(test)]
1160        {
1161            self.last_tmp_dims = None;
1162        }
1163
1164        let src_w = src.width().unwrap();
1165        let src_h = src.height().unwrap();
1166        let dst_w = dst.width().unwrap();
1167        let dst_h = dst.height().unwrap();
1168
1169        crop.check_crop_dims(src_w, src_h, dst_w, dst_h)?;
1170
1171        // Determine intermediate format for the resize step
1172        let intermediate = match (src_fmt, dst_fmt) {
1173            (Nv12, Rgb) => Rgb,
1174            (Nv12, Rgba) => Rgba,
1175            (Nv12, Grey) => Grey,
1176            (Nv12, Yuyv) => Rgba,
1177            (Nv12, Nv16) => Rgba,
1178            (Nv12, PlanarRgb) => Rgb,
1179            (Nv12, PlanarRgba) => Rgba,
1180            (Nv16, PlanarRgb) => Rgb,
1181            (Nv16, PlanarRgba) => Rgba,
1182            (Nv24, PlanarRgb) => Rgb,
1183            (Nv24, PlanarRgba) => Rgba,
1184            (Yuyv, Rgb) => Rgb,
1185            (Yuyv, Rgba) => Rgba,
1186            (Yuyv, Grey) => Grey,
1187            (Yuyv, Yuyv) => Rgba,
1188            (Yuyv, PlanarRgb) => Rgb,
1189            (Yuyv, PlanarRgba) => Rgba,
1190            (Yuyv, Nv16) => Rgba,
1191            (Vyuy, Rgb) => Rgb,
1192            (Vyuy, Rgba) => Rgba,
1193            (Vyuy, Grey) => Grey,
1194            (Vyuy, Vyuy) => Rgba,
1195            (Vyuy, PlanarRgb) => Rgb,
1196            (Vyuy, PlanarRgba) => Rgba,
1197            (Vyuy, Nv16) => Rgba,
1198            (Rgba, Rgb) => Rgba,
1199            (Rgba, Rgba) => Rgba,
1200            (Rgba, Grey) => Grey,
1201            (Rgba, Yuyv) => Rgba,
1202            (Rgba, PlanarRgb) => Rgba,
1203            (Rgba, PlanarRgba) => Rgba,
1204            (Rgba, Nv16) => Rgba,
1205            (Rgb, Rgb) => Rgb,
1206            (Rgb, Rgba) => Rgb,
1207            (Rgb, Grey) => Grey,
1208            (Rgb, Yuyv) => Rgb,
1209            (Rgb, PlanarRgb) => Rgb,
1210            (Rgb, PlanarRgba) => Rgb,
1211            (Rgb, Nv16) => Rgb,
1212            (Grey, Rgb) => Rgb,
1213            (Grey, Rgba) => Rgba,
1214            (Grey, Grey) => Grey,
1215            (Grey, Yuyv) => Grey,
1216            (Grey, PlanarRgb) => Grey,
1217            (Grey, PlanarRgba) => Grey,
1218            (Grey, Nv16) => Grey,
1219            (Nv12, Bgra) => Rgba,
1220            (Yuyv, Bgra) => Rgba,
1221            (Vyuy, Bgra) => Rgba,
1222            (Rgba, Bgra) => Rgba,
1223            (Rgb, Bgra) => Rgb,
1224            (Grey, Bgra) => Grey,
1225            (Bgra, Bgra) => Bgra,
1226            (Nv16, Rgb) => Rgb,
1227            (Nv16, Rgba) => Rgba,
1228            (Nv16, Bgra) => Rgba,
1229            (Nv24, Rgb) => Rgb,
1230            (Nv24, Rgba) => Rgba,
1231            (Nv24, Grey) => Grey,
1232            (Nv24, Bgra) => Rgba,
1233            (PlanarRgb, Rgb) => Rgb,
1234            (PlanarRgb, Rgba) => Rgb,
1235            (PlanarRgb, Bgra) => Rgb,
1236            (PlanarRgba, Rgb) => Rgba,
1237            (PlanarRgba, Rgba) => Rgba,
1238            (PlanarRgba, Bgra) => Rgba,
1239            (s, d) => {
1240                return Err(Error::NotSupported(format!("Conversion from {s} to {d}",)));
1241            }
1242        };
1243
1244        let need_resize_flip_rotation = rotation != Rotation::None
1245            || flip != Flip::None
1246            || src_w != dst_w
1247            || src_h != dst_h
1248            || crop.src_rect.is_some_and(|c| {
1249                c != Rect {
1250                    left: 0,
1251                    top: 0,
1252                    width: src_w,
1253                    height: src_h,
1254                }
1255            })
1256            || crop.dst_rect.is_some_and(|c| {
1257                c != Rect {
1258                    left: 0,
1259                    top: 0,
1260                    width: dst_w,
1261                    height: dst_h,
1262                }
1263            });
1264
1265        // Pick the resolved colorimetry for a single conversion by its YUV
1266        // side: YUV→RGB decodes use the source side, RGB→YUV encodes the dest.
1267        let direct_is_yuv_src = matches!(src_fmt, Nv12 | Nv16 | Nv24 | Yuyv | Vyuy);
1268        let direct_params = if direct_is_yuv_src {
1269            src_params
1270        } else {
1271            dst_params
1272        };
1273
1274        // Fused NV→planar: decode the YUV source into packed RGB one
1275        // cache-resident row strip at a time and NEON-deinterleave each strip
1276        // straight into the destination planes, so the full-size packed-RGB
1277        // intermediate never round-trips through DRAM and is not reallocated
1278        // per frame. JPEG decodes to the NV family and the model wants planar
1279        // RGB, so this is the hot Orin CPU-preprocess path.
1280        //
1281        // Two shapes take this path: the whole-frame case (no crop, dst ==
1282        // src size — the original hot path, unchanged) and a **scale-
1283        // identity** source crop (crop size == destination size, no rotate/
1284        // flip, full destination placement, chroma-aligned origin) — the
1285        // primary CPU-fallback path for uniform tiling (e.g. SAHI tiles cut
1286        // from one frame). Any other shape (an actual resize, a partial
1287        // destination placement, or a chroma-misaligned origin) falls
1288        // through to the general pipeline below — never snapped or shifted
1289        // into alignment.
1290        let full_dst_rect = Rect {
1291            left: 0,
1292            top: 0,
1293            width: dst_w,
1294            height: dst_h,
1295        };
1296        let fused_region = if rotation != Rotation::None || flip != Flip::None {
1297            None
1298        } else {
1299            match crop.src_rect {
1300                None if src_w == dst_w && src_h == dst_h => Some(None),
1301                None => None,
1302                Some(r)
1303                    if r.width == dst_w
1304                        && r.height == dst_h
1305                        && crop.dst_rect.is_none_or(|d| d == full_dst_rect)
1306                        && chroma_alignment_ok(src_fmt, r) =>
1307                {
1308                    Some(Some(r))
1309                }
1310                Some(_) => None,
1311            }
1312        };
1313        if let Some(region) = fused_region {
1314            if matches!(src_fmt, Nv12 | Nv16 | Nv24) && matches!(dst_fmt, PlanarRgb | PlanarRgba) {
1315                #[cfg(test)]
1316                {
1317                    self.fused_hits += 1;
1318                }
1319                return self.convert_nv_to_planar_fused(
1320                    src,
1321                    dst,
1322                    src_fmt,
1323                    dst_fmt,
1324                    direct_params,
1325                    region,
1326                );
1327            }
1328        }
1329
1330        // check if a direct conversion can be done
1331        if !need_resize_flip_rotation && Self::support_conversion_pf(src_fmt, dst_fmt) {
1332            return Self::convert_format_pf(src, dst, src_fmt, dst_fmt, direct_params);
1333        }
1334
1335        // any extra checks
1336        if dst_fmt == Yuyv && !dst_w.is_multiple_of(2) {
1337            return Err(Error::NotSupported(format!(
1338                "{} destination must have width divisible by 2",
1339                dst_fmt,
1340            )));
1341        }
1342
1343        // Take the cached intermediates out of `self` so the resize step can
1344        // borrow `self` exclusively; they are restored before returning on the
1345        // success path. Reused buffers are not re-zeroed — every consumer below
1346        // fully overwrites the region it later reads (the pre-resize convert
1347        // writes all of `tmp`; the resize writes the scaled rect of `tmp2` and
1348        // its letterbox border is either pre-filled from `dst` or overwritten in
1349        // `dst` by the final `fill_image_outside_crop_u8`).
1350        let mut cached_tmp = self.convert_tmp.take();
1351        let mut cached_tmp2 = self.convert_tmp2.take();
1352
1353        // For a cropped convert, size the pre-resize intermediate to the crop
1354        // (grown by the resize filter's halo) instead of the whole frame — see
1355        // `pre_resize_region`. `None` keeps the previous full-frame behaviour,
1356        // which is also the uncropped hot path.
1357        let pre_region = if intermediate != src_fmt {
1358            self.pre_resize_region(src_fmt, (src_w, src_h), (dst_w, dst_h), rotation, crop)
1359        } else {
1360            None
1361        };
1362
1363        // create tmp buffer (reusing the cached one when its geometry matches)
1364        let tmp_holder: Option<Tensor<u8>> = if intermediate != src_fmt {
1365            let _s = tracing::trace_span!(
1366                "image.convert.cpu.format_convert",
1367                from = ?src_fmt,
1368                to = ?intermediate,
1369                pass = "pre_resize",
1370            )
1371            .entered();
1372            let (tmp_w, tmp_h) = pre_region.map_or((src_w, src_h), |g| (g.width, g.height));
1373            let mut t = Self::reuse_or_alloc_image(cached_tmp.take(), tmp_w, tmp_h, intermediate)?;
1374            #[cfg(test)]
1375            {
1376                self.last_tmp_dims = Some((tmp_w, tmp_h));
1377            }
1378            match pre_region {
1379                Some(g) => {
1380                    let mut sub = Self::reuse_or_alloc_image(
1381                        self.convert_src_sub.take(),
1382                        g.width,
1383                        g.height,
1384                        src_fmt,
1385                    )?;
1386                    {
1387                        let _s = tracing::trace_span!(
1388                            "image.convert.cpu.extract_region",
1389                            region_w = g.width,
1390                            region_h = g.height,
1391                        )
1392                        .entered();
1393                        Self::extract_nv_region(src, &mut sub, src_fmt, g)?;
1394                    }
1395                    Self::convert_format_pf(&sub, &mut t, src_fmt, intermediate, src_params)?;
1396                    self.convert_src_sub = Some(sub);
1397                }
1398                None => Self::convert_format_pf(src, &mut t, src_fmt, intermediate, src_params)?,
1399            }
1400            Some(t)
1401        } else {
1402            None
1403        };
1404
1405        // The intermediate now starts at `grown`'s origin rather than the
1406        // frame's, so rebase the source crop into its coordinates for the
1407        // resize step. The crop's size — and therefore the resize scale, the
1408        // filter coefficients, and every output pixel — is unchanged.
1409        let crop = match (pre_region, crop.src_rect) {
1410            (Some(g), Some(r)) => ResolvedCrop {
1411                src_rect: Some(Rect {
1412                    left: r.left - g.left,
1413                    top: r.top - g.top,
1414                    ..r
1415                }),
1416                ..crop
1417            },
1418            _ => crop,
1419        };
1420        let (tmp, tmp_fmt): (&Tensor<u8>, PixelFormat) = match &tmp_holder {
1421            Some(t) => (t, intermediate),
1422            None => (src, src_fmt),
1423        };
1424
1425        // format must be RGB/RGBA/GREY
1426        debug_assert!(matches!(tmp_fmt, Rgb | Rgba | Grey));
1427        if tmp_fmt == dst_fmt {
1428            let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
1429            self.resize_flip_rotate_pf(tmp, dst, dst_fmt, rotation, flip, crop)?;
1430        } else if !need_resize_flip_rotation {
1431            let _s = tracing::trace_span!(
1432                "image.convert.cpu.format_convert",
1433                from = ?tmp_fmt,
1434                to = ?dst_fmt,
1435                pass = "direct",
1436            )
1437            .entered();
1438            Self::convert_format_pf(tmp, dst, tmp_fmt, dst_fmt, dst_params)?;
1439        } else {
1440            let mut tmp2 = Self::reuse_or_alloc_image(cached_tmp2.take(), dst_w, dst_h, tmp_fmt)?;
1441            if crop.dst_rect.is_some_and(|c| {
1442                c != Rect {
1443                    left: 0,
1444                    top: 0,
1445                    width: dst_w,
1446                    height: dst_h,
1447                }
1448            }) && crop.dst_color.is_none()
1449            {
1450                Self::convert_format_pf(dst, &mut tmp2, dst_fmt, tmp_fmt, dst_params)?;
1451            }
1452            {
1453                let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
1454                self.resize_flip_rotate_pf(tmp, &mut tmp2, tmp_fmt, rotation, flip, crop)?;
1455            }
1456            {
1457                let _s = tracing::trace_span!(
1458                    "image.convert.cpu.format_convert",
1459                    from = ?tmp_fmt,
1460                    to = ?dst_fmt,
1461                    pass = "post_resize",
1462                )
1463                .entered();
1464                Self::convert_format_pf(&tmp2, dst, tmp_fmt, dst_fmt, dst_params)?;
1465            }
1466            cached_tmp2 = Some(tmp2);
1467        }
1468        // Restore the intermediates to the cache for the next call (`tmp` — a
1469        // borrow of `tmp_holder` — is no longer used past this point).
1470        if let Some(t) = tmp_holder {
1471            cached_tmp = Some(t);
1472        }
1473        self.convert_tmp = cached_tmp;
1474        self.convert_tmp2 = cached_tmp2;
1475
1476        if let (Some(dst_rect), Some(dst_color)) = (crop.dst_rect, crop.dst_color) {
1477            let full_rect = Rect {
1478                left: 0,
1479                top: 0,
1480                width: dst_w,
1481                height: dst_h,
1482            };
1483            if dst_rect != full_rect {
1484                Self::fill_image_outside_crop_u8(dst, dst_color, dst_rect)?;
1485            }
1486        }
1487
1488        Ok(())
1489    }
1490
1491    fn draw_decoded_masks_impl(
1492        &mut self,
1493        dst: &mut Tensor<u8>,
1494        detect: &[DetectBox],
1495        segmentation: &[Segmentation],
1496        opacity: f32,
1497        color_mode: crate::ColorMode,
1498    ) -> Result<()> {
1499        let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
1500        if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
1501            return Err(crate::Error::NotSupported(
1502                "CPU image rendering only supports RGBA or RGB images".to_string(),
1503            ));
1504        }
1505
1506        let _timer = FunctionTimer::new("CPUProcessor::draw_decoded_masks");
1507
1508        let dst_w = dst.width().unwrap();
1509        let dst_h = dst.height().unwrap();
1510        let dst_rs = tensor_row_stride(dst);
1511        let dst_c = dst_fmt.channels();
1512
1513        let mut map = dst.map_mut()?;
1514        let dst_slice = map.as_mut_slice();
1515
1516        self.render_box(dst_w, dst_h, dst_rs, dst_c, dst_slice, detect, color_mode)?;
1517
1518        if segmentation.is_empty() {
1519            return Ok(());
1520        }
1521
1522        // Rank guard BEFORE indexing shape()[2]. `Segmentation.segmentation`
1523        // is a `TensorDyn` with public fields, so a caller of the public
1524        // `draw_decoded_masks` can hand over a rank-2 mask; under the former
1525        // `Array3` that was structurally impossible. The guards inside the
1526        // render functions are unreachable for bad rank because this dispatch
1527        // indexes first.
1528        if let Some(bad) = segmentation
1529            .iter()
1530            .find(|s| s.segmentation.shape().len() != 3)
1531        {
1532            return Err(crate::Error::InvalidShape(format!(
1533                "segmentation must be [H, W, C], got {:?}",
1534                bad.segmentation.shape()
1535            )));
1536        }
1537
1538        // Semantic segmentation (e.g. ModelPack) has C > 1 (multi-class),
1539        // instance segmentation (e.g. YOLO) has C = 1 (binary per-instance).
1540        let is_semantic = segmentation[0].segmentation.shape()[2] > 1;
1541
1542        if is_semantic {
1543            self.render_modelpack_segmentation(
1544                dst_w,
1545                dst_h,
1546                dst_rs,
1547                dst_c,
1548                dst_slice,
1549                &segmentation[0],
1550                opacity,
1551            )?;
1552        } else {
1553            for (idx, (seg, det)) in segmentation.iter().zip(detect).enumerate() {
1554                let color_index = color_mode.index(idx, det.label);
1555                self.render_yolo_segmentation(
1556                    dst_w,
1557                    dst_h,
1558                    dst_rs,
1559                    dst_c,
1560                    dst_slice,
1561                    seg,
1562                    color_index,
1563                    opacity,
1564                )?;
1565            }
1566        }
1567
1568        Ok(())
1569    }
1570
1571    fn draw_proto_masks_impl(
1572        &mut self,
1573        dst: &mut Tensor<u8>,
1574        detect: &[DetectBox],
1575        proto_data: &ProtoData,
1576        opacity: f32,
1577        letterbox: Option<[f32; 4]>,
1578        color_mode: crate::ColorMode,
1579    ) -> Result<()> {
1580        let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
1581        if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
1582            return Err(crate::Error::NotSupported(
1583                "CPU image rendering only supports RGBA or RGB images".to_string(),
1584            ));
1585        }
1586
1587        let _timer = FunctionTimer::new("CPUProcessor::draw_proto_masks");
1588
1589        let dst_w = dst.width().unwrap();
1590        let dst_h = dst.height().unwrap();
1591        let dst_rs = tensor_row_stride(dst);
1592        let channels = dst_fmt.channels();
1593
1594        let mut map = dst.map_mut()?;
1595        let dst_slice = map.as_mut_slice();
1596
1597        self.render_box(
1598            dst_w, dst_h, dst_rs, channels, dst_slice, detect, color_mode,
1599        )?;
1600
1601        if detect.is_empty() {
1602            return Ok(());
1603        }
1604        let proto_shape = proto_data.protos.shape();
1605        if proto_shape.len() != 3 {
1606            return Err(Error::InvalidShape(format!(
1607                "protos tensor must be rank-3, got {proto_shape:?}"
1608            )));
1609        }
1610        let proto_h = proto_shape[0];
1611        let proto_w = proto_shape[1];
1612        let num_protos = proto_shape[2];
1613        let coeff_shape = proto_data.mask_coefficients.shape();
1614        if coeff_shape.len() != 2 {
1615            return Err(Error::InvalidShape(format!(
1616                "mask_coefficients tensor must be rank-2, got {coeff_shape:?}"
1617            )));
1618        }
1619        // Genuine "no detections this frame" → nothing to render.
1620        if coeff_shape[0] == 0 {
1621            return Ok(());
1622        }
1623        if coeff_shape[1] != num_protos {
1624            return Err(Error::InvalidShape(format!(
1625                "mask_coefficients second dimension must match num_protos \
1626                 ({num_protos}), got {coeff_shape:?}"
1627            )));
1628        }
1629
1630        // Widen coefficients to f32 once; shape [N, num_protos].
1631        let coeff_f32: Vec<f32> = match proto_data.mask_coefficients.dtype() {
1632            DType::F32 => {
1633                let t = proto_data.mask_coefficients.as_typed::<f32>().expect("F32");
1634                let m = t.map_read()?;
1635                m.as_slice().to_vec()
1636            }
1637            DType::F16 => {
1638                let t = proto_data
1639                    .mask_coefficients
1640                    .as_typed::<half::f16>()
1641                    .expect("F16");
1642                let m = t.map_read()?;
1643                m.as_slice().iter().map(|v| v.to_f32()).collect()
1644            }
1645            DType::I8 => {
1646                let t = proto_data.mask_coefficients.as_typed::<i8>().expect("I8");
1647                let m = t.map_read()?;
1648                if let Some(q) = t.quantization() {
1649                    use edgefirst_tensor::QuantMode;
1650                    let (scale, zp) = match q.mode() {
1651                        QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1652                        QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1653                        other => {
1654                            return Err(Error::NotSupported(format!(
1655                                "I8 mask_coefficients quantization mode {other:?} not supported"
1656                            )));
1657                        }
1658                    };
1659                    m.as_slice()
1660                        .iter()
1661                        .map(|&v| (v as f32 - zp) * scale)
1662                        .collect()
1663                } else {
1664                    m.as_slice().iter().map(|&v| v as f32).collect()
1665                }
1666            }
1667            DType::I16 => {
1668                let t = proto_data.mask_coefficients.as_typed::<i16>().expect("I16");
1669                let m = t.map_read()?;
1670                if let Some(q) = t.quantization() {
1671                    use edgefirst_tensor::QuantMode;
1672                    let (scale, zp) = match q.mode() {
1673                        QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1674                        QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1675                        other => {
1676                            return Err(Error::NotSupported(format!(
1677                                "I16 mask_coefficients quantization mode {other:?} not supported"
1678                            )));
1679                        }
1680                    };
1681                    m.as_slice()
1682                        .iter()
1683                        .map(|&v| (v as f32 - zp) * scale)
1684                        .collect()
1685                } else {
1686                    m.as_slice().iter().map(|&v| v as f32).collect()
1687                }
1688            }
1689            other => {
1690                return Err(Error::InvalidShape(format!(
1691                    "mask_coefficients dtype {other:?} not supported"
1692                )));
1693            }
1694        };
1695
1696        // Precompute letterbox scale/offset for output-pixel → proto-pixel mapping.
1697        let (lx0, lx_range, ly0, ly_range) = match letterbox {
1698            Some([lx0, ly0, lx1, ly1]) => (lx0, lx1 - lx0, ly0, ly1 - ly0),
1699            None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
1700        };
1701
1702        // Per-dtype dispatch. Map protos once, call the inner draw loop
1703        // with a dtype-specialized loader closure.
1704        match proto_data.protos.dtype() {
1705            DType::F32 => {
1706                let t = proto_data.protos.as_typed::<f32>().expect("F32");
1707                let m = t.map_read()?;
1708                self.draw_proto_masks_inner(
1709                    dst_slice,
1710                    dst_w,
1711                    dst_h,
1712                    dst_rs,
1713                    channels,
1714                    detect,
1715                    m.as_slice(),
1716                    &coeff_f32,
1717                    proto_h,
1718                    proto_w,
1719                    num_protos,
1720                    opacity,
1721                    (lx0, lx_range, ly0, ly_range),
1722                    color_mode,
1723                    0.0_f32,
1724                    |p: &f32, _| *p,
1725                );
1726            }
1727            DType::F16 => {
1728                let t = proto_data.protos.as_typed::<half::f16>().expect("F16");
1729                let m = t.map_read()?;
1730                self.draw_proto_masks_inner(
1731                    dst_slice,
1732                    dst_w,
1733                    dst_h,
1734                    dst_rs,
1735                    channels,
1736                    detect,
1737                    m.as_slice(),
1738                    &coeff_f32,
1739                    proto_h,
1740                    proto_w,
1741                    num_protos,
1742                    opacity,
1743                    (lx0, lx_range, ly0, ly_range),
1744                    color_mode,
1745                    0.0_f32,
1746                    |p: &half::f16, _| p.to_f32(),
1747                );
1748            }
1749            DType::I8 => {
1750                use edgefirst_tensor::QuantMode;
1751                let t = proto_data.protos.as_typed::<i8>().expect("I8");
1752                let m = t.map_read()?;
1753                let quant = t.quantization().ok_or_else(|| {
1754                    Error::InvalidShape("I8 protos require quantization metadata".into())
1755                })?;
1756                let (scale, zp) = match quant.mode() {
1757                    QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1758                    QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1759                    QuantMode::PerChannel { axis, .. }
1760                    | QuantMode::PerChannelSymmetric { axis, .. } => {
1761                        return Err(Error::NotSupported(format!(
1762                            "per-channel quantization (axis={axis}) in draw_proto_masks \
1763                             CPU path not yet supported"
1764                        )));
1765                    }
1766                };
1767                self.draw_proto_masks_inner(
1768                    dst_slice,
1769                    dst_w,
1770                    dst_h,
1771                    dst_rs,
1772                    channels,
1773                    detect,
1774                    m.as_slice(),
1775                    &coeff_f32,
1776                    proto_h,
1777                    proto_w,
1778                    num_protos,
1779                    opacity,
1780                    (lx0, lx_range, ly0, ly_range),
1781                    color_mode,
1782                    scale,
1783                    move |p: &i8, _| (*p as f32) - zp,
1784                );
1785            }
1786            other => {
1787                return Err(Error::InvalidShape(format!(
1788                    "proto tensor dtype {other:?} not supported"
1789                )));
1790            }
1791        }
1792
1793        Ok(())
1794    }
1795
1796    #[allow(clippy::too_many_arguments)]
1797    fn draw_proto_masks_inner<P: Copy>(
1798        &self,
1799        dst_slice: &mut [u8],
1800        dst_w: usize,
1801        dst_h: usize,
1802        dst_rs: usize,
1803        channels: usize,
1804        detect: &[DetectBox],
1805        protos: &[P],
1806        coeff_all_f32: &[f32],
1807        proto_h: usize,
1808        proto_w: usize,
1809        num_protos: usize,
1810        opacity: f32,
1811        letterbox_xy: (f32, f32, f32, f32),
1812        color_mode: crate::ColorMode,
1813        acc_scale: f32,
1814        load_f32: impl Fn(&P, f32) -> f32 + Copy,
1815    ) {
1816        let (lx0, lx_range, ly0, ly_range) = letterbox_xy;
1817        let stride_y = proto_w * num_protos;
1818        for (idx, det) in detect.iter().enumerate() {
1819            let coeff = &coeff_all_f32[idx * num_protos..(idx + 1) * num_protos];
1820            let color_index = color_mode.index(idx, det.label);
1821            let color = self.colors[color_index % self.colors.len()];
1822            let alpha = if opacity == 1.0 {
1823                color[3] as u16
1824            } else {
1825                (color[3] as f32 * opacity).round() as u16
1826            };
1827
1828            let start_x = (dst_w as f32 * det.bbox.xmin).round() as usize;
1829            let start_y = (dst_h as f32 * det.bbox.ymin).round() as usize;
1830            let end_x = ((dst_w as f32 * det.bbox.xmax).round() as usize).min(dst_w);
1831            let end_y = ((dst_h as f32 * det.bbox.ymax).round() as usize).min(dst_h);
1832
1833            for y in start_y..end_y {
1834                for x in start_x..end_x {
1835                    let px = (lx0 + (x as f32 / dst_w as f32) * lx_range) * proto_w as f32 - 0.5;
1836                    let py = (ly0 + (y as f32 / dst_h as f32) * ly_range) * proto_h as f32 - 0.5;
1837
1838                    // Bilinear interpolation with per-load widening. Inline
1839                    // bilinear-sample since bilinear_dot_slice takes a
1840                    // different closure shape (no `zp` arg).
1841                    let x0 = (px.floor() as isize).clamp(0, proto_w as isize - 1) as usize;
1842                    let y0 = (py.floor() as isize).clamp(0, proto_h as isize - 1) as usize;
1843                    let x1 = (x0 + 1).min(proto_w - 1);
1844                    let y1 = (y0 + 1).min(proto_h - 1);
1845                    let fx = px - px.floor();
1846                    let fy = py - py.floor();
1847                    let w00 = (1.0 - fx) * (1.0 - fy);
1848                    let w10 = fx * (1.0 - fy);
1849                    let w01 = (1.0 - fx) * fy;
1850                    let w11 = fx * fy;
1851                    let b00 = y0 * stride_y + x0 * num_protos;
1852                    let b10 = y0 * stride_y + x1 * num_protos;
1853                    let b01 = y1 * stride_y + x0 * num_protos;
1854                    let b11 = y1 * stride_y + x1 * num_protos;
1855                    let mut acc = 0.0_f32;
1856                    for p in 0..num_protos {
1857                        let v00 = load_f32(&protos[b00 + p], 0.0);
1858                        let v10 = load_f32(&protos[b10 + p], 0.0);
1859                        let v01 = load_f32(&protos[b01 + p], 0.0);
1860                        let v11 = load_f32(&protos[b11 + p], 0.0);
1861                        let val = w00 * v00 + w10 * v10 + w01 * v01 + w11 * v11;
1862                        acc += coeff[p] * val;
1863                    }
1864                    let final_acc = if acc_scale == 0.0 {
1865                        acc
1866                    } else {
1867                        acc_scale * acc
1868                    };
1869                    // Pass-through: acc_scale=0.0 means "no scaling" (f32/f16
1870                    // native); non-zero means "apply scale once" (i8 with
1871                    // per-tensor quant).
1872                    let mask = 1.0 / (1.0 + (-final_acc).exp());
1873                    if mask < 0.5 {
1874                        continue;
1875                    }
1876                    let dst_index = y * dst_rs + x * channels;
1877                    for c in 0..3 {
1878                        dst_slice[dst_index + c] = ((color[c] as u16 * alpha
1879                            + dst_slice[dst_index + c] as u16 * (255 - alpha))
1880                            / 255) as u8;
1881                    }
1882                }
1883            }
1884        }
1885    }
1886}