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_decoder::{DetectBox, ProtoData, Segmentation};
8use edgefirst_tensor::{
9    DType, PixelFormat, Tensor, TensorDyn, TensorMapTrait, TensorMemory, TensorTrait,
10};
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 cache-resident scratch for the strip-fused NV→planar path
71    /// (`convert_nv_to_planar_fused`): holds a single row strip of packed RGB
72    /// (`STRIP_ROWS * width * 3` bytes). Keeping it on the processor avoids a
73    /// per-frame allocation and lets each strip stay hot in L2 between the YUV
74    /// decode and the deinterleave. Grown on demand; never shrunk.
75    nv_strip_scratch: Vec<u8>,
76    /// Reusable intermediate buffers for the multi-step convert pipeline
77    /// (pre-resize format-convert and the resized-RGB scratch). Reused across
78    /// frames when the dimensions/format match so the steady-state
79    /// letterbox/resize loop does not reallocate — and `alloc_zeroed`-clear — a
80    /// full-frame buffer per call. The region each consumer reads is always
81    /// fully overwritten first, so reused (non-zeroed) contents are never read.
82    convert_tmp: Option<Tensor<u8>>,
83    convert_tmp2: Option<Tensor<u8>>,
84}
85
86// `CPUProcessor` was `#[derive(Clone)]` before the `widen_scratch` field was
87// added; `TensorDyn` is not `Clone`, so the derive no longer applies. Restore
88// the public `Clone` impl by hand to avoid a breaking API change. The scratch
89// cache is a private allocation amortiser, not part of the logical value, so a
90// clone starts empty rather than sharing or duplicating it.
91impl Clone for CPUProcessor {
92    fn clone(&self) -> Self {
93        Self {
94            resizer: self.resizer.clone(),
95            options: self.options,
96            colors: self.colors,
97            widen_scratch: None,
98            resize_destride_scratch: Vec::new(),
99            nv_strip_scratch: Vec::new(),
100            convert_tmp: None,
101            convert_tmp2: None,
102        }
103    }
104}
105
106unsafe impl Send for CPUProcessor {}
107unsafe impl Sync for CPUProcessor {}
108
109impl Default for CPUProcessor {
110    fn default() -> Self {
111        Self::new_bilinear()
112    }
113}
114
115/// Write the base layer of `dst` before mask rendering.
116///
117/// This is the terminal fallback: on CPU we have no 2D hardware, so a
118/// direct buffer write is the appropriate primitive. The invariant is that
119/// every call to the CPU draw_* entry points fully initialises dst — we
120/// never rely on "whatever was in the buffer" from the caller.
121///
122/// - `background == Some(bg)` → byte-for-byte copy bg → dst (after shape /
123///   format validation).
124/// - `background == None` → fill dst with 0x00 (transparent black).
125fn prepare_dst_base_cpu(dst: &mut TensorDyn, background: Option<&TensorDyn>) -> Result<()> {
126    match background {
127        Some(bg) => {
128            if bg.shape() != dst.shape() {
129                return Err(Error::InvalidShape(
130                    "background shape does not match dst".into(),
131                ));
132            }
133            if bg.format() != dst.format() {
134                return Err(Error::InvalidShape(
135                    "background pixel format does not match dst".into(),
136                ));
137            }
138            let bg_u8 = bg.as_u8().ok_or(Error::NotAnImage)?;
139            let dst_u8 = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
140            let bg_map = bg_u8.map_read()?;
141            let mut dst_map = dst_u8.map_mut()?;
142            let bg_slice = bg_map.as_slice();
143            let dst_slice = dst_map.as_mut_slice();
144            if bg_slice.len() != dst_slice.len() {
145                return Err(Error::InvalidShape(
146                    "background buffer size does not match dst".into(),
147                ));
148            }
149            dst_slice.copy_from_slice(bg_slice);
150        }
151        None => {
152            let dst_u8 = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
153            let mut dst_map = dst_u8.map_mut()?;
154            dst_map.as_mut_slice().fill(0);
155        }
156    }
157    Ok(())
158}
159
160/// Compute row stride for a packed-format Tensor<u8> image given its format.
161fn row_stride_for(width: usize, fmt: PixelFormat) -> usize {
162    use edgefirst_tensor::PixelLayout;
163    match fmt.layout() {
164        PixelLayout::Packed => width * fmt.channels(),
165        PixelLayout::Planar | PixelLayout::SemiPlanar => width,
166        _ => width, // fallback for non-exhaustive
167    }
168}
169
170/// Read the effective row stride from a tensor, falling back to the computed
171/// minimum stride if the tensor has no explicit stride set. This correctly
172/// handles tensors with GPU pitch-alignment padding (e.g., from
173/// `ImageProcessor::create_image()` or codec strided decode).
174fn tensor_row_stride(tensor: &Tensor<u8>) -> usize {
175    tensor.effective_row_stride().unwrap_or_else(|| {
176        let w = tensor.width().unwrap_or(0);
177        let fmt = tensor.format().unwrap_or(PixelFormat::Rgb);
178        row_stride_for(w, fmt)
179    })
180}
181
182/// Split a contiguous semi-planar (NV12/NV16/NV24) tensor's mapped bytes into
183/// `(luma, chroma)` planes at the `stride * src_h` boundary, validating the map
184/// holds the full combined plane first. A bare `split_at` would panic if an
185/// imported tensor's caller-supplied dimensions/stride exceed its actual buffer
186/// (untrusted input), so this returns `Error::InvalidShape` instead.
187fn split_semi_planar(
188    bytes: &[u8],
189    stride: usize,
190    src_h: usize,
191    fmt: PixelFormat,
192) -> Result<(&[u8], &[u8])> {
193    let total_h = fmt.combined_plane_height(src_h).unwrap_or(src_h);
194    let need = stride.checked_mul(total_h).ok_or_else(|| {
195        Error::InvalidShape(format!(
196            "{fmt:?} plane size overflow (stride={stride}, h={src_h})"
197        ))
198    })?;
199    if bytes.len() < need {
200        return Err(Error::InvalidShape(format!(
201            "{fmt:?} source has {} bytes but needs {need} (stride={stride}, h={src_h})",
202            bytes.len()
203        )));
204    }
205    Ok(bytes.split_at(stride * src_h))
206}
207
208/// Mutable mirror of [`split_semi_planar`]: split a contiguous semi-planar
209/// destination's mapped bytes into `(luma, chroma)` planes at the
210/// `stride * dst_h` boundary, validating the map holds the full combined plane
211/// first. A bare `split_at_mut` would panic if a caller-supplied (untrusted)
212/// destination's declared dimensions/stride exceed its actual buffer, so this
213/// returns `Error::InvalidShape` instead.
214fn split_semi_planar_mut(
215    bytes: &mut [u8],
216    stride: usize,
217    dst_h: usize,
218    fmt: PixelFormat,
219) -> Result<(&mut [u8], &mut [u8])> {
220    let total_h = fmt.combined_plane_height(dst_h).unwrap_or(dst_h);
221    let need = stride.checked_mul(total_h).ok_or_else(|| {
222        Error::InvalidShape(format!(
223            "{fmt:?} plane size overflow (stride={stride}, combined_h={total_h})"
224        ))
225    })?;
226    if bytes.len() < need {
227        return Err(Error::InvalidShape(format!(
228            "{fmt:?} destination has {} bytes but needs {need} (stride={stride}, combined_h={total_h})",
229            bytes.len()
230        )));
231    }
232    Ok(bytes.split_at_mut(stride * dst_h))
233}
234
235/// Validate that a mapped plane buffer of `buf_len` bytes can hold `rows` rows
236/// of `stride` bytes each, with `row_bytes` valid bytes per row. Returns
237/// `Error::InvalidShape` (instead of letting a later row slice panic) when a
238/// caller-supplied stride/shape exceeds the actual allocation. `what` labels the
239/// buffer in the error message.
240fn guard_plane(
241    buf_len: usize,
242    stride: usize,
243    rows: usize,
244    row_bytes: usize,
245    what: &str,
246) -> Result<()> {
247    let need = stride.checked_mul(rows).ok_or_else(|| {
248        Error::InvalidShape(format!(
249            "{what} plane size overflow (stride={stride}, rows={rows})"
250        ))
251    })?;
252    if row_bytes > stride || buf_len < need {
253        return Err(Error::InvalidShape(format!(
254            "{what} buffer too small: {buf_len} bytes, need {need} (stride={stride}, rows={rows}, row_bytes={row_bytes})"
255        )));
256    }
257    Ok(())
258}
259
260/// Apply XOR 0x80 bias to color channels only, preserving alpha.
261///
262/// Matches GL int8 shader behavior: `vec4(int8_bias(c.rgb), c.a)`.
263/// For formats without alpha, XORs every byte (fast path).
264pub(crate) fn apply_int8_xor_bias(data: &mut [u8], fmt: PixelFormat) {
265    use edgefirst_tensor::PixelLayout;
266    if !fmt.has_alpha() {
267        for b in data.iter_mut() {
268            *b ^= 0x80;
269        }
270    } else if fmt.layout() == PixelLayout::Planar {
271        // Planar with alpha (e.g. PlanarRgba): XOR color planes, skip alpha plane.
272        let channels = fmt.channels();
273        let plane_size = data.len() / channels;
274        for b in data[..plane_size * (channels - 1)].iter_mut() {
275            *b ^= 0x80;
276        }
277    } else {
278        // Packed with alpha (Rgba, Bgra): XOR color bytes, skip alpha byte.
279        let channels = fmt.channels();
280        for pixel in data.chunks_exact_mut(channels) {
281            for b in &mut pixel[..channels - 1] {
282                *b ^= 0x80;
283            }
284        }
285    }
286}
287
288impl CPUProcessor {
289    /// Creates a new CPUConverter with bilinear resizing.
290    pub fn new() -> Self {
291        Self::new_bilinear()
292    }
293
294    /// Creates a new CPUConverter with bilinear resizing.
295    fn new_bilinear() -> Self {
296        let resizer = fast_image_resize::Resizer::new();
297        let options = fast_image_resize::ResizeOptions::new()
298            .resize_alg(fast_image_resize::ResizeAlg::Convolution(
299                fast_image_resize::FilterType::Bilinear,
300            ))
301            .use_alpha(false);
302
303        log::debug!("CPUConverter created");
304        Self {
305            resizer,
306            options,
307            colors: crate::DEFAULT_COLORS_U8,
308            widen_scratch: None,
309            resize_destride_scratch: Vec::new(),
310            nv_strip_scratch: Vec::new(),
311            convert_tmp: None,
312            convert_tmp2: None,
313        }
314    }
315
316    /// Creates a new CPUConverter with nearest neighbor resizing.
317    pub fn new_nearest() -> Self {
318        let resizer = fast_image_resize::Resizer::new();
319        let options = fast_image_resize::ResizeOptions::new()
320            .resize_alg(fast_image_resize::ResizeAlg::Nearest)
321            .use_alpha(false);
322        log::debug!("CPUConverter created");
323        Self {
324            resizer,
325            options,
326            colors: crate::DEFAULT_COLORS_U8,
327            widen_scratch: None,
328            resize_destride_scratch: Vec::new(),
329            nv_strip_scratch: Vec::new(),
330            convert_tmp: None,
331            convert_tmp2: None,
332        }
333    }
334
335    pub(crate) fn support_conversion_pf(src: PixelFormat, dst: PixelFormat) -> bool {
336        use PixelFormat::*;
337        matches!(
338            (src, dst),
339            (Nv12, Rgb)
340                | (Nv12, Rgba)
341                | (Nv12, Grey)
342                | (Nv16, Rgb)
343                | (Nv16, Rgba)
344                | (Nv16, Bgra)
345                | (Nv24, Rgb)
346                | (Nv24, Rgba)
347                | (Nv24, Grey)
348                | (Nv24, Bgra)
349                | (Yuyv, Rgb)
350                | (Yuyv, Rgba)
351                | (Yuyv, Grey)
352                | (Yuyv, Yuyv)
353                | (Yuyv, PlanarRgb)
354                | (Yuyv, PlanarRgba)
355                | (Yuyv, Nv16)
356                | (Vyuy, Rgb)
357                | (Vyuy, Rgba)
358                | (Vyuy, Grey)
359                | (Vyuy, Vyuy)
360                | (Vyuy, PlanarRgb)
361                | (Vyuy, PlanarRgba)
362                | (Vyuy, Nv16)
363                | (Rgba, Rgb)
364                | (Rgba, Rgba)
365                | (Rgba, Grey)
366                | (Rgba, Yuyv)
367                | (Rgba, PlanarRgb)
368                | (Rgba, PlanarRgba)
369                | (Rgba, Nv16)
370                | (Rgb, Rgb)
371                | (Rgb, Rgba)
372                | (Rgb, Grey)
373                | (Rgb, Yuyv)
374                | (Rgb, PlanarRgb)
375                | (Rgb, PlanarRgba)
376                | (Rgb, Nv16)
377                | (Grey, Rgb)
378                | (Grey, Rgba)
379                | (Grey, Grey)
380                | (Grey, Yuyv)
381                | (Grey, PlanarRgb)
382                | (Grey, PlanarRgba)
383                | (Grey, Nv16)
384                | (Nv12, Bgra)
385                | (Yuyv, Bgra)
386                | (Vyuy, Bgra)
387                | (Rgba, Bgra)
388                | (Rgb, Bgra)
389                | (Grey, Bgra)
390                | (Bgra, Bgra)
391                | (PlanarRgb, Rgb)
392                | (PlanarRgb, Rgba)
393                | (PlanarRgba, Rgb)
394                | (PlanarRgba, Rgba)
395                | (PlanarRgb, Bgra)
396                | (PlanarRgba, Bgra)
397        )
398    }
399
400    /// Format conversion dispatch for Tensor<u8> with PixelFormat metadata.
401    pub(crate) fn convert_format_pf(
402        src: &Tensor<u8>,
403        dst: &mut Tensor<u8>,
404        src_fmt: PixelFormat,
405        dst_fmt: PixelFormat,
406        cp: ColorParams,
407    ) -> Result<()> {
408        let _timer = FunctionTimer::new(format!(
409            "ImageProcessor::convert_format {} to {}",
410            src_fmt, dst_fmt,
411        ));
412
413        use PixelFormat::*;
414        match (src_fmt, dst_fmt) {
415            (Nv12, Rgb) => Self::convert_nv12_to_rgb(src, dst, cp),
416            (Nv12, Rgba) => Self::convert_nv12_to_rgba(src, dst, cp),
417            (Nv12, Grey) => Self::convert_nv12_to_grey(src, dst, cp),
418            (Yuyv, Rgb) => Self::convert_yuyv_to_rgb(src, dst, cp),
419            (Yuyv, Rgba) => Self::convert_yuyv_to_rgba(src, dst, cp),
420            (Yuyv, Grey) => Self::convert_yuyv_to_grey(src, dst, cp),
421            (Yuyv, Yuyv) => Self::copy_image(src, dst),
422            (Yuyv, PlanarRgb) => Self::convert_yuyv_to_8bps(src, dst, cp),
423            (Yuyv, PlanarRgba) => Self::convert_yuyv_to_prgba(src, dst, cp),
424            (Yuyv, Nv16) => Self::convert_yuyv_to_nv16(src, dst),
425            (Vyuy, Rgb) => Self::convert_vyuy_to_rgb(src, dst, cp),
426            (Vyuy, Rgba) => Self::convert_vyuy_to_rgba(src, dst, cp),
427            (Vyuy, Grey) => Self::convert_vyuy_to_grey(src, dst, cp),
428            (Vyuy, Vyuy) => Self::copy_image(src, dst),
429            (Vyuy, PlanarRgb) => Self::convert_vyuy_to_8bps(src, dst, cp),
430            (Vyuy, PlanarRgba) => Self::convert_vyuy_to_prgba(src, dst, cp),
431            (Vyuy, Nv16) => Self::convert_vyuy_to_nv16(src, dst),
432            (Rgba, Rgb) => Self::convert_rgba_to_rgb(src, dst),
433            (Rgba, Rgba) => Self::copy_image(src, dst),
434            (Rgba, Grey) => Self::convert_rgba_to_grey(src, dst),
435            (Rgba, Yuyv) => Self::convert_rgba_to_yuyv(src, dst, cp),
436            (Rgba, PlanarRgb) => Self::convert_rgba_to_8bps(src, dst),
437            (Rgba, PlanarRgba) => Self::convert_rgba_to_prgba(src, dst),
438            (Rgba, Nv16) => Self::convert_rgba_to_nv16(src, dst, cp),
439            (Rgb, Rgb) => Self::copy_image(src, dst),
440            (Rgb, Rgba) => Self::convert_rgb_to_rgba(src, dst),
441            (Rgb, Grey) => Self::convert_rgb_to_grey(src, dst),
442            (Rgb, Yuyv) => Self::convert_rgb_to_yuyv(src, dst, cp),
443            (Rgb, PlanarRgb) => Self::convert_rgb_to_8bps(src, dst),
444            (Rgb, PlanarRgba) => Self::convert_rgb_to_prgba(src, dst),
445            (Rgb, Nv16) => Self::convert_rgb_to_nv16(src, dst, cp),
446            (Grey, Rgb) => Self::convert_grey_to_rgb(src, dst),
447            (Grey, Rgba) => Self::convert_grey_to_rgba(src, dst),
448            (Grey, Grey) => Self::copy_image(src, dst),
449            (Grey, Yuyv) => Self::convert_grey_to_yuyv(src, dst, cp),
450            (Grey, PlanarRgb) => Self::convert_grey_to_8bps(src, dst),
451            (Grey, PlanarRgba) => Self::convert_grey_to_prgba(src, dst),
452            (Grey, Nv16) => Self::convert_grey_to_nv16(src, dst, cp),
453
454            // the following converts are added for use in testing
455            (Nv16, Rgb) => Self::convert_nv16_to_rgb(src, dst, cp),
456            (Nv16, Rgba) => Self::convert_nv16_to_rgba(src, dst, cp),
457            (Nv24, Rgb) => Self::convert_nv24_to_rgb(src, dst, cp),
458            (Nv24, Rgba) => Self::convert_nv24_to_rgba(src, dst, cp),
459            (Nv24, Grey) => Self::convert_nv24_to_grey(src, dst, cp),
460            (PlanarRgb, Rgb) => Self::convert_8bps_to_rgb(src, dst),
461            (PlanarRgb, Rgba) => Self::convert_8bps_to_rgba(src, dst),
462            (PlanarRgba, Rgb) => Self::convert_prgba_to_rgb(src, dst),
463            (PlanarRgba, Rgba) => Self::convert_prgba_to_rgba(src, dst),
464
465            // BGRA destination: convert to RGBA layout, then swap R and B
466            (Bgra, Bgra) => Self::copy_image(src, dst),
467            (Nv12, Bgra) => {
468                Self::convert_nv12_to_rgba(src, dst, cp)?;
469                Self::swizzle_rb_4chan(dst)
470            }
471            (Nv16, Bgra) => {
472                Self::convert_nv16_to_rgba(src, dst, cp)?;
473                Self::swizzle_rb_4chan(dst)
474            }
475            (Nv24, Bgra) => {
476                Self::convert_nv24_to_rgba(src, dst, cp)?;
477                Self::swizzle_rb_4chan(dst)
478            }
479            (Yuyv, Bgra) => {
480                Self::convert_yuyv_to_rgba(src, dst, cp)?;
481                Self::swizzle_rb_4chan(dst)
482            }
483            (Vyuy, Bgra) => {
484                Self::convert_vyuy_to_rgba(src, dst, cp)?;
485                Self::swizzle_rb_4chan(dst)
486            }
487            (Rgba, Bgra) => {
488                dst.map_mut()?.copy_from_slice(&src.map_read()?);
489                Self::swizzle_rb_4chan(dst)
490            }
491            (Rgb, Bgra) => {
492                Self::convert_rgb_to_rgba(src, dst)?;
493                Self::swizzle_rb_4chan(dst)
494            }
495            (Grey, Bgra) => {
496                Self::convert_grey_to_rgba(src, dst)?;
497                Self::swizzle_rb_4chan(dst)
498            }
499            (PlanarRgb, Bgra) => {
500                Self::convert_8bps_to_rgba(src, dst)?;
501                Self::swizzle_rb_4chan(dst)
502            }
503            (PlanarRgba, Bgra) => {
504                Self::convert_prgba_to_rgba(src, dst)?;
505                Self::swizzle_rb_4chan(dst)
506            }
507
508            (s, d) => Err(Error::NotSupported(format!("Conversion from {s} to {d}",))),
509        }
510    }
511
512    /// Tensor<u8>-based fill_image_outside_crop.
513    pub(crate) fn fill_image_outside_crop_u8(
514        dst: &mut Tensor<u8>,
515        rgba: [u8; 4],
516        crop: Rect,
517    ) -> Result<()> {
518        let dst_fmt = dst.format().unwrap();
519        let dst_w = dst.width().unwrap();
520        let dst_h = dst.height().unwrap();
521        // Resolve the YUV fill encoding from the destination tensor so the
522        // border color matches the same matrix/range as a later YUV→RGB
523        // decode of this image. RGB/Grey fills ignore these params.
524        let cm = crate::colorimetry::resolve_colorimetry(dst.colorimetry(), dst.height());
525        let cp = ColorParams {
526            matrix: crate::colorimetry::yuv_matrix(cm.encoding.unwrap()),
527            range: crate::colorimetry::yuv_range(cm.range.unwrap()),
528            encoding: cm.encoding.unwrap(),
529            range_kind: cm.range.unwrap(),
530            src_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
531            dst_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
532        };
533        let mut dst_map = dst.map_mut()?;
534        let dst_tup = (dst_map.as_mut_slice(), dst_w, dst_h);
535        Self::fill_outside_crop_dispatch(dst_tup, dst_fmt, rgba, crop, cp)
536    }
537
538    /// Common fill dispatch by format.
539    fn fill_outside_crop_dispatch(
540        dst: (&mut [u8], usize, usize),
541        fmt: PixelFormat,
542        rgba: [u8; 4],
543        crop: Rect,
544        cp: ColorParams,
545    ) -> Result<()> {
546        use PixelFormat::*;
547        match fmt {
548            Rgba | Bgra => Self::fill_image_outside_crop_(dst, rgba, crop),
549            Rgb => Self::fill_image_outside_crop_(dst, Self::rgba_to_rgb(rgba), crop),
550            Grey => Self::fill_image_outside_crop_(dst, Self::rgba_to_grey(rgba), crop),
551            Yuyv => Self::fill_image_outside_crop_(
552                (dst.0, dst.1 / 2, dst.2),
553                Self::rgba_to_yuyv(rgba, cp),
554                Rect::new(crop.left / 2, crop.top, crop.width.div_ceil(2), crop.height),
555            ),
556            PlanarRgb => Self::fill_image_outside_crop_planar(dst, Self::rgba_to_rgb(rgba), crop),
557            PlanarRgba => Self::fill_image_outside_crop_planar(dst, rgba, crop),
558            Nv16 => {
559                let yuyv = Self::rgba_to_yuyv(rgba, cp);
560                Self::fill_image_outside_crop_yuv_semiplanar(dst, yuyv[0], [yuyv[1], yuyv[3]], crop)
561            }
562            _ => Err(Error::Internal(format!(
563                "Found unexpected destination {fmt}",
564            ))),
565        }
566    }
567}
568
569impl ImageProcessorTrait for CPUProcessor {
570    fn convert(
571        &mut self,
572        src: &TensorDyn,
573        dst: &mut TensorDyn,
574        rotation: Rotation,
575        flip: Flip,
576        crop: Crop,
577    ) -> Result<()> {
578        let crop = crop.resolve(
579            src.width().unwrap_or(0),
580            src.height().unwrap_or(0),
581            dst.width().unwrap_or(0),
582            dst.height().unwrap_or(0),
583        )?;
584        self.convert_impl(src, dst, rotation, flip, crop)
585    }
586
587    fn draw_decoded_masks(
588        &mut self,
589        dst: &mut TensorDyn,
590        detect: &[DetectBox],
591        segmentation: &[Segmentation],
592        overlay: crate::MaskOverlay<'_>,
593    ) -> Result<()> {
594        // CPU is the terminal fallback — it must always produce the full
595        // output, never assume the caller cleared dst. Every call writes
596        // the base layer first (bg copy or zero fill) and then the masks.
597        prepare_dst_base_cpu(dst, overlay.background)?;
598        let dst = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
599        self.draw_decoded_masks_impl(
600            dst,
601            detect,
602            segmentation,
603            overlay.opacity,
604            overlay.color_mode,
605        )
606    }
607
608    fn draw_proto_masks(
609        &mut self,
610        dst: &mut TensorDyn,
611        detect: &[DetectBox],
612        proto_data: &ProtoData,
613        overlay: crate::MaskOverlay<'_>,
614    ) -> Result<()> {
615        prepare_dst_base_cpu(dst, overlay.background)?;
616        let dst = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
617        self.draw_proto_masks_impl(
618            dst,
619            detect,
620            proto_data,
621            overlay.opacity,
622            overlay.letterbox,
623            overlay.color_mode,
624        )
625    }
626
627    fn set_class_colors(&mut self, colors: &[[u8; 4]]) -> Result<()> {
628        for (c, new_c) in self.colors.iter_mut().zip(colors.iter()) {
629            *c = *new_c;
630        }
631        Ok(())
632    }
633}
634
635// Internal methods — dtype-aware dispatch layer.
636impl CPUProcessor {
637    /// Top-level conversion dispatcher: handles dtype combinations.
638    pub(crate) fn convert_impl(
639        &mut self,
640        src: &TensorDyn,
641        dst: &mut TensorDyn,
642        rotation: Rotation,
643        flip: Flip,
644        crop: ResolvedCrop,
645    ) -> Result<()> {
646        let src_fmt = src.format().ok_or(Error::NotAnImage)?;
647        let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
648
649        // Resolve per-tensor colorimetry once, at the use site, without
650        // mutating either tensor. YUV→RGB conversions take their matrix/range
651        // from the source; RGB→YUV conversions from the destination. The grey/
652        // luma expansion is gated on the resolved range of the relevant side.
653        let src_cm = crate::colorimetry::effective_colorimetry(src);
654        let dst_cm = crate::colorimetry::effective_colorimetry(dst);
655        let src_full = src_cm.range == Some(edgefirst_tensor::ColorRange::Full);
656        let dst_full = dst_cm.range == Some(edgefirst_tensor::ColorRange::Full);
657        let src_params = ColorParams {
658            matrix: crate::colorimetry::yuv_matrix(src_cm.encoding.unwrap()),
659            range: crate::colorimetry::yuv_range(src_cm.range.unwrap()),
660            encoding: src_cm.encoding.unwrap(),
661            range_kind: src_cm.range.unwrap(),
662            src_full_range: src_full,
663            dst_full_range: dst_full,
664        };
665        let dst_params = ColorParams {
666            matrix: crate::colorimetry::yuv_matrix(dst_cm.encoding.unwrap()),
667            range: crate::colorimetry::yuv_range(dst_cm.range.unwrap()),
668            encoding: dst_cm.encoding.unwrap(),
669            range_kind: dst_cm.range.unwrap(),
670            src_full_range: src_full,
671            dst_full_range: dst_full,
672        };
673        match (src.dtype(), dst.dtype()) {
674            (DType::U8, DType::U8) => {
675                let src = src.as_u8().unwrap();
676                let dst = dst.as_u8_mut().unwrap();
677                self.convert_u8(
678                    src, dst, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
679                )
680            }
681            (DType::U8, DType::I8) => {
682                // Int8 output: reinterpret the i8 destination as u8 (layout-
683                // identical), convert directly into it, then XOR 0x80 in-place.
684                let src_u8 = src.as_u8().unwrap();
685                let dst_i8 = dst.as_i8_mut().unwrap();
686                // SAFETY: Tensor<i8> and Tensor<u8> are layout-identical
687                // (same element size, no T-dependent drop glue). Same
688                // rationale as gl::processor::tensor_i8_as_u8_mut.
689                let dst_u8 = unsafe { &mut *(dst_i8 as *mut Tensor<i8> as *mut Tensor<u8>) };
690                self.convert_u8(
691                    src_u8, dst_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
692                )?;
693                // Apply XOR 0x80 bias in-place (u8 → i8 conversion)
694                let mut map = dst_u8.map_mut()?;
695                apply_int8_xor_bias(map.as_mut_slice(), dst_fmt);
696                Ok(())
697            }
698            (DType::U8, d @ (DType::F32 | DType::F16)) => {
699                let src_u8 = src.as_u8().unwrap();
700                let dw = dst.width().ok_or(Error::NotAnImage)?;
701                let dh = dst.height().ok_or(Error::NotAnImage)?;
702                // Reuse the scratch tensor when format and dimensions match;
703                // otherwise reallocate and cache the new scratch.  Take the
704                // scratch out of `self` so that `convert_u8` can borrow `self`
705                // exclusively, then restore it afterwards.
706                let scratch_matches = self.widen_scratch.as_ref().is_some_and(|t| {
707                    t.width() == Some(dw) && t.height() == Some(dh) && t.format() == Some(dst_fmt)
708                });
709                let mut tmp = if scratch_matches {
710                    self.widen_scratch.take().unwrap()
711                } else {
712                    TensorDyn::image(
713                        dw,
714                        dh,
715                        dst_fmt,
716                        DType::U8,
717                        Some(TensorMemory::Mem),
718                        edgefirst_tensor::CpuAccess::ReadWrite,
719                    )?
720                };
721                {
722                    let tmp_u8 = tmp.as_u8_mut().unwrap();
723                    self.convert_u8(
724                        src_u8, tmp_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params,
725                        dst_params,
726                    )?;
727                }
728                // Widen the u8 scratch into the float destination, then restore
729                // the scratch for reuse on the next call.
730                {
731                    let tmp_u8 = tmp.as_u8().unwrap();
732                    let src_map = tmp_u8.map_read()?;
733                    match d {
734                        DType::F32 => {
735                            let dst_t = dst.as_f32_mut().unwrap();
736                            let mut dst_map = dst_t.map_mut()?;
737                            debug_assert_eq!(src_map.as_slice().len(), dst_map.as_slice().len());
738                            // NEON-accelerated u8→f32 `/255` widen (bit-identical
739                            // to the scalar `b as f32 / 255.0`); the scalar
740                            // iterator form did not vectorise. See cpu::simd.
741                            simd::widen_u8_to_f32_norm(src_map.as_slice(), dst_map.as_mut_slice());
742                        }
743                        DType::F16 => {
744                            let dst_t = dst.as_f16_mut().unwrap();
745                            let mut dst_map = dst_t.map_mut()?;
746                            debug_assert_eq!(src_map.as_slice().len(), dst_map.as_slice().len());
747                            // u8→f16 `/255` widen; uses native FP16
748                            // (`ucvtf`+`fdiv`) at runtime on FEAT_FP16 CPUs
749                            // (Orin), scalar `half::f16::from_f32` elsewhere.
750                            // See cpu::simd.
751                            simd::widen_u8_to_f16_norm(src_map.as_slice(), dst_map.as_mut_slice());
752                        }
753                        _ => unreachable!(),
754                    }
755                }
756                self.widen_scratch = Some(tmp);
757                Ok(())
758            }
759            (s, d) => Err(Error::NotSupported(format!("dtype {s} -> {d}",))),
760        }
761    }
762
763    /// Reuse `cached` if it already has dimensions `(w, h)` and pixel format
764    /// `fmt`, otherwise allocate a fresh `Mem` image. The returned buffer's
765    /// contents are **not** zeroed on reuse — callers must fully overwrite the
766    /// region they later read (see the note at the convert pipeline's tmp/tmp2
767    /// site). This amortises the per-frame `Tensor::image` `alloc_zeroed`.
768    fn reuse_or_alloc_image(
769        cached: Option<Tensor<u8>>,
770        w: usize,
771        h: usize,
772        fmt: PixelFormat,
773    ) -> Result<Tensor<u8>> {
774        if let Some(t) = cached {
775            if t.width() == Some(w) && t.height() == Some(h) && t.format() == Some(fmt) {
776                return Ok(t);
777            }
778        }
779        Ok(Tensor::<u8>::image(
780            w,
781            h,
782            fmt,
783            Some(TensorMemory::Mem),
784            edgefirst_tensor::CpuAccess::ReadWrite,
785        )?)
786    }
787
788    /// U8-to-U8 conversion: the full format conversion + resize pipeline.
789    #[allow(clippy::too_many_arguments)]
790    fn convert_u8(
791        &mut self,
792        src: &Tensor<u8>,
793        dst: &mut Tensor<u8>,
794        src_fmt: PixelFormat,
795        dst_fmt: PixelFormat,
796        rotation: Rotation,
797        flip: Flip,
798        crop: ResolvedCrop,
799        src_params: ColorParams,
800        dst_params: ColorParams,
801    ) -> Result<()> {
802        use PixelFormat::*;
803
804        let src_w = src.width().unwrap();
805        let src_h = src.height().unwrap();
806        let dst_w = dst.width().unwrap();
807        let dst_h = dst.height().unwrap();
808
809        crop.check_crop_dims(src_w, src_h, dst_w, dst_h)?;
810
811        // Determine intermediate format for the resize step
812        let intermediate = match (src_fmt, dst_fmt) {
813            (Nv12, Rgb) => Rgb,
814            (Nv12, Rgba) => Rgba,
815            (Nv12, Grey) => Grey,
816            (Nv12, Yuyv) => Rgba,
817            (Nv12, Nv16) => Rgba,
818            (Nv12, PlanarRgb) => Rgb,
819            (Nv12, PlanarRgba) => Rgba,
820            (Nv16, PlanarRgb) => Rgb,
821            (Nv16, PlanarRgba) => Rgba,
822            (Nv24, PlanarRgb) => Rgb,
823            (Nv24, PlanarRgba) => Rgba,
824            (Yuyv, Rgb) => Rgb,
825            (Yuyv, Rgba) => Rgba,
826            (Yuyv, Grey) => Grey,
827            (Yuyv, Yuyv) => Rgba,
828            (Yuyv, PlanarRgb) => Rgb,
829            (Yuyv, PlanarRgba) => Rgba,
830            (Yuyv, Nv16) => Rgba,
831            (Vyuy, Rgb) => Rgb,
832            (Vyuy, Rgba) => Rgba,
833            (Vyuy, Grey) => Grey,
834            (Vyuy, Vyuy) => Rgba,
835            (Vyuy, PlanarRgb) => Rgb,
836            (Vyuy, PlanarRgba) => Rgba,
837            (Vyuy, Nv16) => Rgba,
838            (Rgba, Rgb) => Rgba,
839            (Rgba, Rgba) => Rgba,
840            (Rgba, Grey) => Grey,
841            (Rgba, Yuyv) => Rgba,
842            (Rgba, PlanarRgb) => Rgba,
843            (Rgba, PlanarRgba) => Rgba,
844            (Rgba, Nv16) => Rgba,
845            (Rgb, Rgb) => Rgb,
846            (Rgb, Rgba) => Rgb,
847            (Rgb, Grey) => Grey,
848            (Rgb, Yuyv) => Rgb,
849            (Rgb, PlanarRgb) => Rgb,
850            (Rgb, PlanarRgba) => Rgb,
851            (Rgb, Nv16) => Rgb,
852            (Grey, Rgb) => Rgb,
853            (Grey, Rgba) => Rgba,
854            (Grey, Grey) => Grey,
855            (Grey, Yuyv) => Grey,
856            (Grey, PlanarRgb) => Grey,
857            (Grey, PlanarRgba) => Grey,
858            (Grey, Nv16) => Grey,
859            (Nv12, Bgra) => Rgba,
860            (Yuyv, Bgra) => Rgba,
861            (Vyuy, Bgra) => Rgba,
862            (Rgba, Bgra) => Rgba,
863            (Rgb, Bgra) => Rgb,
864            (Grey, Bgra) => Grey,
865            (Bgra, Bgra) => Bgra,
866            (Nv16, Rgb) => Rgb,
867            (Nv16, Rgba) => Rgba,
868            (Nv16, Bgra) => Rgba,
869            (Nv24, Rgb) => Rgb,
870            (Nv24, Rgba) => Rgba,
871            (Nv24, Grey) => Grey,
872            (Nv24, Bgra) => Rgba,
873            (PlanarRgb, Rgb) => Rgb,
874            (PlanarRgb, Rgba) => Rgb,
875            (PlanarRgb, Bgra) => Rgb,
876            (PlanarRgba, Rgb) => Rgba,
877            (PlanarRgba, Rgba) => Rgba,
878            (PlanarRgba, Bgra) => Rgba,
879            (s, d) => {
880                return Err(Error::NotSupported(format!("Conversion from {s} to {d}",)));
881            }
882        };
883
884        let need_resize_flip_rotation = rotation != Rotation::None
885            || flip != Flip::None
886            || src_w != dst_w
887            || src_h != dst_h
888            || crop.src_rect.is_some_and(|c| {
889                c != Rect {
890                    left: 0,
891                    top: 0,
892                    width: src_w,
893                    height: src_h,
894                }
895            })
896            || crop.dst_rect.is_some_and(|c| {
897                c != Rect {
898                    left: 0,
899                    top: 0,
900                    width: dst_w,
901                    height: dst_h,
902                }
903            });
904
905        // Pick the resolved colorimetry for a single conversion by its YUV
906        // side: YUV→RGB decodes use the source side, RGB→YUV encodes the dest.
907        let direct_is_yuv_src = matches!(src_fmt, Nv12 | Nv16 | Nv24 | Yuyv | Vyuy);
908        let direct_params = if direct_is_yuv_src {
909            src_params
910        } else {
911            dst_params
912        };
913
914        // Fused NV→planar (no resize/flip/rotation/crop): decode the YUV source
915        // into packed RGB one cache-resident row strip at a time and
916        // NEON-deinterleave each strip straight into the destination planes, so
917        // the full-size packed-RGB intermediate never round-trips through DRAM
918        // and is not reallocated per frame. JPEG decodes to the NV family and the
919        // model wants planar RGB, so this is the hot Orin CPU-preprocess path.
920        if !need_resize_flip_rotation
921            && matches!(src_fmt, Nv12 | Nv16 | Nv24)
922            && matches!(dst_fmt, PlanarRgb | PlanarRgba)
923        {
924            return self.convert_nv_to_planar_fused(src, dst, src_fmt, dst_fmt, direct_params);
925        }
926
927        // check if a direct conversion can be done
928        if !need_resize_flip_rotation && Self::support_conversion_pf(src_fmt, dst_fmt) {
929            return Self::convert_format_pf(src, dst, src_fmt, dst_fmt, direct_params);
930        }
931
932        // any extra checks
933        if dst_fmt == Yuyv && !dst_w.is_multiple_of(2) {
934            return Err(Error::NotSupported(format!(
935                "{} destination must have width divisible by 2",
936                dst_fmt,
937            )));
938        }
939
940        // Take the cached intermediates out of `self` so the resize step can
941        // borrow `self` exclusively; they are restored before returning on the
942        // success path. Reused buffers are not re-zeroed — every consumer below
943        // fully overwrites the region it later reads (the pre-resize convert
944        // writes all of `tmp`; the resize writes the scaled rect of `tmp2` and
945        // its letterbox border is either pre-filled from `dst` or overwritten in
946        // `dst` by the final `fill_image_outside_crop_u8`).
947        let mut cached_tmp = self.convert_tmp.take();
948        let mut cached_tmp2 = self.convert_tmp2.take();
949
950        // create tmp buffer (reusing the cached one when its geometry matches)
951        let tmp_holder: Option<Tensor<u8>> = if intermediate != src_fmt {
952            let _s = tracing::trace_span!(
953                "image.convert.cpu.format_convert",
954                from = ?src_fmt,
955                to = ?intermediate,
956                pass = "pre_resize",
957            )
958            .entered();
959            let mut t = Self::reuse_or_alloc_image(cached_tmp.take(), src_w, src_h, intermediate)?;
960            Self::convert_format_pf(src, &mut t, src_fmt, intermediate, src_params)?;
961            Some(t)
962        } else {
963            None
964        };
965        let (tmp, tmp_fmt): (&Tensor<u8>, PixelFormat) = match &tmp_holder {
966            Some(t) => (t, intermediate),
967            None => (src, src_fmt),
968        };
969
970        // format must be RGB/RGBA/GREY
971        debug_assert!(matches!(tmp_fmt, Rgb | Rgba | Grey));
972        if tmp_fmt == dst_fmt {
973            let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
974            self.resize_flip_rotate_pf(tmp, dst, dst_fmt, rotation, flip, crop)?;
975        } else if !need_resize_flip_rotation {
976            let _s = tracing::trace_span!(
977                "image.convert.cpu.format_convert",
978                from = ?tmp_fmt,
979                to = ?dst_fmt,
980                pass = "direct",
981            )
982            .entered();
983            Self::convert_format_pf(tmp, dst, tmp_fmt, dst_fmt, dst_params)?;
984        } else {
985            let mut tmp2 = Self::reuse_or_alloc_image(cached_tmp2.take(), dst_w, dst_h, tmp_fmt)?;
986            if crop.dst_rect.is_some_and(|c| {
987                c != Rect {
988                    left: 0,
989                    top: 0,
990                    width: dst_w,
991                    height: dst_h,
992                }
993            }) && crop.dst_color.is_none()
994            {
995                Self::convert_format_pf(dst, &mut tmp2, dst_fmt, tmp_fmt, dst_params)?;
996            }
997            {
998                let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
999                self.resize_flip_rotate_pf(tmp, &mut tmp2, tmp_fmt, rotation, flip, crop)?;
1000            }
1001            {
1002                let _s = tracing::trace_span!(
1003                    "image.convert.cpu.format_convert",
1004                    from = ?tmp_fmt,
1005                    to = ?dst_fmt,
1006                    pass = "post_resize",
1007                )
1008                .entered();
1009                Self::convert_format_pf(&tmp2, dst, tmp_fmt, dst_fmt, dst_params)?;
1010            }
1011            cached_tmp2 = Some(tmp2);
1012        }
1013        // Restore the intermediates to the cache for the next call (`tmp` — a
1014        // borrow of `tmp_holder` — is no longer used past this point).
1015        if let Some(t) = tmp_holder {
1016            cached_tmp = Some(t);
1017        }
1018        self.convert_tmp = cached_tmp;
1019        self.convert_tmp2 = cached_tmp2;
1020
1021        if let (Some(dst_rect), Some(dst_color)) = (crop.dst_rect, crop.dst_color) {
1022            let full_rect = Rect {
1023                left: 0,
1024                top: 0,
1025                width: dst_w,
1026                height: dst_h,
1027            };
1028            if dst_rect != full_rect {
1029                Self::fill_image_outside_crop_u8(dst, dst_color, dst_rect)?;
1030            }
1031        }
1032
1033        Ok(())
1034    }
1035
1036    fn draw_decoded_masks_impl(
1037        &mut self,
1038        dst: &mut Tensor<u8>,
1039        detect: &[DetectBox],
1040        segmentation: &[Segmentation],
1041        opacity: f32,
1042        color_mode: crate::ColorMode,
1043    ) -> Result<()> {
1044        let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
1045        if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
1046            return Err(crate::Error::NotSupported(
1047                "CPU image rendering only supports RGBA or RGB images".to_string(),
1048            ));
1049        }
1050
1051        let _timer = FunctionTimer::new("CPUProcessor::draw_decoded_masks");
1052
1053        let dst_w = dst.width().unwrap();
1054        let dst_h = dst.height().unwrap();
1055        let dst_rs = tensor_row_stride(dst);
1056        let dst_c = dst_fmt.channels();
1057
1058        let mut map = dst.map_mut()?;
1059        let dst_slice = map.as_mut_slice();
1060
1061        self.render_box(dst_w, dst_h, dst_rs, dst_c, dst_slice, detect, color_mode)?;
1062
1063        if segmentation.is_empty() {
1064            return Ok(());
1065        }
1066
1067        // Semantic segmentation (e.g. ModelPack) has C > 1 (multi-class),
1068        // instance segmentation (e.g. YOLO) has C = 1 (binary per-instance).
1069        let is_semantic = segmentation[0].segmentation.shape()[2] > 1;
1070
1071        if is_semantic {
1072            self.render_modelpack_segmentation(
1073                dst_w,
1074                dst_h,
1075                dst_rs,
1076                dst_c,
1077                dst_slice,
1078                &segmentation[0],
1079                opacity,
1080            )?;
1081        } else {
1082            for (idx, (seg, det)) in segmentation.iter().zip(detect).enumerate() {
1083                let color_index = color_mode.index(idx, det.label);
1084                self.render_yolo_segmentation(
1085                    dst_w,
1086                    dst_h,
1087                    dst_rs,
1088                    dst_c,
1089                    dst_slice,
1090                    seg,
1091                    color_index,
1092                    opacity,
1093                )?;
1094            }
1095        }
1096
1097        Ok(())
1098    }
1099
1100    fn draw_proto_masks_impl(
1101        &mut self,
1102        dst: &mut Tensor<u8>,
1103        detect: &[DetectBox],
1104        proto_data: &ProtoData,
1105        opacity: f32,
1106        letterbox: Option<[f32; 4]>,
1107        color_mode: crate::ColorMode,
1108    ) -> Result<()> {
1109        let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
1110        if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
1111            return Err(crate::Error::NotSupported(
1112                "CPU image rendering only supports RGBA or RGB images".to_string(),
1113            ));
1114        }
1115
1116        let _timer = FunctionTimer::new("CPUProcessor::draw_proto_masks");
1117
1118        let dst_w = dst.width().unwrap();
1119        let dst_h = dst.height().unwrap();
1120        let dst_rs = tensor_row_stride(dst);
1121        let channels = dst_fmt.channels();
1122
1123        let mut map = dst.map_mut()?;
1124        let dst_slice = map.as_mut_slice();
1125
1126        self.render_box(
1127            dst_w, dst_h, dst_rs, channels, dst_slice, detect, color_mode,
1128        )?;
1129
1130        if detect.is_empty() {
1131            return Ok(());
1132        }
1133        let proto_shape = proto_data.protos.shape();
1134        if proto_shape.len() != 3 {
1135            return Err(Error::InvalidShape(format!(
1136                "protos tensor must be rank-3, got {proto_shape:?}"
1137            )));
1138        }
1139        let proto_h = proto_shape[0];
1140        let proto_w = proto_shape[1];
1141        let num_protos = proto_shape[2];
1142        let coeff_shape = proto_data.mask_coefficients.shape();
1143        if coeff_shape.len() != 2 {
1144            return Err(Error::InvalidShape(format!(
1145                "mask_coefficients tensor must be rank-2, got {coeff_shape:?}"
1146            )));
1147        }
1148        // Genuine "no detections this frame" → nothing to render.
1149        if coeff_shape[0] == 0 {
1150            return Ok(());
1151        }
1152        if coeff_shape[1] != num_protos {
1153            return Err(Error::InvalidShape(format!(
1154                "mask_coefficients second dimension must match num_protos \
1155                 ({num_protos}), got {coeff_shape:?}"
1156            )));
1157        }
1158
1159        // Widen coefficients to f32 once; shape [N, num_protos].
1160        let coeff_f32: Vec<f32> = match proto_data.mask_coefficients.dtype() {
1161            DType::F32 => {
1162                let t = proto_data.mask_coefficients.as_f32().expect("F32");
1163                let m = t.map_read()?;
1164                m.as_slice().to_vec()
1165            }
1166            DType::F16 => {
1167                let t = proto_data.mask_coefficients.as_f16().expect("F16");
1168                let m = t.map_read()?;
1169                m.as_slice().iter().map(|v| v.to_f32()).collect()
1170            }
1171            DType::I8 => {
1172                let t = proto_data.mask_coefficients.as_i8().expect("I8");
1173                let m = t.map_read()?;
1174                if let Some(q) = t.quantization() {
1175                    use edgefirst_tensor::QuantMode;
1176                    let (scale, zp) = match q.mode() {
1177                        QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1178                        QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1179                        other => {
1180                            return Err(Error::NotSupported(format!(
1181                                "I8 mask_coefficients quantization mode {other:?} not supported"
1182                            )));
1183                        }
1184                    };
1185                    m.as_slice()
1186                        .iter()
1187                        .map(|&v| (v as f32 - zp) * scale)
1188                        .collect()
1189                } else {
1190                    m.as_slice().iter().map(|&v| v as f32).collect()
1191                }
1192            }
1193            DType::I16 => {
1194                let t = proto_data.mask_coefficients.as_i16().expect("I16");
1195                let m = t.map_read()?;
1196                if let Some(q) = t.quantization() {
1197                    use edgefirst_tensor::QuantMode;
1198                    let (scale, zp) = match q.mode() {
1199                        QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1200                        QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1201                        other => {
1202                            return Err(Error::NotSupported(format!(
1203                                "I16 mask_coefficients quantization mode {other:?} not supported"
1204                            )));
1205                        }
1206                    };
1207                    m.as_slice()
1208                        .iter()
1209                        .map(|&v| (v as f32 - zp) * scale)
1210                        .collect()
1211                } else {
1212                    m.as_slice().iter().map(|&v| v as f32).collect()
1213                }
1214            }
1215            other => {
1216                return Err(Error::InvalidShape(format!(
1217                    "mask_coefficients dtype {other:?} not supported"
1218                )));
1219            }
1220        };
1221
1222        // Precompute letterbox scale/offset for output-pixel → proto-pixel mapping.
1223        let (lx0, lx_range, ly0, ly_range) = match letterbox {
1224            Some([lx0, ly0, lx1, ly1]) => (lx0, lx1 - lx0, ly0, ly1 - ly0),
1225            None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
1226        };
1227
1228        // Per-dtype dispatch. Map protos once, call the inner draw loop
1229        // with a dtype-specialized loader closure.
1230        match proto_data.protos.dtype() {
1231            DType::F32 => {
1232                let t = proto_data.protos.as_f32().expect("F32");
1233                let m = t.map_read()?;
1234                self.draw_proto_masks_inner(
1235                    dst_slice,
1236                    dst_w,
1237                    dst_h,
1238                    dst_rs,
1239                    channels,
1240                    detect,
1241                    m.as_slice(),
1242                    &coeff_f32,
1243                    proto_h,
1244                    proto_w,
1245                    num_protos,
1246                    opacity,
1247                    (lx0, lx_range, ly0, ly_range),
1248                    color_mode,
1249                    0.0_f32,
1250                    |p: &f32, _| *p,
1251                );
1252            }
1253            DType::F16 => {
1254                let t = proto_data.protos.as_f16().expect("F16");
1255                let m = t.map_read()?;
1256                self.draw_proto_masks_inner(
1257                    dst_slice,
1258                    dst_w,
1259                    dst_h,
1260                    dst_rs,
1261                    channels,
1262                    detect,
1263                    m.as_slice(),
1264                    &coeff_f32,
1265                    proto_h,
1266                    proto_w,
1267                    num_protos,
1268                    opacity,
1269                    (lx0, lx_range, ly0, ly_range),
1270                    color_mode,
1271                    0.0_f32,
1272                    |p: &half::f16, _| p.to_f32(),
1273                );
1274            }
1275            DType::I8 => {
1276                use edgefirst_tensor::QuantMode;
1277                let t = proto_data.protos.as_i8().expect("I8");
1278                let m = t.map_read()?;
1279                let quant = t.quantization().ok_or_else(|| {
1280                    Error::InvalidShape("I8 protos require quantization metadata".into())
1281                })?;
1282                let (scale, zp) = match quant.mode() {
1283                    QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1284                    QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1285                    QuantMode::PerChannel { axis, .. }
1286                    | QuantMode::PerChannelSymmetric { axis, .. } => {
1287                        return Err(Error::NotSupported(format!(
1288                            "per-channel quantization (axis={axis}) in draw_proto_masks \
1289                             CPU path not yet supported"
1290                        )));
1291                    }
1292                };
1293                self.draw_proto_masks_inner(
1294                    dst_slice,
1295                    dst_w,
1296                    dst_h,
1297                    dst_rs,
1298                    channels,
1299                    detect,
1300                    m.as_slice(),
1301                    &coeff_f32,
1302                    proto_h,
1303                    proto_w,
1304                    num_protos,
1305                    opacity,
1306                    (lx0, lx_range, ly0, ly_range),
1307                    color_mode,
1308                    scale,
1309                    move |p: &i8, _| (*p as f32) - zp,
1310                );
1311            }
1312            other => {
1313                return Err(Error::InvalidShape(format!(
1314                    "proto tensor dtype {other:?} not supported"
1315                )));
1316            }
1317        }
1318
1319        Ok(())
1320    }
1321
1322    #[allow(clippy::too_many_arguments)]
1323    fn draw_proto_masks_inner<P: Copy>(
1324        &self,
1325        dst_slice: &mut [u8],
1326        dst_w: usize,
1327        dst_h: usize,
1328        dst_rs: usize,
1329        channels: usize,
1330        detect: &[DetectBox],
1331        protos: &[P],
1332        coeff_all_f32: &[f32],
1333        proto_h: usize,
1334        proto_w: usize,
1335        num_protos: usize,
1336        opacity: f32,
1337        letterbox_xy: (f32, f32, f32, f32),
1338        color_mode: crate::ColorMode,
1339        acc_scale: f32,
1340        load_f32: impl Fn(&P, f32) -> f32 + Copy,
1341    ) {
1342        let (lx0, lx_range, ly0, ly_range) = letterbox_xy;
1343        let stride_y = proto_w * num_protos;
1344        for (idx, det) in detect.iter().enumerate() {
1345            let coeff = &coeff_all_f32[idx * num_protos..(idx + 1) * num_protos];
1346            let color_index = color_mode.index(idx, det.label);
1347            let color = self.colors[color_index % self.colors.len()];
1348            let alpha = if opacity == 1.0 {
1349                color[3] as u16
1350            } else {
1351                (color[3] as f32 * opacity).round() as u16
1352            };
1353
1354            let start_x = (dst_w as f32 * det.bbox.xmin).round() as usize;
1355            let start_y = (dst_h as f32 * det.bbox.ymin).round() as usize;
1356            let end_x = ((dst_w as f32 * det.bbox.xmax).round() as usize).min(dst_w);
1357            let end_y = ((dst_h as f32 * det.bbox.ymax).round() as usize).min(dst_h);
1358
1359            for y in start_y..end_y {
1360                for x in start_x..end_x {
1361                    let px = (lx0 + (x as f32 / dst_w as f32) * lx_range) * proto_w as f32 - 0.5;
1362                    let py = (ly0 + (y as f32 / dst_h as f32) * ly_range) * proto_h as f32 - 0.5;
1363
1364                    // Bilinear interpolation with per-load widening. Inline
1365                    // bilinear-sample since bilinear_dot_slice takes a
1366                    // different closure shape (no `zp` arg).
1367                    let x0 = (px.floor() as isize).clamp(0, proto_w as isize - 1) as usize;
1368                    let y0 = (py.floor() as isize).clamp(0, proto_h as isize - 1) as usize;
1369                    let x1 = (x0 + 1).min(proto_w - 1);
1370                    let y1 = (y0 + 1).min(proto_h - 1);
1371                    let fx = px - px.floor();
1372                    let fy = py - py.floor();
1373                    let w00 = (1.0 - fx) * (1.0 - fy);
1374                    let w10 = fx * (1.0 - fy);
1375                    let w01 = (1.0 - fx) * fy;
1376                    let w11 = fx * fy;
1377                    let b00 = y0 * stride_y + x0 * num_protos;
1378                    let b10 = y0 * stride_y + x1 * num_protos;
1379                    let b01 = y1 * stride_y + x0 * num_protos;
1380                    let b11 = y1 * stride_y + x1 * num_protos;
1381                    let mut acc = 0.0_f32;
1382                    for p in 0..num_protos {
1383                        let v00 = load_f32(&protos[b00 + p], 0.0);
1384                        let v10 = load_f32(&protos[b10 + p], 0.0);
1385                        let v01 = load_f32(&protos[b01 + p], 0.0);
1386                        let v11 = load_f32(&protos[b11 + p], 0.0);
1387                        let val = w00 * v00 + w10 * v10 + w01 * v01 + w11 * v11;
1388                        acc += coeff[p] * val;
1389                    }
1390                    let final_acc = if acc_scale == 0.0 {
1391                        acc
1392                    } else {
1393                        acc_scale * acc
1394                    };
1395                    // Pass-through: acc_scale=0.0 means "no scaling" (f32/f16
1396                    // native); non-zero means "apply scale once" (i8 with
1397                    // per-tensor quant).
1398                    let mask = 1.0 / (1.0 + (-final_acc).exp());
1399                    if mask < 0.5 {
1400                        continue;
1401                    }
1402                    let dst_index = y * dst_rs + x * channels;
1403                    for c in 0..3 {
1404                        dst_slice[dst_index + c] = ((color[c] as u16 * alpha
1405                            + dst_slice[dst_index + c] as u16 * (255 - alpha))
1406                            / 255) as u8;
1407                    }
1408                }
1409            }
1410        }
1411    }
1412}