Skip to main content

edgefirst_decoder/
lib.rs

1// SPDX-FileCopyrightText: Copyright 2025 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4/*!
5## EdgeFirst HAL - Decoders
6This crate provides decoding utilities for YOLOobject detection and segmentation models, and ModelPack detection and segmentation models.
7It supports both floating-point and quantized model outputs, allowing for efficient processing on edge devices. The crate includes functions
8for efficient post-processing model outputs into usable detection boxes and segmentation masks, as well as utilities for dequantizing model outputs..
9
10For general usage, use the `Decoder` struct which provides functions for decoding various model outputs based on the model configuration.
11If you already know the model type and output formats, you can use the lower-level functions directly from the `yolo` and `modelpack` modules.
12
13
14### Quick Example
15```rust,no_run
16use edgefirst_decoder::{DecoderBuilder, DecoderResult, configs::{self, DecoderVersion}};
17use edgefirst_tensor::TensorDyn;
18
19fn main() -> DecoderResult<()> {
20    // Create a decoder for a YOLOv8 model with quantized int8 output
21    let decoder = DecoderBuilder::new()
22        .with_config_yolo_det(configs::Detection {
23            anchors: None,
24            decoder: configs::DecoderType::Ultralytics,
25            quantization: Some(configs::QuantTuple(0.012345, 26)),
26            shape: vec![1, 84, 8400],
27            dshape: Vec::new(),
28            normalized: Some(true),
29        },
30        Some(DecoderVersion::Yolov8))
31        .with_score_threshold(0.25)
32        .with_iou_threshold(0.7)
33        .build()?;
34
35    // Get the model output tensors from inference
36    let model_output: Vec<TensorDyn> = vec![/* tensors from inference */];
37    let tensor_refs: Vec<&TensorDyn> = model_output.iter().collect();
38
39    let mut output_boxes = Vec::with_capacity(10);
40    let mut output_masks = Vec::with_capacity(10);
41
42    // Decode model output into detection boxes and segmentation masks
43    decoder.decode(&tensor_refs, &mut output_boxes, &mut output_masks)?;
44    Ok(())
45}
46```
47
48# Overview
49
50The primary components of this crate are:
51- `Decoder`/`DecoderBuilder` struct: Provides high-level functions to decode model outputs based on the model configuration.
52- `yolo` module: Contains functions specific to decoding YOLO model outputs.
53- `modelpack` module: Contains functions specific to decoding ModelPack model outputs.
54
55The `Decoder` supports both floating-point and quantized model outputs, allowing for efficient processing on edge devices.
56It also supports mixed integer types for quantized outputs, such as when one output tensor is int8 and another is uint8.
57When decoding quantized outputs, the appropriate quantization parameters must be provided for each output tensor.
58If the integer types used in the model output is not supported by the decoder, the user can manually dequantize the model outputs using
59the `dequantize` functions provided in this crate, and then use the floating-point decoding functions. However, it is recommended
60to not dequantize the model outputs manually before passing them to the decoder, as the quantized decoder functions are optimized for performance.
61
62The `yolo` and `modelpack` modules provide lower-level functions for decoding model outputs directly,
63which can be used if the model type and output formats are known in advance.
64
65
66*/
67#![cfg_attr(coverage_nightly, feature(coverage_attribute))]
68
69use ndarray::{Array, Array2, ArrayView, ArrayView1, ArrayView3, Dimension};
70use num_traits::{AsPrimitive, Float, PrimInt};
71
72pub mod byte;
73pub mod error;
74pub mod float;
75pub mod infer;
76pub mod modelpack;
77pub mod per_scale;
78pub mod schema;
79pub mod tiling;
80pub mod yolo;
81
82pub use infer::{
83    infer_ultralytics_schema, InferError, InferredSchema, ModelSignals, ModelSource, TensorInfo,
84};
85
86mod decoder;
87pub use decoder::*;
88#[cfg(test)]
89mod test_support;
90
91// The detection vocabulary lives in edgefirst-tensor so edgefirst-image can
92// render from it without depending on this crate. Re-exported here so the
93// decoder's own modules keep resolving `crate::Segmentation`.
94pub use edgefirst_tensor::{BoundingBox, DetectBox, ProtoData, ProtoLayout, Segmentation};
95
96// TilePlacement is the contract between edgefirst-image (producer) and this
97// crate (consumer); it lives in edgefirst-decoder-abi so image can name it
98// without linking decoder. Re-exported here so the decoder's own modules
99// keep resolving `crate::tiling::TilePlacement`.
100pub use edgefirst_decoder_abi::TilePlacement;
101
102pub use configs::{DecoderVersion, Nms};
103pub use error::{DecoderError, DecoderResult};
104pub use per_scale::DecodeDtype;
105
106use crate::{
107    decoder::configs::QuantTuple, modelpack::modelpack_segmentation_to_mask,
108    yolo::yolo_segmentation_to_mask,
109};
110
111/// Materialise an ndarray mask into the [`TensorDyn`] that
112/// [`Segmentation`] carries.
113///
114/// The decoder computes mask values with ndarray and converts once, here at
115/// the boundary. `Segmentation` holds a `TensorDyn` rather than an
116/// `Array3<u8>` so that `edgefirst-tensor` can own the type without taking an
117/// ndarray dependency, which is what lets `edgefirst-image` render masks
118/// without depending on this crate.
119fn mask_to_tensor(mask: ndarray::ArrayView3<'_, u8>) -> DecoderResult<edgefirst_tensor::TensorDyn> {
120    edgefirst_tensor::Tensor::from_arrayview3(mask)
121        .map(edgefirst_tensor::TensorDyn::from)
122        .map_err(|e| DecoderError::Internal(format!("mask -> TensorDyn: {e}")))
123}
124
125/// Float-epsilon comparison for `DetectBox`, as a test-only extension trait.
126///
127/// This is a TEST ASSERTION helper -- all 98 call sites are in this crate's
128/// test modules and there are none in production code -- so it does not belong
129/// on `edgefirst-tensor`'s public API, where the type move briefly put it. An
130/// extension trait keeps every call site written as
131/// `a.equal_within_delta(b, eps)` while confining the method to tests.
132#[cfg(test)]
133pub(crate) trait DetectBoxTestExt {
134    fn equal_within_delta(&self, rhs: &DetectBox, eps: f32) -> bool;
135}
136
137#[cfg(test)]
138impl DetectBoxTestExt for DetectBox {
139    fn equal_within_delta(&self, rhs: &DetectBox, eps: f32) -> bool {
140        let eq_delta = |a: f32, b: f32| (a - b).abs() <= eps;
141        self.label == rhs.label
142            && eq_delta(self.score, rhs.score)
143            && eq_delta(self.bbox.xmin, rhs.bbox.xmin)
144            && eq_delta(self.bbox.ymin, rhs.bbox.ymin)
145            && eq_delta(self.bbox.xmax, rhs.bbox.xmax)
146            && eq_delta(self.bbox.ymax, rhs.bbox.ymax)
147    }
148}
149
150/// Trait to convert bounding box formats to XYXY float format
151pub trait BBoxTypeTrait {
152    /// Converts the bbox into XYXY float format.
153    fn to_xyxy_float<A: Float + 'static, B: AsPrimitive<A>>(input: &[B; 4]) -> [A; 4];
154
155    /// Converts the bbox into XYXY float format.
156    fn to_xyxy_dequant<A: Float + 'static, B: AsPrimitive<A>>(
157        input: &[B; 4],
158        quant: Quantization,
159    ) -> [A; 4]
160    where
161        f32: AsPrimitive<A>,
162        i32: AsPrimitive<A>;
163
164    /// Converts the bbox into XYXY float format.
165    ///
166    /// # Examples
167    /// ```rust
168    /// # use edgefirst_decoder::{BBoxTypeTrait, XYWH};
169    /// # use ndarray::array;
170    /// let arr = array![10.0_f32, 20.0, 20.0, 20.0];
171    /// let xyxy: [f32; 4] = XYWH::ndarray_to_xyxy_float(arr.view());
172    /// assert_eq!(xyxy, [0.0_f32, 10.0, 20.0, 30.0]);
173    /// ```
174    #[inline(always)]
175    fn ndarray_to_xyxy_float<A: Float + 'static, B: AsPrimitive<A>>(
176        input: ArrayView1<B>,
177    ) -> [A; 4] {
178        Self::to_xyxy_float(&[input[0], input[1], input[2], input[3]])
179    }
180
181    #[inline(always)]
182    /// Converts the bbox into XYXY float format.
183    fn ndarray_to_xyxy_dequant<A: Float + 'static, B: AsPrimitive<A>>(
184        input: ArrayView1<B>,
185        quant: Quantization,
186    ) -> [A; 4]
187    where
188        f32: AsPrimitive<A>,
189        i32: AsPrimitive<A>,
190    {
191        Self::to_xyxy_dequant(&[input[0], input[1], input[2], input[3]], quant)
192    }
193}
194
195/// Converts XYXY bounding boxes to XYXY
196#[derive(Debug, Clone, Copy, PartialEq, Eq)]
197pub struct XYXY {}
198
199impl BBoxTypeTrait for XYXY {
200    fn to_xyxy_float<A: Float + 'static, B: AsPrimitive<A>>(input: &[B; 4]) -> [A; 4] {
201        input.map(|b| b.as_())
202    }
203
204    fn to_xyxy_dequant<A: Float + 'static, B: AsPrimitive<A>>(
205        input: &[B; 4],
206        quant: Quantization,
207    ) -> [A; 4]
208    where
209        f32: AsPrimitive<A>,
210        i32: AsPrimitive<A>,
211    {
212        let scale = quant.scale.as_();
213        let zp = quant.zero_point.as_();
214        input.map(|b| (b.as_() - zp) * scale)
215    }
216}
217
218/// Converts XYWH bounding boxes to XYXY. The XY values are the center of the
219/// box
220#[derive(Debug, Clone, Copy, PartialEq, Eq)]
221pub struct XYWH {}
222
223impl BBoxTypeTrait for XYWH {
224    #[inline(always)]
225    fn to_xyxy_float<A: Float + 'static, B: AsPrimitive<A>>(input: &[B; 4]) -> [A; 4] {
226        let half = A::one() / (A::one() + A::one());
227        [
228            (input[0].as_()) - (input[2].as_() * half),
229            (input[1].as_()) - (input[3].as_() * half),
230            (input[0].as_()) + (input[2].as_() * half),
231            (input[1].as_()) + (input[3].as_() * half),
232        ]
233    }
234
235    #[inline(always)]
236    fn to_xyxy_dequant<A: Float + 'static, B: AsPrimitive<A>>(
237        input: &[B; 4],
238        quant: Quantization,
239    ) -> [A; 4]
240    where
241        f32: AsPrimitive<A>,
242        i32: AsPrimitive<A>,
243    {
244        let scale = quant.scale.as_();
245        let half_scale = (quant.scale * 0.5).as_();
246        let zp = quant.zero_point.as_();
247        let [x, y, w, h] = [
248            (input[0].as_() - zp) * scale,
249            (input[1].as_() - zp) * scale,
250            (input[2].as_() - zp) * half_scale,
251            (input[3].as_() - zp) * half_scale,
252        ];
253
254        [x - w, y - h, x + w, y + h]
255    }
256}
257
258/// Describes the quantization parameters for a tensor
259#[derive(Debug, Clone, Copy, PartialEq)]
260pub struct Quantization {
261    pub scale: f32,
262    pub zero_point: i32,
263}
264
265impl Quantization {
266    /// Creates a new Quantization struct
267    /// # Examples
268    /// ```
269    /// # use edgefirst_decoder::Quantization;
270    /// let quant = Quantization::new(0.1, -128);
271    /// assert_eq!(quant.scale, 0.1);
272    /// assert_eq!(quant.zero_point, -128);
273    /// ```
274    pub fn new(scale: f32, zero_point: i32) -> Self {
275        Self { scale, zero_point }
276    }
277
278    /// Returns a quantization that's a no-op: scale=1.0, zero_point=0.
279    ///
280    /// Used by the per-scale pipeline as a sentinel for float-to-float
281    /// passthrough cells where no affine transform is required.
282    /// # Examples
283    /// ```
284    /// # use edgefirst_decoder::Quantization;
285    /// let quant = Quantization::identity();
286    /// assert_eq!(quant.scale, 1.0);
287    /// assert_eq!(quant.zero_point, 0);
288    /// ```
289    pub fn identity() -> Self {
290        Self {
291            scale: 1.0,
292            zero_point: 0,
293        }
294    }
295}
296
297impl From<QuantTuple> for Quantization {
298    /// Creates a new Quantization struct from a QuantTuple
299    /// # Examples
300    /// ```
301    /// # use edgefirst_decoder::Quantization;
302    /// # use edgefirst_decoder::configs::QuantTuple;
303    /// let quant_tuple = QuantTuple(0.1_f32, -128_i32);
304    /// let quant = Quantization::from(quant_tuple);
305    /// assert_eq!(quant.scale, 0.1);
306    /// assert_eq!(quant.zero_point, -128);
307    /// ```
308    fn from(quant_tuple: QuantTuple) -> Quantization {
309        Quantization {
310            scale: quant_tuple.0,
311            zero_point: quant_tuple.1,
312        }
313    }
314}
315
316impl<S, Z> From<(S, Z)> for Quantization
317where
318    S: AsPrimitive<f32>,
319    Z: AsPrimitive<i32>,
320{
321    /// Creates a new Quantization struct from a tuple
322    /// # Examples
323    /// ```
324    /// # use edgefirst_decoder::Quantization;
325    /// let quant = Quantization::from((0.1_f64, -128_i64));
326    /// assert_eq!(quant.scale, 0.1);
327    /// assert_eq!(quant.zero_point, -128);
328    /// ```
329    fn from((scale, zp): (S, Z)) -> Quantization {
330        Self {
331            scale: scale.as_(),
332            zero_point: zp.as_(),
333        }
334    }
335}
336
337impl Default for Quantization {
338    /// Creates a default Quantization struct with scale 1.0 and zero_point 0
339    /// # Examples
340    /// ```rust
341    /// # use edgefirst_decoder::Quantization;
342    /// let quant = Quantization::default();
343    /// assert_eq!(quant.scale, 1.0);
344    /// assert_eq!(quant.zero_point, 0);
345    /// ```
346    fn default() -> Self {
347        Self {
348            scale: 1.0,
349            zero_point: 0,
350        }
351    }
352}
353
354/// Turns a DetectBoxQuantized into a DetectBox by dequantizing the score.
355///
356///  # Examples
357/// ```
358/// # use edgefirst_decoder::{BoundingBox, DetectBoxQuantized, Quantization, dequant_detect_box};
359/// let quant = Quantization::new(0.1, -128);
360/// let bbox = BoundingBox::new(0.1, 0.2, 0.3, 0.4);
361/// let detect_quant = DetectBoxQuantized {
362///     bbox,
363///     score: 100_i8,
364///     label: 1,
365/// };
366/// let detect = dequant_detect_box(&detect_quant, quant);
367/// assert_eq!(detect.score, 0.1 * 100.0 + 12.8);
368/// assert_eq!(detect.label, 1);
369/// assert_eq!(detect.bbox, bbox);
370/// ```
371pub fn dequant_detect_box<SCORE: PrimInt + AsPrimitive<f32>>(
372    detect: &DetectBoxQuantized<SCORE>,
373    quant_scores: Quantization,
374) -> DetectBox {
375    let scaled_zp = -quant_scores.scale * quant_scores.zero_point as f32;
376    DetectBox {
377        bbox: detect.bbox,
378        score: quant_scores.scale * detect.score.as_() + scaled_zp,
379        label: detect.label,
380    }
381}
382/// A detection box with a f32 bbox and quantized score
383#[derive(Debug, Clone, Copy, PartialEq)]
384pub struct DetectBoxQuantized<
385    // BOX: Signed + PrimInt + AsPrimitive<f32>,
386    SCORE: PrimInt + AsPrimitive<f32>,
387> {
388    // pub bbox: BoundingBoxQuantized<BOX>,
389    pub bbox: BoundingBox,
390    /// model-specific score for this detection, higher implies more
391    /// confidence.
392    pub score: SCORE,
393    /// label index for this detect
394    pub label: usize,
395}
396
397/// Dequantizes an ndarray from quantized values to f32 values using the given
398/// quantization parameters
399///
400/// # Examples
401/// ```
402/// # use edgefirst_decoder::{dequantize_ndarray, Quantization};
403/// let quant = Quantization::new(0.1, -128);
404/// let input: Vec<i8> = vec![0, 127, -128, 64];
405/// let input_array = ndarray::Array1::from(input);
406/// let output_array: ndarray::Array1<f32> = dequantize_ndarray(input_array.view(), quant);
407/// assert_eq!(output_array, ndarray::array![12.8, 25.5, 0.0, 19.2]);
408/// ```
409pub fn dequantize_ndarray<T: AsPrimitive<F>, D: Dimension, F: Float + 'static>(
410    input: ArrayView<T, D>,
411    quant: Quantization,
412) -> Array<F, D>
413where
414    i32: num_traits::AsPrimitive<F>,
415    f32: num_traits::AsPrimitive<F>,
416{
417    let zero_point = quant.zero_point.as_();
418    let scale = quant.scale.as_();
419    if zero_point != F::zero() {
420        let scaled_zero = -zero_point * scale;
421        input.mapv(|d| d.as_() * scale + scaled_zero)
422    } else {
423        input.mapv(|d| d.as_() * scale)
424    }
425}
426
427/// Dequantizes a slice from quantized values to float values using the given
428/// quantization parameters
429///
430/// # Examples
431/// ```
432/// # use edgefirst_decoder::{dequantize_cpu, Quantization};
433/// let quant = Quantization::new(0.1, -128);
434/// let input: Vec<i8> = vec![0, 127, -128, 64];
435/// let mut output: Vec<f32> = vec![0.0; input.len()];
436/// dequantize_cpu(&input, quant, &mut output);
437/// assert_eq!(output, vec![12.8, 25.5, 0.0, 19.2]);
438/// ```
439pub fn dequantize_cpu<T: AsPrimitive<F>, F: Float + 'static>(
440    input: &[T],
441    quant: Quantization,
442    output: &mut [F],
443) where
444    f32: num_traits::AsPrimitive<F>,
445    i32: num_traits::AsPrimitive<F>,
446{
447    assert!(input.len() == output.len());
448    let zero_point = quant.zero_point.as_();
449    let scale = quant.scale.as_();
450    if zero_point != F::zero() {
451        let scaled_zero = -zero_point * scale; // scale * (d - zero_point) = d * scale - zero_point * scale
452        input
453            .iter()
454            .zip(output)
455            .for_each(|(d, deq)| *deq = d.as_() * scale + scaled_zero);
456    } else {
457        input
458            .iter()
459            .zip(output)
460            .for_each(|(d, deq)| *deq = d.as_() * scale);
461    }
462}
463
464/// Dequantizes a slice from quantized values to float values using the given
465/// quantization parameters, using chunked processing. This is around 5% faster
466/// than `dequantize_cpu` for large slices.
467///
468/// # Examples
469/// ```
470/// # use edgefirst_decoder::{dequantize_cpu_chunked, Quantization};
471/// let quant = Quantization::new(0.1, -128);
472/// let input: Vec<i8> = vec![0, 127, -128, 64];
473/// let mut output: Vec<f32> = vec![0.0; input.len()];
474/// dequantize_cpu_chunked(&input, quant, &mut output);
475/// assert_eq!(output, vec![12.8, 25.5, 0.0, 19.2]);
476/// ```
477pub fn dequantize_cpu_chunked<T: AsPrimitive<F>, F: Float + 'static>(
478    input: &[T],
479    quant: Quantization,
480    output: &mut [F],
481) where
482    f32: num_traits::AsPrimitive<F>,
483    i32: num_traits::AsPrimitive<F>,
484{
485    assert!(input.len() == output.len());
486    let zero_point = quant.zero_point.as_();
487    let scale = quant.scale.as_();
488
489    let input = input.as_chunks::<4>();
490    let output = output.as_chunks_mut::<4>();
491
492    if zero_point != F::zero() {
493        let scaled_zero = -zero_point * scale; // scale * (d - zero_point) = d * scale - zero_point * scale
494
495        input
496            .0
497            .iter()
498            .zip(output.0)
499            .for_each(|(d, deq)| *deq = d.map(|d| d.as_() * scale + scaled_zero));
500        input
501            .1
502            .iter()
503            .zip(output.1)
504            .for_each(|(d, deq)| *deq = d.as_() * scale + scaled_zero);
505    } else {
506        input
507            .0
508            .iter()
509            .zip(output.0)
510            .for_each(|(d, deq)| *deq = d.map(|d| d.as_() * scale));
511        input
512            .1
513            .iter()
514            .zip(output.1)
515            .for_each(|(d, deq)| *deq = d.as_() * scale);
516    }
517}
518
519/// Converts a segmentation tensor into a 2D mask
520/// If the last dimension of the segmentation tensor is 1, values equal or
521/// above 128 are considered objects. Otherwise the object is the argmax index
522///
523/// # Errors
524///
525/// Returns `DecoderError::InvalidShape` if the segmentation tensor has an
526/// invalid shape.
527///
528/// # Examples
529/// ```
530/// # use edgefirst_decoder::segmentation_to_mask;
531/// let segmentation =
532///     ndarray::Array3::<u8>::from_shape_vec((2, 2, 1), vec![0, 255, 128, 127]).unwrap();
533/// let mask = segmentation_to_mask(segmentation.view()).unwrap();
534/// assert_eq!(mask, ndarray::array![[0, 1], [1, 0]]);
535/// ```
536pub fn segmentation_to_mask(segmentation: ArrayView3<u8>) -> Result<Array2<u8>, DecoderError> {
537    if segmentation.shape()[2] == 0 {
538        return Err(DecoderError::InvalidShape(
539            "Segmentation tensor must have non-zero depth".to_string(),
540        ));
541    }
542    if segmentation.shape()[2] == 1 {
543        yolo_segmentation_to_mask(segmentation, 128)
544    } else {
545        Ok(modelpack_segmentation_to_mask(segmentation))
546    }
547}
548
549/// Returns the maximum value and its index from a 1D array
550fn arg_max<T: PartialOrd + Copy>(score: ArrayView1<T>) -> (T, usize) {
551    score
552        .iter()
553        .enumerate()
554        .fold((score[0], 0), |(max, arg_max), (ind, s)| {
555            if max > *s {
556                (max, arg_max)
557            } else {
558                (*s, ind)
559            }
560        })
561}
562
563/// NEON-accelerated argmax for i8 slices on aarch64.
564///
565/// Finds the maximum value and its index (last index wins on ties, matching
566/// the semantics of [`arg_max`]).  Falls back to the scalar implementation
567/// on non-aarch64 targets or when the slice is too short to benefit from NEON.
568#[cfg(target_arch = "aarch64")]
569pub(crate) fn arg_max_i8(scores: &[i8]) -> (i8, usize) {
570    use std::arch::aarch64::*; // NOSONAR
571
572    let n = scores.len();
573    if n < 16 {
574        // Scalar fallback for very short slices.
575        let mut max = scores[0];
576        let mut idx = 0;
577        for (i, &s) in scores.iter().enumerate().skip(1) {
578            if s >= max {
579                max = s;
580                idx = i;
581            }
582        }
583        return (max, idx);
584    }
585
586    unsafe {
587        // Step 1: Find the global max value using NEON horizontal max.
588        let chunks = n / 16;
589        let mut vmax = vld1q_s8(scores.as_ptr());
590        for i in 1..chunks {
591            let v = vld1q_s8(scores.as_ptr().add(i * 16));
592            vmax = vmaxq_s8(vmax, v);
593        }
594        let global_max = vmaxvq_s8(vmax);
595
596        // Handle remainder with scalar
597        let remainder_start = chunks * 16;
598        let mut final_max = global_max;
599        for &s in &scores[remainder_start..] {
600            if s > final_max {
601                final_max = s;
602            }
603        }
604
605        // Step 2: Find the LAST index of `final_max` (preserves tie semantics).
606        // Scan backwards for the last occurrence.
607        let mut idx = 0;
608        for i in (0..n).rev() {
609            if scores[i] == final_max {
610                idx = i;
611                break;
612            }
613        }
614        (final_max, idx)
615    }
616}
617#[cfg(test)]
618#[cfg_attr(coverage_nightly, coverage(off))]
619mod decoder_tests {
620    #![allow(clippy::excessive_precision)]
621    use crate::DetectBoxTestExt;
622    use crate::{
623        configs::{DecoderType, DimName, Protos},
624        modelpack::{decode_modelpack_det, decode_modelpack_split_quant},
625        yolo::{
626            decode_yolo_det, decode_yolo_det_float, decode_yolo_segdet_float,
627            decode_yolo_segdet_quant,
628        },
629        *,
630    };
631    use edgefirst_tensor::{Tensor, TensorMapTrait, TensorTrait};
632    use ndarray::Dimension;
633    use ndarray::{array, s, Array2, Array3, Array4, Axis};
634    use num_traits::{AsPrimitive, PrimInt};
635
636    fn compare_outputs(
637        boxes: (&[DetectBox], &[DetectBox]),
638        masks: (&[Segmentation], &[Segmentation]),
639    ) {
640        let (boxes0, boxes1) = boxes;
641        let (masks0, masks1) = masks;
642
643        assert_eq!(boxes0.len(), boxes1.len());
644        assert_eq!(masks0.len(), masks1.len());
645
646        for (b_i8, b_f32) in boxes0.iter().zip(boxes1) {
647            assert!(
648                b_i8.equal_within_delta(b_f32, 1e-6),
649                "{b_i8:?} is not equal to {b_f32:?}"
650            );
651        }
652
653        for (m_i8, m_f32) in masks0.iter().zip(masks1) {
654            assert_eq!(
655                [m_i8.xmin, m_i8.ymin, m_i8.xmax, m_i8.ymax],
656                [m_f32.xmin, m_f32.ymin, m_f32.xmax, m_f32.ymax],
657            );
658            let shape = m_i8.segmentation.shape().to_vec();
659            assert_eq!(shape, m_f32.segmentation.shape());
660            // `Segmentation` carries a TensorDyn, so compare the mapped bytes
661            // rather than going through ndarray element ops.
662            let (t_i8, t_f32) = (
663                m_i8.segmentation.as_u8().expect("mask must be U8"),
664                m_f32.segmentation.as_u8().expect("mask must be U8"),
665            );
666            let (map_i8, map_f32) = (
667                t_i8.map_read().expect("map i8 mask"),
668                t_f32.map_read().expect("map f32 mask"),
669            );
670            let (a, b) = (map_i8.as_slice(), map_f32.as_slice());
671            let (w, c) = (shape[1], shape[2]);
672            let mut sq_err = 0f64;
673            for (i, (&va, &vb)) in a.iter().zip(b.iter()).enumerate() {
674                let val = va as i32 - vb as i32;
675                assert!(
676                    val.abs() <= 1,
677                    "Difference between mask0 and mask1 is greater than 1 at ({}, {}, {}): {}",
678                    i / (w * c),
679                    (i / c) % w,
680                    i % c,
681                    val
682                );
683                sq_err += (val as f64) * (val as f64);
684            }
685            let mean_sq_err = sq_err / a.len() as f64;
686            assert!(
687                mean_sq_err < 1e-2,
688                "Mean Square Error between masks was greater than 1%: {:.2}%",
689                mean_sq_err * 100.0
690            );
691        }
692    }
693
694    /// Rebuild an owned `Array3<u8>` from a `Segmentation`'s `TensorDyn`, for
695    /// the ndarray-shaped test helpers that still want a view.
696    fn seg_view(s: &Segmentation) -> ndarray::Array3<u8> {
697        let t = s.segmentation.as_u8().expect("mask must be U8");
698        let m = t.map_read().expect("map mask");
699        let sh = s.segmentation.shape();
700        ndarray::Array3::from_shape_vec((sh[0], sh[1], sh[2]), m.as_slice().to_vec())
701            .expect("mask shape")
702    }
703
704    // ─── Shared test data loaders ────────────────────────
705
706    fn load_yolov8_boxes() -> Array3<i8> {
707        let raw = edgefirst_bench::testdata::read("yolov8_boxes_116x8400.bin");
708        let raw = unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const i8, raw.len()) };
709        Array3::from_shape_vec((1, 116, 8400), raw.to_vec()).unwrap()
710    }
711
712    fn load_yolov8_protos() -> Array4<i8> {
713        let raw = edgefirst_bench::testdata::read("yolov8_protos_160x160x32.bin");
714        let raw = unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const i8, raw.len()) };
715        Array4::from_shape_vec((1, 160, 160, 32), raw.to_vec()).unwrap()
716    }
717
718    fn load_yolov8s_det() -> Array3<i8> {
719        let raw = edgefirst_bench::testdata::read("yolov8s_80_classes.bin");
720        let raw = unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const i8, raw.len()) };
721        Array3::from_shape_vec((1, 84, 8400), raw.to_vec()).unwrap()
722    }
723
724    #[test]
725    fn test_decoder_modelpack() {
726        let score_threshold = 0.45;
727        let iou_threshold = 0.45;
728        let boxes = edgefirst_bench::testdata::read("modelpack_boxes_1935x1x4.bin");
729        let boxes = ndarray::Array4::from_shape_vec((1, 1935, 1, 4), boxes.to_vec()).unwrap();
730
731        let scores = edgefirst_bench::testdata::read("modelpack_scores_1935x1.bin");
732        let scores = ndarray::Array3::from_shape_vec((1, 1935, 1), scores.to_vec()).unwrap();
733
734        let quant_boxes = (0.004656755365431309, 21).into();
735        let quant_scores = (0.0019603664986789227, 0).into();
736
737        let decoder = DecoderBuilder::default()
738            .with_config_modelpack_det(
739                configs::Boxes {
740                    decoder: DecoderType::ModelPack,
741                    quantization: Some(quant_boxes),
742                    shape: vec![1, 1935, 1, 4],
743                    dshape: vec![
744                        (DimName::Batch, 1),
745                        (DimName::NumBoxes, 1935),
746                        (DimName::Padding, 1),
747                        (DimName::BoxCoords, 4),
748                    ],
749                    normalized: Some(true),
750                },
751                configs::Scores {
752                    decoder: DecoderType::ModelPack,
753                    quantization: Some(quant_scores),
754                    shape: vec![1, 1935, 1],
755                    dshape: vec![
756                        (DimName::Batch, 1),
757                        (DimName::NumBoxes, 1935),
758                        (DimName::NumClasses, 1),
759                    ],
760                },
761            )
762            .with_score_threshold(score_threshold)
763            .with_iou_threshold(iou_threshold)
764            .build()
765            .unwrap();
766
767        let quant_boxes = quant_boxes.into();
768        let quant_scores = quant_scores.into();
769
770        let mut output_boxes: Vec<_> = Vec::with_capacity(50);
771        decode_modelpack_det(
772            (boxes.slice(s![0, .., 0, ..]), quant_boxes),
773            (scores.slice(s![0, .., ..]), quant_scores),
774            score_threshold,
775            iou_threshold,
776            300,
777            &mut output_boxes,
778        );
779        assert!(output_boxes[0].equal_within_delta(
780            &DetectBox {
781                bbox: BoundingBox {
782                    xmin: 0.40513772,
783                    ymin: 0.6379755,
784                    xmax: 0.5122431,
785                    ymax: 0.7730214,
786                },
787                score: 0.4861709,
788                label: 0
789            },
790            1e-6
791        ));
792
793        let mut output_boxes1 = Vec::with_capacity(50);
794        let mut output_masks1 = Vec::with_capacity(50);
795
796        decoder
797            .decode_quantized(
798                &[boxes.view().into(), scores.view().into()],
799                &mut output_boxes1,
800                &mut output_masks1,
801            )
802            .unwrap();
803
804        let mut output_boxes_float = Vec::with_capacity(50);
805        let mut output_masks_float = Vec::with_capacity(50);
806
807        let boxes = dequantize_ndarray(boxes.view(), quant_boxes);
808        let scores = dequantize_ndarray(scores.view(), quant_scores);
809
810        decoder
811            .decode_float::<f32>(
812                &[boxes.view().into_dyn(), scores.view().into_dyn()],
813                &mut output_boxes_float,
814                &mut output_masks_float,
815            )
816            .unwrap();
817
818        compare_outputs((&output_boxes, &output_boxes1), (&[], &output_masks1));
819        compare_outputs(
820            (&output_boxes, &output_boxes_float),
821            (&[], &output_masks_float),
822        );
823    }
824
825    #[test]
826    fn test_decoder_modelpack_split_u8() {
827        let score_threshold = 0.45;
828        let iou_threshold = 0.45;
829        let detect0 = edgefirst_bench::testdata::read("modelpack_split_9x15x18.bin");
830        let detect0 = ndarray::Array4::from_shape_vec((1, 9, 15, 18), detect0.to_vec()).unwrap();
831
832        let detect1 = edgefirst_bench::testdata::read("modelpack_split_17x30x18.bin");
833        let detect1 = ndarray::Array4::from_shape_vec((1, 17, 30, 18), detect1.to_vec()).unwrap();
834
835        let quant0 = (0.08547406643629074, 174).into();
836        let quant1 = (0.09929127991199493, 183).into();
837        let anchors0 = vec![
838            [0.36666667461395264, 0.31481480598449707],
839            [0.38749998807907104, 0.4740740656852722],
840            [0.5333333611488342, 0.644444465637207],
841        ];
842        let anchors1 = vec![
843            [0.13750000298023224, 0.2074074000120163],
844            [0.2541666626930237, 0.21481481194496155],
845            [0.23125000298023224, 0.35185185074806213],
846        ];
847
848        let detect_config0 = configs::Detection {
849            decoder: DecoderType::ModelPack,
850            shape: vec![1, 9, 15, 18],
851            anchors: Some(anchors0.clone()),
852            quantization: Some(quant0),
853            dshape: vec![
854                (DimName::Batch, 1),
855                (DimName::Height, 9),
856                (DimName::Width, 15),
857                (DimName::NumAnchorsXFeatures, 18),
858            ],
859            normalized: Some(true),
860        };
861
862        let detect_config1 = configs::Detection {
863            decoder: DecoderType::ModelPack,
864            shape: vec![1, 17, 30, 18],
865            anchors: Some(anchors1.clone()),
866            quantization: Some(quant1),
867            dshape: vec![
868                (DimName::Batch, 1),
869                (DimName::Height, 17),
870                (DimName::Width, 30),
871                (DimName::NumAnchorsXFeatures, 18),
872            ],
873            normalized: Some(true),
874        };
875
876        let config0 = (&detect_config0).try_into().unwrap();
877        let config1 = (&detect_config1).try_into().unwrap();
878
879        let decoder = DecoderBuilder::default()
880            .with_config_modelpack_det_split(vec![detect_config1, detect_config0])
881            .with_score_threshold(score_threshold)
882            .with_iou_threshold(iou_threshold)
883            .build()
884            .unwrap();
885
886        let quant0 = quant0.into();
887        let quant1 = quant1.into();
888
889        let mut output_boxes: Vec<_> = Vec::with_capacity(10);
890        decode_modelpack_split_quant(
891            &[
892                detect0.slice(s![0, .., .., ..]),
893                detect1.slice(s![0, .., .., ..]),
894            ],
895            &[config0, config1],
896            score_threshold,
897            iou_threshold,
898            300,
899            &mut output_boxes,
900        );
901        assert!(output_boxes[0].equal_within_delta(
902            &DetectBox {
903                bbox: BoundingBox {
904                    xmin: 0.43171933,
905                    ymin: 0.68243736,
906                    xmax: 0.5626645,
907                    ymax: 0.808863,
908                },
909                score: 0.99240804,
910                label: 0
911            },
912            1e-6
913        ));
914
915        let mut output_boxes1: Vec<_> = Vec::with_capacity(10);
916        let mut output_masks1: Vec<_> = Vec::with_capacity(10);
917        decoder
918            .decode_quantized(
919                &[detect0.view().into(), detect1.view().into()],
920                &mut output_boxes1,
921                &mut output_masks1,
922            )
923            .unwrap();
924
925        let mut output_boxes1_f32: Vec<_> = Vec::with_capacity(10);
926        let mut output_masks1_f32: Vec<_> = Vec::with_capacity(10);
927
928        let detect0 = dequantize_ndarray(detect0.view(), quant0);
929        let detect1 = dequantize_ndarray(detect1.view(), quant1);
930        decoder
931            .decode_float::<f32>(
932                &[detect0.view().into_dyn(), detect1.view().into_dyn()],
933                &mut output_boxes1_f32,
934                &mut output_masks1_f32,
935            )
936            .unwrap();
937
938        compare_outputs((&output_boxes, &output_boxes1), (&[], &output_masks1));
939        compare_outputs(
940            (&output_boxes, &output_boxes1_f32),
941            (&[], &output_masks1_f32),
942        );
943    }
944
945    #[test]
946    fn test_decoder_parse_config_modelpack_split_u8() {
947        let score_threshold = 0.45;
948        let iou_threshold = 0.45;
949        let detect0 = edgefirst_bench::testdata::read("modelpack_split_9x15x18.bin");
950        let detect0 = ndarray::Array4::from_shape_vec((1, 9, 15, 18), detect0.to_vec()).unwrap();
951
952        let detect1 = edgefirst_bench::testdata::read("modelpack_split_17x30x18.bin");
953        let detect1 = ndarray::Array4::from_shape_vec((1, 17, 30, 18), detect1.to_vec()).unwrap();
954
955        let decoder = DecoderBuilder::default()
956            .with_config_yaml_str(
957                edgefirst_bench::testdata::read_to_string("modelpack_split.yaml").to_string(),
958            )
959            .with_score_threshold(score_threshold)
960            .with_iou_threshold(iou_threshold)
961            .build()
962            .unwrap();
963
964        let mut output_boxes: Vec<_> = Vec::with_capacity(10);
965        let mut output_masks: Vec<_> = Vec::with_capacity(10);
966        decoder
967            .decode_quantized(
968                &[
969                    ArrayViewDQuantized::from(detect1.view()),
970                    ArrayViewDQuantized::from(detect0.view()),
971                ],
972                &mut output_boxes,
973                &mut output_masks,
974            )
975            .unwrap();
976        assert!(output_boxes[0].equal_within_delta(
977            &DetectBox {
978                bbox: BoundingBox {
979                    xmin: 0.43171933,
980                    ymin: 0.68243736,
981                    xmax: 0.5626645,
982                    ymax: 0.808863,
983                },
984                score: 0.99240804,
985                label: 0
986            },
987            1e-6
988        ));
989    }
990
991    #[test]
992    fn test_modelpack_seg() {
993        let out = edgefirst_bench::testdata::read("modelpack_seg_2x160x160.bin");
994        let out = ndarray::Array4::from_shape_vec((1, 2, 160, 160), out.to_vec()).unwrap();
995        let quant = (1.0 / 255.0, 0).into();
996
997        let decoder = DecoderBuilder::default()
998            .with_config_modelpack_seg(configs::Segmentation {
999                decoder: DecoderType::ModelPack,
1000                quantization: Some(quant),
1001                shape: vec![1, 2, 160, 160],
1002                dshape: vec![
1003                    (DimName::Batch, 1),
1004                    (DimName::NumClasses, 2),
1005                    (DimName::Height, 160),
1006                    (DimName::Width, 160),
1007                ],
1008            })
1009            .build()
1010            .unwrap();
1011        let mut output_boxes: Vec<_> = Vec::with_capacity(10);
1012        let mut output_masks: Vec<_> = Vec::with_capacity(10);
1013        decoder
1014            .decode_quantized(&[out.view().into()], &mut output_boxes, &mut output_masks)
1015            .unwrap();
1016
1017        let mut mask = out.slice(s![0, .., .., ..]);
1018        mask.swap_axes(0, 1);
1019        mask.swap_axes(1, 2);
1020        let mask = [Segmentation {
1021            xmin: 0.0,
1022            ymin: 0.0,
1023            xmax: 1.0,
1024            ymax: 1.0,
1025            segmentation: crate::mask_to_tensor(mask.view()).unwrap(),
1026        }];
1027        compare_outputs((&[], &output_boxes), (&mask, &output_masks));
1028
1029        decoder
1030            .decode_float::<f32>(
1031                &[dequantize_ndarray(out.view(), quant.into())
1032                    .view()
1033                    .into_dyn()],
1034                &mut output_boxes,
1035                &mut output_masks,
1036            )
1037            .unwrap();
1038
1039        // not expected for float decoder to have same values as quantized decoder, as
1040        // float decoder ensures the data fills 0-255, quantized decoder uses whatever
1041        // the model output. Thus the float output is the same as the quantized output
1042        // but scaled differently. However, it is expected that the mask after argmax
1043        // will be the same.
1044        compare_outputs((&[], &output_boxes), (&[], &[]));
1045        let mask0 = segmentation_to_mask(seg_view(&mask[0]).view()).unwrap();
1046        let mask1 = segmentation_to_mask(seg_view(&output_masks[0]).view()).unwrap();
1047
1048        assert_eq!(mask0, mask1);
1049    }
1050    #[test]
1051    fn test_modelpack_seg_quant() {
1052        let out = edgefirst_bench::testdata::read("modelpack_seg_2x160x160.bin");
1053        let out_u8 = ndarray::Array4::from_shape_vec((1, 2, 160, 160), out.to_vec()).unwrap();
1054        let out_i8 = out_u8.mapv(|x| (x as i16 - 128) as i8);
1055        let out_u16 = out_u8.mapv(|x| (x as u16) << 8);
1056        let out_i16 = out_u8.mapv(|x| (((x as i32) << 8) - 32768) as i16);
1057        let out_u32 = out_u8.mapv(|x| (x as u32) << 24);
1058        let out_i32 = out_u8.mapv(|x| (((x as i64) << 24) - 2147483648) as i32);
1059
1060        let quant = (1.0 / 255.0, 0).into();
1061
1062        let decoder = DecoderBuilder::default()
1063            .with_config_modelpack_seg(configs::Segmentation {
1064                decoder: DecoderType::ModelPack,
1065                quantization: Some(quant),
1066                shape: vec![1, 2, 160, 160],
1067                dshape: vec![
1068                    (DimName::Batch, 1),
1069                    (DimName::NumClasses, 2),
1070                    (DimName::Height, 160),
1071                    (DimName::Width, 160),
1072                ],
1073            })
1074            .build()
1075            .unwrap();
1076        let mut output_boxes: Vec<_> = Vec::with_capacity(10);
1077        let mut output_masks_u8: Vec<_> = Vec::with_capacity(10);
1078        decoder
1079            .decode_quantized(
1080                &[out_u8.view().into()],
1081                &mut output_boxes,
1082                &mut output_masks_u8,
1083            )
1084            .unwrap();
1085
1086        let mut output_masks_i8: Vec<_> = Vec::with_capacity(10);
1087        decoder
1088            .decode_quantized(
1089                &[out_i8.view().into()],
1090                &mut output_boxes,
1091                &mut output_masks_i8,
1092            )
1093            .unwrap();
1094
1095        let mut output_masks_u16: Vec<_> = Vec::with_capacity(10);
1096        decoder
1097            .decode_quantized(
1098                &[out_u16.view().into()],
1099                &mut output_boxes,
1100                &mut output_masks_u16,
1101            )
1102            .unwrap();
1103
1104        let mut output_masks_i16: Vec<_> = Vec::with_capacity(10);
1105        decoder
1106            .decode_quantized(
1107                &[out_i16.view().into()],
1108                &mut output_boxes,
1109                &mut output_masks_i16,
1110            )
1111            .unwrap();
1112
1113        let mut output_masks_u32: Vec<_> = Vec::with_capacity(10);
1114        decoder
1115            .decode_quantized(
1116                &[out_u32.view().into()],
1117                &mut output_boxes,
1118                &mut output_masks_u32,
1119            )
1120            .unwrap();
1121
1122        let mut output_masks_i32: Vec<_> = Vec::with_capacity(10);
1123        decoder
1124            .decode_quantized(
1125                &[out_i32.view().into()],
1126                &mut output_boxes,
1127                &mut output_masks_i32,
1128            )
1129            .unwrap();
1130
1131        compare_outputs((&[], &output_boxes), (&[], &[]));
1132        let mask_u8 = segmentation_to_mask(seg_view(&output_masks_u8[0]).view()).unwrap();
1133        let mask_i8 = segmentation_to_mask(seg_view(&output_masks_i8[0]).view()).unwrap();
1134        let mask_u16 = segmentation_to_mask(seg_view(&output_masks_u16[0]).view()).unwrap();
1135        let mask_i16 = segmentation_to_mask(seg_view(&output_masks_i16[0]).view()).unwrap();
1136        let mask_u32 = segmentation_to_mask(seg_view(&output_masks_u32[0]).view()).unwrap();
1137        let mask_i32 = segmentation_to_mask(seg_view(&output_masks_i32[0]).view()).unwrap();
1138        assert_eq!(mask_u8, mask_i8);
1139        assert_eq!(mask_u8, mask_u16);
1140        assert_eq!(mask_u8, mask_i16);
1141        assert_eq!(mask_u8, mask_u32);
1142        assert_eq!(mask_u8, mask_i32);
1143    }
1144
1145    #[test]
1146    fn test_modelpack_segdet() {
1147        let score_threshold = 0.45;
1148        let iou_threshold = 0.45;
1149
1150        let boxes = edgefirst_bench::testdata::read("modelpack_boxes_1935x1x4.bin");
1151        let boxes = Array4::from_shape_vec((1, 1935, 1, 4), boxes.to_vec()).unwrap();
1152
1153        let scores = edgefirst_bench::testdata::read("modelpack_scores_1935x1.bin");
1154        let scores = Array3::from_shape_vec((1, 1935, 1), scores.to_vec()).unwrap();
1155
1156        let seg = edgefirst_bench::testdata::read("modelpack_seg_2x160x160.bin");
1157        let seg = Array4::from_shape_vec((1, 2, 160, 160), seg.to_vec()).unwrap();
1158
1159        let quant_boxes = (0.004656755365431309, 21).into();
1160        let quant_scores = (0.0019603664986789227, 0).into();
1161        let quant_seg = (1.0 / 255.0, 0).into();
1162
1163        let decoder = DecoderBuilder::default()
1164            .with_config_modelpack_segdet(
1165                configs::Boxes {
1166                    decoder: DecoderType::ModelPack,
1167                    quantization: Some(quant_boxes),
1168                    shape: vec![1, 1935, 1, 4],
1169                    dshape: vec![
1170                        (DimName::Batch, 1),
1171                        (DimName::NumBoxes, 1935),
1172                        (DimName::Padding, 1),
1173                        (DimName::BoxCoords, 4),
1174                    ],
1175                    normalized: Some(true),
1176                },
1177                configs::Scores {
1178                    decoder: DecoderType::ModelPack,
1179                    quantization: Some(quant_scores),
1180                    shape: vec![1, 1935, 1],
1181                    dshape: vec![
1182                        (DimName::Batch, 1),
1183                        (DimName::NumBoxes, 1935),
1184                        (DimName::NumClasses, 1),
1185                    ],
1186                },
1187                configs::Segmentation {
1188                    decoder: DecoderType::ModelPack,
1189                    quantization: Some(quant_seg),
1190                    shape: vec![1, 2, 160, 160],
1191                    dshape: vec![
1192                        (DimName::Batch, 1),
1193                        (DimName::NumClasses, 2),
1194                        (DimName::Height, 160),
1195                        (DimName::Width, 160),
1196                    ],
1197                },
1198            )
1199            .with_iou_threshold(iou_threshold)
1200            .with_score_threshold(score_threshold)
1201            .build()
1202            .unwrap();
1203        let mut output_boxes: Vec<_> = Vec::with_capacity(10);
1204        let mut output_masks: Vec<_> = Vec::with_capacity(10);
1205        decoder
1206            .decode_quantized(
1207                &[scores.view().into(), boxes.view().into(), seg.view().into()],
1208                &mut output_boxes,
1209                &mut output_masks,
1210            )
1211            .unwrap();
1212
1213        let mut mask = seg.slice(s![0, .., .., ..]);
1214        mask.swap_axes(0, 1);
1215        mask.swap_axes(1, 2);
1216        let mask = [Segmentation {
1217            xmin: 0.0,
1218            ymin: 0.0,
1219            xmax: 1.0,
1220            ymax: 1.0,
1221            segmentation: crate::mask_to_tensor(mask.view()).unwrap(),
1222        }];
1223        let correct_boxes = [DetectBox {
1224            bbox: BoundingBox {
1225                xmin: 0.40513772,
1226                ymin: 0.6379755,
1227                xmax: 0.5122431,
1228                ymax: 0.7730214,
1229            },
1230            score: 0.4861709,
1231            label: 0,
1232        }];
1233        compare_outputs((&correct_boxes, &output_boxes), (&mask, &output_masks));
1234
1235        let scores = dequantize_ndarray(scores.view(), quant_scores.into());
1236        let boxes = dequantize_ndarray(boxes.view(), quant_boxes.into());
1237        let seg = dequantize_ndarray(seg.view(), quant_seg.into());
1238        decoder
1239            .decode_float::<f32>(
1240                &[
1241                    scores.view().into_dyn(),
1242                    boxes.view().into_dyn(),
1243                    seg.view().into_dyn(),
1244                ],
1245                &mut output_boxes,
1246                &mut output_masks,
1247            )
1248            .unwrap();
1249
1250        // not expected for float segmentation decoder to have same values as quantized
1251        // segmentation decoder, as float decoder ensures the data fills 0-255,
1252        // quantized decoder uses whatever the model output. Thus the float
1253        // output is the same as the quantized output but scaled differently.
1254        // However, it is expected that the mask after argmax will be the same.
1255        compare_outputs((&correct_boxes, &output_boxes), (&[], &[]));
1256        let mask0 = segmentation_to_mask(seg_view(&mask[0]).view()).unwrap();
1257        let mask1 = segmentation_to_mask(seg_view(&output_masks[0]).view()).unwrap();
1258
1259        assert_eq!(mask0, mask1);
1260    }
1261
1262    #[test]
1263    fn test_modelpack_segdet_split() {
1264        let score_threshold = 0.8;
1265        let iou_threshold = 0.5;
1266
1267        let seg = edgefirst_bench::testdata::read("modelpack_seg_2x160x160.bin");
1268        let seg = ndarray::Array4::from_shape_vec((1, 2, 160, 160), seg.to_vec()).unwrap();
1269
1270        let detect0 = edgefirst_bench::testdata::read("modelpack_split_9x15x18.bin");
1271        let detect0 = ndarray::Array4::from_shape_vec((1, 9, 15, 18), detect0.to_vec()).unwrap();
1272
1273        let detect1 = edgefirst_bench::testdata::read("modelpack_split_17x30x18.bin");
1274        let detect1 = ndarray::Array4::from_shape_vec((1, 17, 30, 18), detect1.to_vec()).unwrap();
1275
1276        let quant0 = (0.08547406643629074, 174).into();
1277        let quant1 = (0.09929127991199493, 183).into();
1278        let quant_seg = (1.0 / 255.0, 0).into();
1279
1280        let anchors0 = vec![
1281            [0.36666667461395264, 0.31481480598449707],
1282            [0.38749998807907104, 0.4740740656852722],
1283            [0.5333333611488342, 0.644444465637207],
1284        ];
1285        let anchors1 = vec![
1286            [0.13750000298023224, 0.2074074000120163],
1287            [0.2541666626930237, 0.21481481194496155],
1288            [0.23125000298023224, 0.35185185074806213],
1289        ];
1290
1291        let decoder = DecoderBuilder::default()
1292            .with_config_modelpack_segdet_split(
1293                vec![
1294                    configs::Detection {
1295                        decoder: DecoderType::ModelPack,
1296                        shape: vec![1, 17, 30, 18],
1297                        anchors: Some(anchors1),
1298                        quantization: Some(quant1),
1299                        dshape: vec![
1300                            (DimName::Batch, 1),
1301                            (DimName::Height, 17),
1302                            (DimName::Width, 30),
1303                            (DimName::NumAnchorsXFeatures, 18),
1304                        ],
1305                        normalized: Some(true),
1306                    },
1307                    configs::Detection {
1308                        decoder: DecoderType::ModelPack,
1309                        shape: vec![1, 9, 15, 18],
1310                        anchors: Some(anchors0),
1311                        quantization: Some(quant0),
1312                        dshape: vec![
1313                            (DimName::Batch, 1),
1314                            (DimName::Height, 9),
1315                            (DimName::Width, 15),
1316                            (DimName::NumAnchorsXFeatures, 18),
1317                        ],
1318                        normalized: Some(true),
1319                    },
1320                ],
1321                configs::Segmentation {
1322                    decoder: DecoderType::ModelPack,
1323                    quantization: Some(quant_seg),
1324                    shape: vec![1, 2, 160, 160],
1325                    dshape: vec![
1326                        (DimName::Batch, 1),
1327                        (DimName::NumClasses, 2),
1328                        (DimName::Height, 160),
1329                        (DimName::Width, 160),
1330                    ],
1331                },
1332            )
1333            .with_score_threshold(score_threshold)
1334            .with_iou_threshold(iou_threshold)
1335            .build()
1336            .unwrap();
1337        let mut output_boxes: Vec<_> = Vec::with_capacity(10);
1338        let mut output_masks: Vec<_> = Vec::with_capacity(10);
1339        decoder
1340            .decode_quantized(
1341                &[
1342                    detect0.view().into(),
1343                    detect1.view().into(),
1344                    seg.view().into(),
1345                ],
1346                &mut output_boxes,
1347                &mut output_masks,
1348            )
1349            .unwrap();
1350
1351        let mut mask = seg.slice(s![0, .., .., ..]);
1352        mask.swap_axes(0, 1);
1353        mask.swap_axes(1, 2);
1354        let mask = [Segmentation {
1355            xmin: 0.0,
1356            ymin: 0.0,
1357            xmax: 1.0,
1358            ymax: 1.0,
1359            segmentation: crate::mask_to_tensor(mask.view()).unwrap(),
1360        }];
1361        let correct_boxes = [DetectBox {
1362            bbox: BoundingBox {
1363                xmin: 0.43171933,
1364                ymin: 0.68243736,
1365                xmax: 0.5626645,
1366                ymax: 0.808863,
1367            },
1368            score: 0.99240804,
1369            label: 0,
1370        }];
1371        println!("Output Boxes: {:?}", output_boxes);
1372        compare_outputs((&correct_boxes, &output_boxes), (&mask, &output_masks));
1373
1374        let detect0 = dequantize_ndarray(detect0.view(), quant0.into());
1375        let detect1 = dequantize_ndarray(detect1.view(), quant1.into());
1376        let seg = dequantize_ndarray(seg.view(), quant_seg.into());
1377        decoder
1378            .decode_float::<f32>(
1379                &[
1380                    detect0.view().into_dyn(),
1381                    detect1.view().into_dyn(),
1382                    seg.view().into_dyn(),
1383                ],
1384                &mut output_boxes,
1385                &mut output_masks,
1386            )
1387            .unwrap();
1388
1389        // not expected for float segmentation decoder to have same values as quantized
1390        // segmentation decoder, as float decoder ensures the data fills 0-255,
1391        // quantized decoder uses whatever the model output. Thus the float
1392        // output is the same as the quantized output but scaled differently.
1393        // However, it is expected that the mask after argmax will be the same.
1394        compare_outputs((&correct_boxes, &output_boxes), (&[], &[]));
1395        let mask0 = segmentation_to_mask(seg_view(&mask[0]).view()).unwrap();
1396        let mask1 = segmentation_to_mask(seg_view(&output_masks[0]).view()).unwrap();
1397
1398        assert_eq!(mask0, mask1);
1399    }
1400
1401    #[test]
1402    fn test_dequant_chunked() {
1403        let mut out = load_yolov8s_det().into_raw_vec_and_offset().0;
1404        out.push(123); // make sure to test non multiple of 16 length
1405
1406        let mut out_dequant = vec![0.0; 84 * 8400 + 1];
1407        let mut out_dequant_simd = vec![0.0; 84 * 8400 + 1];
1408        let quant = Quantization::new(0.0040811873, -123);
1409        dequantize_cpu(&out, quant, &mut out_dequant);
1410
1411        dequantize_cpu_chunked(&out, quant, &mut out_dequant_simd);
1412        assert_eq!(out_dequant, out_dequant_simd);
1413
1414        let quant = Quantization::new(0.0040811873, 0);
1415        dequantize_cpu(&out, quant, &mut out_dequant);
1416
1417        dequantize_cpu_chunked(&out, quant, &mut out_dequant_simd);
1418        assert_eq!(out_dequant, out_dequant_simd);
1419    }
1420
1421    #[test]
1422    fn test_dequant_ground_truth() {
1423        // Formula: output = (input - zero_point) * scale
1424        // Verify both dequantize_cpu and dequantize_cpu_chunked against hand-computed values.
1425
1426        // Case 1: scale=0.1, zero_point=-128 (from doc example)
1427        let quant = Quantization::new(0.1, -128);
1428        let input: Vec<i8> = vec![0, 127, -128, 64];
1429        let mut output = vec![0.0f32; 4];
1430        let mut output_chunked = vec![0.0f32; 4];
1431        dequantize_cpu(&input, quant, &mut output);
1432        dequantize_cpu_chunked(&input, quant, &mut output_chunked);
1433        // (0 - (-128)) * 0.1 = 12.8
1434        // (127 - (-128)) * 0.1 = 25.5
1435        // (-128 - (-128)) * 0.1 = 0.0
1436        // (64 - (-128)) * 0.1 = 19.2
1437        let expected: Vec<f32> = vec![12.8, 25.5, 0.0, 19.2];
1438        for (i, (&out, &exp)) in output.iter().zip(expected.iter()).enumerate() {
1439            assert!((out - exp).abs() < 1e-5, "cpu[{i}]: {out} != {exp}");
1440        }
1441        for (i, (&out, &exp)) in output_chunked.iter().zip(expected.iter()).enumerate() {
1442            assert!((out - exp).abs() < 1e-5, "chunked[{i}]: {out} != {exp}");
1443        }
1444
1445        // Case 2: scale=1.0, zero_point=0 (identity-like)
1446        let quant = Quantization::new(1.0, 0);
1447        dequantize_cpu(&input, quant, &mut output);
1448        dequantize_cpu_chunked(&input, quant, &mut output_chunked);
1449        let expected: Vec<f32> = vec![0.0, 127.0, -128.0, 64.0];
1450        assert_eq!(output, expected);
1451        assert_eq!(output_chunked, expected);
1452
1453        // Case 3: scale=0.5, zero_point=0
1454        let quant = Quantization::new(0.5, 0);
1455        dequantize_cpu(&input, quant, &mut output);
1456        dequantize_cpu_chunked(&input, quant, &mut output_chunked);
1457        let expected: Vec<f32> = vec![0.0, 63.5, -64.0, 32.0];
1458        assert_eq!(output, expected);
1459        assert_eq!(output_chunked, expected);
1460
1461        // Case 4: i8 min/max boundaries with typical quantization params
1462        let quant = Quantization::new(0.021287762, 31);
1463        let input: Vec<i8> = vec![-128, -1, 0, 1, 31, 127];
1464        let mut output = vec![0.0f32; 6];
1465        let mut output_chunked = vec![0.0f32; 6];
1466        dequantize_cpu(&input, quant, &mut output);
1467        dequantize_cpu_chunked(&input, quant, &mut output_chunked);
1468        for i in 0..6 {
1469            let expected = (input[i] as f32 - 31.0) * 0.021287762;
1470            assert!(
1471                (output[i] - expected).abs() < 1e-5,
1472                "cpu[{i}]: {} != {expected}",
1473                output[i]
1474            );
1475            assert!(
1476                (output_chunked[i] - expected).abs() < 1e-5,
1477                "chunked[{i}]: {} != {expected}",
1478                output_chunked[i]
1479            );
1480        }
1481    }
1482
1483    #[test]
1484    fn test_decoder_yolo_det() {
1485        let score_threshold = 0.25;
1486        let iou_threshold = 0.7;
1487        let out = load_yolov8s_det();
1488        let quant = (0.0040811873, -123).into();
1489
1490        let decoder = DecoderBuilder::default()
1491            .with_config_yolo_det(
1492                configs::Detection {
1493                    decoder: DecoderType::Ultralytics,
1494                    shape: vec![1, 84, 8400],
1495                    anchors: None,
1496                    quantization: Some(quant),
1497                    dshape: vec![
1498                        (DimName::Batch, 1),
1499                        (DimName::NumFeatures, 84),
1500                        (DimName::NumBoxes, 8400),
1501                    ],
1502                    normalized: Some(true),
1503                },
1504                Some(DecoderVersion::Yolo11),
1505            )
1506            .with_score_threshold(score_threshold)
1507            .with_iou_threshold(iou_threshold)
1508            .build()
1509            .unwrap();
1510
1511        let mut output_boxes: Vec<_> = Vec::with_capacity(50);
1512        decode_yolo_det(
1513            (out.slice(s![0, .., ..]), quant.into()),
1514            score_threshold,
1515            iou_threshold,
1516            Some(configs::Nms::ClassAgnostic),
1517            &mut output_boxes,
1518        );
1519        assert!(output_boxes[0].equal_within_delta(
1520            &DetectBox {
1521                bbox: BoundingBox {
1522                    xmin: 0.5285137,
1523                    ymin: 0.05305544,
1524                    xmax: 0.87541467,
1525                    ymax: 0.9998909,
1526                },
1527                score: 0.5591227,
1528                label: 0
1529            },
1530            1e-6
1531        ));
1532
1533        assert!(output_boxes[1].equal_within_delta(
1534            &DetectBox {
1535                bbox: BoundingBox {
1536                    xmin: 0.130598,
1537                    ymin: 0.43260583,
1538                    xmax: 0.35098213,
1539                    ymax: 0.9958097,
1540                },
1541                score: 0.33057618,
1542                label: 75
1543            },
1544            1e-6
1545        ));
1546
1547        let mut output_boxes1: Vec<_> = Vec::with_capacity(50);
1548        let mut output_masks1: Vec<_> = Vec::with_capacity(50);
1549        decoder
1550            .decode_quantized(&[out.view().into()], &mut output_boxes1, &mut output_masks1)
1551            .unwrap();
1552
1553        let out = dequantize_ndarray(out.view(), quant.into());
1554        let mut output_boxes_f32: Vec<_> = Vec::with_capacity(50);
1555        let mut output_masks_f32: Vec<_> = Vec::with_capacity(50);
1556        decoder
1557            .decode_float::<f32>(
1558                &[out.view().into_dyn()],
1559                &mut output_boxes_f32,
1560                &mut output_masks_f32,
1561            )
1562            .unwrap();
1563
1564        compare_outputs((&output_boxes, &output_boxes1), (&[], &output_masks1));
1565        compare_outputs((&output_boxes, &output_boxes_f32), (&[], &output_masks_f32));
1566    }
1567
1568    #[test]
1569    fn test_decoder_masks() {
1570        let score_threshold = 0.45;
1571        let iou_threshold = 0.45;
1572        let boxes = load_yolov8_boxes();
1573        let quant_boxes = Quantization::new(0.021287761628627777, 31);
1574
1575        let protos = load_yolov8_protos();
1576        let quant_protos = Quantization::new(0.02491161972284317, -117);
1577        let protos = dequantize_ndarray::<_, _, f32>(protos.view(), quant_protos);
1578        let seg = dequantize_ndarray::<_, _, f32>(boxes.view(), quant_boxes);
1579        let mut output_boxes: Vec<_> = Vec::with_capacity(10);
1580        let mut output_masks: Vec<_> = Vec::with_capacity(10);
1581        decode_yolo_segdet_float(
1582            seg.slice(s![0, .., ..]),
1583            protos.slice(s![0, .., .., ..]),
1584            score_threshold,
1585            iou_threshold,
1586            Some(configs::Nms::ClassAgnostic),
1587            &mut output_boxes,
1588            &mut output_masks,
1589        )
1590        .unwrap();
1591        assert_eq!(output_boxes.len(), 2);
1592        assert_eq!(output_boxes.len(), output_masks.len());
1593
1594        for (b, m) in output_boxes.iter().zip(&output_masks) {
1595            // Mask region is the proto-grid-aligned crop (floor for min,
1596            // ceil for max), so it encloses the post-NMS bbox (EDGEAI-1304).
1597            assert!(b.bbox.xmin >= m.xmin);
1598            assert!(b.bbox.ymin >= m.ymin);
1599            assert!(b.bbox.xmax <= m.xmax);
1600            assert!(b.bbox.ymax <= m.ymax);
1601        }
1602        assert!(output_boxes[0].equal_within_delta(
1603            &DetectBox {
1604                bbox: BoundingBox {
1605                    xmin: 0.08515105,
1606                    ymin: 0.7131401,
1607                    xmax: 0.29802868,
1608                    ymax: 0.8195788,
1609                },
1610                score: 0.91537374,
1611                label: 23
1612            },
1613            1.0 / 160.0, // wider range because mask will expand the box
1614        ));
1615
1616        assert!(output_boxes[1].equal_within_delta(
1617            &DetectBox {
1618                bbox: BoundingBox {
1619                    xmin: 0.59605736,
1620                    ymin: 0.25545314,
1621                    xmax: 0.93666154,
1622                    ymax: 0.72378385,
1623                },
1624                score: 0.91537374,
1625                label: 23
1626            },
1627            1.0 / 160.0, // wider range because mask will expand the box
1628        ));
1629
1630        let full_mask = edgefirst_bench::testdata::read("yolov8_mask_results.bin");
1631        let full_mask = ndarray::Array2::from_shape_vec((160, 160), full_mask.to_vec()).unwrap();
1632
1633        let cropped_mask = full_mask.slice(ndarray::s![
1634            (output_masks[1].ymin * 160.0) as usize..(output_masks[1].ymax * 160.0) as usize,
1635            (output_masks[1].xmin * 160.0) as usize..(output_masks[1].xmax * 160.0) as usize,
1636        ]);
1637
1638        assert_eq!(
1639            cropped_mask,
1640            segmentation_to_mask(seg_view(&output_masks[1]).view()).unwrap()
1641        );
1642    }
1643
1644    /// Regression test: config-driven path with physically-NCHW protos.
1645    /// Simulates YOLOv8-seg ONNX outputs where the producer emits protos
1646    /// as `(1, 32, 160, 160)` in CHW memory order — the caller declares
1647    /// shape and dshape matching that physical order and HAL permutes to
1648    /// canonical HWC via `swap_axes_if_needed`.
1649    ///
1650    /// This is the counterpart to the NHWC-producer case (TFLite /
1651    /// Ara-2) where shape+dshape are `(1, 160, 160, 32)` +
1652    /// `[batch, height, width, num_protos]` and no reordering is needed.
1653    #[test]
1654    fn test_decoder_masks_nchw_protos() {
1655        let score_threshold = 0.45;
1656        let iou_threshold = 0.45;
1657
1658        // Load test data — boxes as [116, 8400]
1659        let boxes_2d = load_yolov8_boxes().slice_move(s![0, .., ..]);
1660        let quant_boxes = Quantization::new(0.021287761628627777, 31);
1661
1662        // Load protos as HWC [160, 160, 32] (file layout) then dequantize
1663        let protos_hwc = load_yolov8_protos().slice_move(s![0, .., .., ..]);
1664        let quant_protos = Quantization::new(0.02491161972284317, -117);
1665        let protos_f32_hwc = dequantize_ndarray::<_, _, f32>(protos_hwc.view(), quant_protos);
1666
1667        // ---- Reference: direct call with HWC protos (known working) ----
1668        let seg = dequantize_ndarray::<_, _, f32>(boxes_2d.view(), quant_boxes);
1669        let mut ref_boxes: Vec<_> = Vec::with_capacity(10);
1670        let mut ref_masks: Vec<_> = Vec::with_capacity(10);
1671        decode_yolo_segdet_float(
1672            seg.view(),
1673            protos_f32_hwc.view(),
1674            score_threshold,
1675            iou_threshold,
1676            Some(configs::Nms::ClassAgnostic),
1677            &mut ref_boxes,
1678            &mut ref_masks,
1679        )
1680        .unwrap();
1681        assert_eq!(ref_boxes.len(), 2);
1682
1683        // ---- Config-driven path: NCHW protos declared in physical order ----
1684        // Permute the HWC test data to CHW memory order — this is what an
1685        // ONNX-style producer would emit into the tensor buffer.
1686        // `to_owned` materialises a C-contiguous Array3<f32> with CHW
1687        // strides, matching a producer that writes channels-outer.
1688        let protos_f32_chw_view = protos_f32_hwc.view().permuted_axes([2, 0, 1]); // [32, 160, 160]
1689        let protos_f32_chw = protos_f32_chw_view.to_owned();
1690        let protos_nchw = protos_f32_chw.insert_axis(ndarray::Axis(0)); // [1, 32, 160, 160]
1691
1692        // Build boxes as [1, 116, 8400] f32
1693        let seg_3d = seg.insert_axis(ndarray::Axis(0)); // [1, 116, 8400]
1694
1695        // Declare shape and dshape in physical memory order (outermost
1696        // first). `swap_axes_if_needed` uses dshape role lookup to
1697        // permute the stride tuple into canonical HWC.
1698        let decoder = DecoderBuilder::default()
1699            .with_config_yolo_segdet(
1700                configs::Detection {
1701                    decoder: configs::DecoderType::Ultralytics,
1702                    quantization: None,
1703                    shape: vec![1, 116, 8400],
1704                    dshape: vec![
1705                        (configs::DimName::Batch, 1),
1706                        (configs::DimName::NumFeatures, 116),
1707                        (configs::DimName::NumBoxes, 8400),
1708                    ],
1709                    normalized: Some(true),
1710                    anchors: None,
1711                },
1712                configs::Protos {
1713                    decoder: configs::DecoderType::Ultralytics,
1714                    quantization: None,
1715                    shape: vec![1, 32, 160, 160],
1716                    dshape: vec![
1717                        (configs::DimName::Batch, 1),
1718                        (configs::DimName::NumProtos, 32),
1719                        (configs::DimName::Height, 160),
1720                        (configs::DimName::Width, 160),
1721                    ],
1722                },
1723                None, // decoder version
1724            )
1725            .with_score_threshold(score_threshold)
1726            .with_iou_threshold(iou_threshold)
1727            .build()
1728            .unwrap();
1729
1730        let mut cfg_boxes: Vec<_> = Vec::with_capacity(10);
1731        let mut cfg_masks: Vec<_> = Vec::with_capacity(10);
1732        decoder
1733            .decode_float(
1734                &[seg_3d.view().into_dyn(), protos_nchw.view().into_dyn()],
1735                &mut cfg_boxes,
1736                &mut cfg_masks,
1737            )
1738            .unwrap();
1739
1740        // Must produce the same number of detections
1741        assert_eq!(
1742            cfg_boxes.len(),
1743            ref_boxes.len(),
1744            "config path produced {} boxes, reference produced {}",
1745            cfg_boxes.len(),
1746            ref_boxes.len()
1747        );
1748
1749        // Boxes must match
1750        for (i, (cb, rb)) in cfg_boxes.iter().zip(&ref_boxes).enumerate() {
1751            assert!(
1752                cb.equal_within_delta(rb, 0.01),
1753                "box {i} mismatch: config={cb:?}, reference={rb:?}"
1754            );
1755        }
1756
1757        // Masks must match pixel-for-pixel
1758        for (i, (cm, rm)) in cfg_masks.iter().zip(&ref_masks).enumerate() {
1759            let cm_arr = segmentation_to_mask(seg_view(cm).view()).unwrap();
1760            let rm_arr = segmentation_to_mask(seg_view(rm).view()).unwrap();
1761            assert_eq!(
1762                cm_arr, rm_arr,
1763                "mask {i} pixel mismatch between config-driven and reference paths"
1764            );
1765        }
1766    }
1767
1768    #[test]
1769    fn test_decoder_masks_i8() {
1770        let score_threshold = 0.45;
1771        let iou_threshold = 0.45;
1772        let boxes = load_yolov8_boxes();
1773        let quant_boxes = (0.021287761628627777, 31).into();
1774
1775        let protos = load_yolov8_protos();
1776        let quant_protos = (0.02491161972284317, -117).into();
1777        let mut output_boxes: Vec<_> = Vec::with_capacity(500);
1778        let mut output_masks: Vec<_> = Vec::with_capacity(500);
1779
1780        let decoder = DecoderBuilder::default()
1781            .with_config_yolo_segdet(
1782                configs::Detection {
1783                    decoder: configs::DecoderType::Ultralytics,
1784                    quantization: Some(quant_boxes),
1785                    shape: vec![1, 116, 8400],
1786                    anchors: None,
1787                    dshape: vec![
1788                        (DimName::Batch, 1),
1789                        (DimName::NumFeatures, 116),
1790                        (DimName::NumBoxes, 8400),
1791                    ],
1792                    normalized: Some(true),
1793                },
1794                Protos {
1795                    decoder: configs::DecoderType::Ultralytics,
1796                    quantization: Some(quant_protos),
1797                    shape: vec![1, 160, 160, 32],
1798                    dshape: vec![
1799                        (DimName::Batch, 1),
1800                        (DimName::Height, 160),
1801                        (DimName::Width, 160),
1802                        (DimName::NumProtos, 32),
1803                    ],
1804                },
1805                Some(DecoderVersion::Yolo11),
1806            )
1807            .with_score_threshold(score_threshold)
1808            .with_iou_threshold(iou_threshold)
1809            .build()
1810            .unwrap();
1811
1812        let quant_boxes = quant_boxes.into();
1813        let quant_protos = quant_protos.into();
1814
1815        decode_yolo_segdet_quant(
1816            (boxes.slice(s![0, .., ..]), quant_boxes),
1817            (protos.slice(s![0, .., .., ..]), quant_protos),
1818            score_threshold,
1819            iou_threshold,
1820            Some(configs::Nms::ClassAgnostic),
1821            &mut output_boxes,
1822            &mut output_masks,
1823        )
1824        .unwrap();
1825
1826        let mut output_boxes1: Vec<_> = Vec::with_capacity(500);
1827        let mut output_masks1: Vec<_> = Vec::with_capacity(500);
1828
1829        decoder
1830            .decode_quantized(
1831                &[boxes.view().into(), protos.view().into()],
1832                &mut output_boxes1,
1833                &mut output_masks1,
1834            )
1835            .unwrap();
1836
1837        let protos = dequantize_ndarray::<_, _, f32>(protos.view(), quant_protos);
1838        let seg = dequantize_ndarray::<_, _, f32>(boxes.view(), quant_boxes);
1839
1840        let mut output_boxes_f32: Vec<_> = Vec::with_capacity(500);
1841        let mut output_masks_f32: Vec<_> = Vec::with_capacity(500);
1842        decode_yolo_segdet_float(
1843            seg.slice(s![0, .., ..]),
1844            protos.slice(s![0, .., .., ..]),
1845            score_threshold,
1846            iou_threshold,
1847            Some(configs::Nms::ClassAgnostic),
1848            &mut output_boxes_f32,
1849            &mut output_masks_f32,
1850        )
1851        .unwrap();
1852
1853        let mut output_boxes1_f32: Vec<_> = Vec::with_capacity(500);
1854        let mut output_masks1_f32: Vec<_> = Vec::with_capacity(500);
1855
1856        decoder
1857            .decode_float(
1858                &[seg.view().into_dyn(), protos.view().into_dyn()],
1859                &mut output_boxes1_f32,
1860                &mut output_masks1_f32,
1861            )
1862            .unwrap();
1863
1864        compare_outputs(
1865            (&output_boxes, &output_boxes1),
1866            (&output_masks, &output_masks1),
1867        );
1868
1869        compare_outputs(
1870            (&output_boxes, &output_boxes_f32),
1871            (&output_masks, &output_masks_f32),
1872        );
1873
1874        compare_outputs(
1875            (&output_boxes_f32, &output_boxes1_f32),
1876            (&output_masks_f32, &output_masks1_f32),
1877        );
1878    }
1879
1880    #[test]
1881    fn test_decoder_yolo_split() {
1882        let score_threshold = 0.45;
1883        let iou_threshold = 0.45;
1884        let boxes = load_yolov8_boxes();
1885        let boxes: Vec<_> = boxes.iter().map(|x| *x as i16 * 256).collect();
1886        let boxes = ndarray::Array3::from_shape_vec((1, 116, 8400), boxes).unwrap();
1887
1888        let quant_boxes = Quantization::new(0.021287761628627777 / 256.0, 31 * 256);
1889
1890        let decoder = DecoderBuilder::default()
1891            .with_config_yolo_split_det(
1892                configs::Boxes {
1893                    decoder: configs::DecoderType::Ultralytics,
1894                    quantization: Some(QuantTuple(quant_boxes.scale, quant_boxes.zero_point)),
1895                    shape: vec![1, 4, 8400],
1896                    dshape: vec![
1897                        (DimName::Batch, 1),
1898                        (DimName::BoxCoords, 4),
1899                        (DimName::NumBoxes, 8400),
1900                    ],
1901                    normalized: Some(true),
1902                },
1903                configs::Scores {
1904                    decoder: configs::DecoderType::Ultralytics,
1905                    quantization: Some(QuantTuple(quant_boxes.scale, quant_boxes.zero_point)),
1906                    shape: vec![1, 80, 8400],
1907                    dshape: vec![
1908                        (DimName::Batch, 1),
1909                        (DimName::NumClasses, 80),
1910                        (DimName::NumBoxes, 8400),
1911                    ],
1912                },
1913            )
1914            .with_score_threshold(score_threshold)
1915            .with_iou_threshold(iou_threshold)
1916            .build()
1917            .unwrap();
1918
1919        let mut output_boxes: Vec<_> = Vec::with_capacity(500);
1920        let mut output_masks: Vec<_> = Vec::with_capacity(500);
1921
1922        decoder
1923            .decode_quantized(
1924                &[
1925                    boxes.slice(s![.., ..4, ..]).into(),
1926                    boxes.slice(s![.., 4..84, ..]).into(),
1927                ],
1928                &mut output_boxes,
1929                &mut output_masks,
1930            )
1931            .unwrap();
1932
1933        let seg = dequantize_ndarray::<_, _, f32>(boxes.view(), quant_boxes);
1934        let mut output_boxes_f32: Vec<_> = Vec::with_capacity(500);
1935        decode_yolo_det_float(
1936            seg.slice(s![0, ..84, ..]),
1937            score_threshold,
1938            iou_threshold,
1939            Some(configs::Nms::ClassAgnostic),
1940            &mut output_boxes_f32,
1941        );
1942
1943        let mut output_boxes1: Vec<_> = Vec::with_capacity(500);
1944        let mut output_masks1: Vec<_> = Vec::with_capacity(500);
1945
1946        decoder
1947            .decode_float(
1948                &[
1949                    seg.slice(s![.., ..4, ..]).into_dyn(),
1950                    seg.slice(s![.., 4..84, ..]).into_dyn(),
1951                ],
1952                &mut output_boxes1,
1953                &mut output_masks1,
1954            )
1955            .unwrap();
1956        compare_outputs((&output_boxes, &output_boxes_f32), (&output_masks, &[]));
1957        compare_outputs((&output_boxes_f32, &output_boxes1), (&[], &output_masks1));
1958    }
1959
1960    #[test]
1961    fn test_decoder_masks_config_mixed() {
1962        let score_threshold = 0.45;
1963        let iou_threshold = 0.45;
1964        let boxes_raw = load_yolov8_boxes();
1965        let boxes: Vec<_> = boxes_raw.iter().map(|x| *x as i16 * 256).collect();
1966        let boxes = ndarray::Array3::from_shape_vec((1, 116, 8400), boxes).unwrap();
1967
1968        let quant_boxes = (0.021287761628627777 / 256.0, 31 * 256);
1969
1970        let protos = load_yolov8_protos();
1971        let quant_protos = (0.02491161972284317, -117);
1972
1973        let decoder = build_yolo_split_segdet_decoder(
1974            score_threshold,
1975            iou_threshold,
1976            quant_boxes,
1977            quant_protos,
1978        );
1979        let mut output_boxes: Vec<_> = Vec::with_capacity(500);
1980        let mut output_masks: Vec<_> = Vec::with_capacity(500);
1981
1982        decoder
1983            .decode_quantized(
1984                &[
1985                    boxes.slice(s![.., ..4, ..]).into(),
1986                    boxes.slice(s![.., 4..84, ..]).into(),
1987                    boxes.slice(s![.., 84.., ..]).into(),
1988                    protos.view().into(),
1989                ],
1990                &mut output_boxes,
1991                &mut output_masks,
1992            )
1993            .unwrap();
1994
1995        let protos = dequantize_ndarray::<_, _, f32>(protos.view(), quant_protos.into());
1996        let seg = dequantize_ndarray::<_, _, f32>(boxes.view(), quant_boxes.into());
1997        let mut output_boxes_f32: Vec<_> = Vec::with_capacity(500);
1998        let mut output_masks_f32: Vec<_> = Vec::with_capacity(500);
1999        decode_yolo_segdet_float(
2000            seg.slice(s![0, .., ..]),
2001            protos.slice(s![0, .., .., ..]),
2002            score_threshold,
2003            iou_threshold,
2004            Some(configs::Nms::ClassAgnostic),
2005            &mut output_boxes_f32,
2006            &mut output_masks_f32,
2007        )
2008        .unwrap();
2009
2010        let mut output_boxes1: Vec<_> = Vec::with_capacity(500);
2011        let mut output_masks1: Vec<_> = Vec::with_capacity(500);
2012
2013        decoder
2014            .decode_float(
2015                &[
2016                    seg.slice(s![.., ..4, ..]).into_dyn(),
2017                    seg.slice(s![.., 4..84, ..]).into_dyn(),
2018                    seg.slice(s![.., 84.., ..]).into_dyn(),
2019                    protos.view().into_dyn(),
2020                ],
2021                &mut output_boxes1,
2022                &mut output_masks1,
2023            )
2024            .unwrap();
2025        compare_outputs(
2026            (&output_boxes, &output_boxes_f32),
2027            (&output_masks, &output_masks_f32),
2028        );
2029        compare_outputs(
2030            (&output_boxes_f32, &output_boxes1),
2031            (&output_masks_f32, &output_masks1),
2032        );
2033    }
2034
2035    fn build_yolo_split_segdet_decoder(
2036        score_threshold: f32,
2037        iou_threshold: f32,
2038        quant_boxes: (f32, i32),
2039        quant_protos: (f32, i32),
2040    ) -> crate::Decoder {
2041        DecoderBuilder::default()
2042            .with_config_yolo_split_segdet(
2043                configs::Boxes {
2044                    decoder: configs::DecoderType::Ultralytics,
2045                    quantization: Some(quant_boxes.into()),
2046                    shape: vec![1, 4, 8400],
2047                    dshape: vec![
2048                        (DimName::Batch, 1),
2049                        (DimName::BoxCoords, 4),
2050                        (DimName::NumBoxes, 8400),
2051                    ],
2052                    normalized: Some(true),
2053                },
2054                configs::Scores {
2055                    decoder: configs::DecoderType::Ultralytics,
2056                    quantization: Some(quant_boxes.into()),
2057                    shape: vec![1, 80, 8400],
2058                    dshape: vec![
2059                        (DimName::Batch, 1),
2060                        (DimName::NumClasses, 80),
2061                        (DimName::NumBoxes, 8400),
2062                    ],
2063                },
2064                configs::MaskCoefficients {
2065                    decoder: configs::DecoderType::Ultralytics,
2066                    quantization: Some(quant_boxes.into()),
2067                    shape: vec![1, 32, 8400],
2068                    dshape: vec![
2069                        (DimName::Batch, 1),
2070                        (DimName::NumProtos, 32),
2071                        (DimName::NumBoxes, 8400),
2072                    ],
2073                },
2074                configs::Protos {
2075                    decoder: configs::DecoderType::Ultralytics,
2076                    quantization: Some(quant_protos.into()),
2077                    shape: vec![1, 160, 160, 32],
2078                    dshape: vec![
2079                        (DimName::Batch, 1),
2080                        (DimName::Height, 160),
2081                        (DimName::Width, 160),
2082                        (DimName::NumProtos, 32),
2083                    ],
2084                },
2085            )
2086            .with_score_threshold(score_threshold)
2087            .with_iou_threshold(iou_threshold)
2088            .build()
2089            .unwrap()
2090    }
2091
2092    fn build_yolov8_seg_decoder(score_threshold: f32, iou_threshold: f32) -> crate::Decoder {
2093        let config_yaml = edgefirst_bench::testdata::read_to_string("yolov8_seg.yaml");
2094        DecoderBuilder::default()
2095            .with_config_yaml_str(config_yaml.to_string())
2096            .with_score_threshold(score_threshold)
2097            .with_iou_threshold(iou_threshold)
2098            .build()
2099            .unwrap()
2100    }
2101    #[test]
2102    fn test_decoder_masks_config_i32() {
2103        let score_threshold = 0.45;
2104        let iou_threshold = 0.45;
2105        let boxes_raw = load_yolov8_boxes();
2106        let scale = 1 << 23;
2107        let boxes: Vec<_> = boxes_raw.iter().map(|x| *x as i32 * scale).collect();
2108        let boxes = ndarray::Array3::from_shape_vec((1, 116, 8400), boxes).unwrap();
2109
2110        let quant_boxes = (0.021287761628627777 / scale as f32, 31 * scale);
2111
2112        let protos_raw = load_yolov8_protos();
2113        let protos: Vec<_> = protos_raw.iter().map(|x| *x as i32 * scale).collect();
2114        let protos = ndarray::Array4::from_shape_vec((1, 160, 160, 32), protos).unwrap();
2115        let quant_protos = (0.02491161972284317 / scale as f32, -117 * scale);
2116
2117        let decoder = build_yolo_split_segdet_decoder(
2118            score_threshold,
2119            iou_threshold,
2120            quant_boxes,
2121            quant_protos,
2122        );
2123
2124        let mut output_boxes: Vec<_> = Vec::with_capacity(500);
2125        let mut output_masks: Vec<_> = Vec::with_capacity(500);
2126
2127        decoder
2128            .decode_quantized(
2129                &[
2130                    boxes.slice(s![.., ..4, ..]).into(),
2131                    boxes.slice(s![.., 4..84, ..]).into(),
2132                    boxes.slice(s![.., 84.., ..]).into(),
2133                    protos.view().into(),
2134                ],
2135                &mut output_boxes,
2136                &mut output_masks,
2137            )
2138            .unwrap();
2139
2140        let protos = dequantize_ndarray::<_, _, f32>(protos.view(), quant_protos.into());
2141        let seg = dequantize_ndarray::<_, _, f32>(boxes.view(), quant_boxes.into());
2142        let mut output_boxes_f32: Vec<_> = Vec::with_capacity(500);
2143        let mut output_masks_f32: Vec<Segmentation> = Vec::with_capacity(500);
2144        decode_yolo_segdet_float(
2145            seg.slice(s![0, .., ..]),
2146            protos.slice(s![0, .., .., ..]),
2147            score_threshold,
2148            iou_threshold,
2149            Some(configs::Nms::ClassAgnostic),
2150            &mut output_boxes_f32,
2151            &mut output_masks_f32,
2152        )
2153        .unwrap();
2154
2155        assert_eq!(output_boxes.len(), output_boxes_f32.len());
2156        assert_eq!(output_masks.len(), output_masks_f32.len());
2157
2158        compare_outputs(
2159            (&output_boxes, &output_boxes_f32),
2160            (&output_masks, &output_masks_f32),
2161        );
2162    }
2163
2164    /// test running multiple decoders concurrently
2165    #[test]
2166    fn test_context_switch() {
2167        let yolo_det = || {
2168            let score_threshold = 0.25;
2169            let iou_threshold = 0.7;
2170            let out = load_yolov8s_det();
2171            let quant = (0.0040811873, -123).into();
2172
2173            let decoder = DecoderBuilder::default()
2174                .with_config_yolo_det(
2175                    configs::Detection {
2176                        decoder: DecoderType::Ultralytics,
2177                        shape: vec![1, 84, 8400],
2178                        anchors: None,
2179                        quantization: Some(quant),
2180                        dshape: vec![
2181                            (DimName::Batch, 1),
2182                            (DimName::NumFeatures, 84),
2183                            (DimName::NumBoxes, 8400),
2184                        ],
2185                        normalized: None,
2186                    },
2187                    None,
2188                )
2189                .with_score_threshold(score_threshold)
2190                .with_iou_threshold(iou_threshold)
2191                .build()
2192                .unwrap();
2193
2194            let mut output_boxes: Vec<_> = Vec::with_capacity(50);
2195            let mut output_masks: Vec<_> = Vec::with_capacity(50);
2196
2197            for _ in 0..100 {
2198                decoder
2199                    .decode_quantized(&[out.view().into()], &mut output_boxes, &mut output_masks)
2200                    .unwrap();
2201
2202                assert!(output_boxes[0].equal_within_delta(
2203                    &DetectBox {
2204                        bbox: BoundingBox {
2205                            xmin: 0.5285137,
2206                            ymin: 0.05305544,
2207                            xmax: 0.87541467,
2208                            ymax: 0.9998909,
2209                        },
2210                        score: 0.5591227,
2211                        label: 0
2212                    },
2213                    1e-6
2214                ));
2215
2216                assert!(output_boxes[1].equal_within_delta(
2217                    &DetectBox {
2218                        bbox: BoundingBox {
2219                            xmin: 0.130598,
2220                            ymin: 0.43260583,
2221                            xmax: 0.35098213,
2222                            ymax: 0.9958097,
2223                        },
2224                        score: 0.33057618,
2225                        label: 75
2226                    },
2227                    1e-6
2228                ));
2229                assert!(output_masks.is_empty());
2230            }
2231        };
2232
2233        let modelpack_det_split = || {
2234            let score_threshold = 0.8;
2235            let iou_threshold = 0.5;
2236
2237            let seg = edgefirst_bench::testdata::read("modelpack_seg_2x160x160.bin");
2238            let seg = ndarray::Array4::from_shape_vec((1, 2, 160, 160), seg.to_vec()).unwrap();
2239
2240            let detect0 = edgefirst_bench::testdata::read("modelpack_split_9x15x18.bin");
2241            let detect0 =
2242                ndarray::Array4::from_shape_vec((1, 9, 15, 18), detect0.to_vec()).unwrap();
2243
2244            let detect1 = edgefirst_bench::testdata::read("modelpack_split_17x30x18.bin");
2245            let detect1 =
2246                ndarray::Array4::from_shape_vec((1, 17, 30, 18), detect1.to_vec()).unwrap();
2247
2248            let mut mask = seg.slice(s![0, .., .., ..]);
2249            mask.swap_axes(0, 1);
2250            mask.swap_axes(1, 2);
2251            let mask = [Segmentation {
2252                xmin: 0.0,
2253                ymin: 0.0,
2254                xmax: 1.0,
2255                ymax: 1.0,
2256                segmentation: crate::mask_to_tensor(mask.view()).unwrap(),
2257            }];
2258            let correct_boxes = [DetectBox {
2259                bbox: BoundingBox {
2260                    xmin: 0.43171933,
2261                    ymin: 0.68243736,
2262                    xmax: 0.5626645,
2263                    ymax: 0.808863,
2264                },
2265                score: 0.99240804,
2266                label: 0,
2267            }];
2268
2269            let quant0 = (0.08547406643629074, 174).into();
2270            let quant1 = (0.09929127991199493, 183).into();
2271            let quant_seg = (1.0 / 255.0, 0).into();
2272
2273            let anchors0 = vec![
2274                [0.36666667461395264, 0.31481480598449707],
2275                [0.38749998807907104, 0.4740740656852722],
2276                [0.5333333611488342, 0.644444465637207],
2277            ];
2278            let anchors1 = vec![
2279                [0.13750000298023224, 0.2074074000120163],
2280                [0.2541666626930237, 0.21481481194496155],
2281                [0.23125000298023224, 0.35185185074806213],
2282            ];
2283
2284            let decoder = DecoderBuilder::default()
2285                .with_config_modelpack_segdet_split(
2286                    vec![
2287                        configs::Detection {
2288                            decoder: DecoderType::ModelPack,
2289                            shape: vec![1, 17, 30, 18],
2290                            anchors: Some(anchors1),
2291                            quantization: Some(quant1),
2292                            dshape: vec![
2293                                (DimName::Batch, 1),
2294                                (DimName::Height, 17),
2295                                (DimName::Width, 30),
2296                                (DimName::NumAnchorsXFeatures, 18),
2297                            ],
2298                            normalized: None,
2299                        },
2300                        configs::Detection {
2301                            decoder: DecoderType::ModelPack,
2302                            shape: vec![1, 9, 15, 18],
2303                            anchors: Some(anchors0),
2304                            quantization: Some(quant0),
2305                            dshape: vec![
2306                                (DimName::Batch, 1),
2307                                (DimName::Height, 9),
2308                                (DimName::Width, 15),
2309                                (DimName::NumAnchorsXFeatures, 18),
2310                            ],
2311                            normalized: None,
2312                        },
2313                    ],
2314                    configs::Segmentation {
2315                        decoder: DecoderType::ModelPack,
2316                        quantization: Some(quant_seg),
2317                        shape: vec![1, 2, 160, 160],
2318                        dshape: vec![
2319                            (DimName::Batch, 1),
2320                            (DimName::NumClasses, 2),
2321                            (DimName::Height, 160),
2322                            (DimName::Width, 160),
2323                        ],
2324                    },
2325                )
2326                .with_score_threshold(score_threshold)
2327                .with_iou_threshold(iou_threshold)
2328                .build()
2329                .unwrap();
2330            let mut output_boxes: Vec<_> = Vec::with_capacity(10);
2331            let mut output_masks: Vec<_> = Vec::with_capacity(10);
2332
2333            for _ in 0..100 {
2334                decoder
2335                    .decode_quantized(
2336                        &[
2337                            detect0.view().into(),
2338                            detect1.view().into(),
2339                            seg.view().into(),
2340                        ],
2341                        &mut output_boxes,
2342                        &mut output_masks,
2343                    )
2344                    .unwrap();
2345
2346                compare_outputs((&correct_boxes, &output_boxes), (&mask, &output_masks));
2347            }
2348        };
2349
2350        let handles = vec![
2351            std::thread::spawn(yolo_det),
2352            std::thread::spawn(modelpack_det_split),
2353            std::thread::spawn(yolo_det),
2354            std::thread::spawn(modelpack_det_split),
2355            std::thread::spawn(yolo_det),
2356            std::thread::spawn(modelpack_det_split),
2357            std::thread::spawn(yolo_det),
2358            std::thread::spawn(modelpack_det_split),
2359        ];
2360        for handle in handles {
2361            handle.join().unwrap();
2362        }
2363    }
2364
2365    #[test]
2366    fn test_ndarray_to_xyxy_float() {
2367        let arr = array![10.0_f32, 20.0, 20.0, 20.0];
2368        let xyxy: [f32; 4] = XYWH::ndarray_to_xyxy_float(arr.view());
2369        assert_eq!(xyxy, [0.0_f32, 10.0, 20.0, 30.0]);
2370
2371        let arr = array![10.0_f32, 20.0, 20.0, 20.0];
2372        let xyxy: [f32; 4] = XYXY::ndarray_to_xyxy_float(arr.view());
2373        assert_eq!(xyxy, [10.0_f32, 20.0, 20.0, 20.0]);
2374    }
2375
2376    /// Covers the slice conversion directly. It is the single implementation
2377    /// per type now that the hand-inlined `ndarray_to_xyxy_float` overrides
2378    /// are gone, so the float postprocessors reach it through the trait
2379    /// default; while those overrides existed nothing reached it at all, which
2380    /// is why its twenty mutants survived.
2381    ///
2382    /// Centre 10,20 with a 6x8 box gives a distinct value in all four slots,
2383    /// so no swapped operator lands on the same answer.
2384    #[test]
2385    fn to_xyxy_float_converts_centre_and_size_to_corners() {
2386        let xyxy: [f32; 4] = XYWH::to_xyxy_float(&[10.0_f32, 20.0, 6.0, 8.0]);
2387        assert_eq!(xyxy, [7.0_f32, 16.0, 13.0, 24.0]);
2388
2389        let xyxy: [f32; 4] = XYXY::to_xyxy_float(&[7.0_f32, 16.0, 13.0, 24.0]);
2390        assert_eq!(xyxy, [7.0_f32, 16.0, 13.0, 24.0]);
2391    }
2392
2393    /// `arg_max` documents that the last index wins on ties, which is the
2394    /// semantics `arg_max_i8` mirrors on aarch64. A strict `>` is what makes
2395    /// the fold take the later index; `>=` would keep the earlier one.
2396    #[test]
2397    fn arg_max_returns_the_last_index_on_a_tie() {
2398        let scores = array![1_i32, 5, 5, 2];
2399        assert_eq!(crate::arg_max(scores.view()), (5, 2));
2400
2401        let single = array![7_i32];
2402        assert_eq!(crate::arg_max(single.view()), (7, 0));
2403    }
2404
2405    #[test]
2406    fn test_class_aware_nms_float() {
2407        use crate::float::nms_class_aware_float;
2408
2409        // Create two overlapping boxes with different classes
2410        let boxes = vec![
2411            DetectBox {
2412                bbox: BoundingBox {
2413                    xmin: 0.0,
2414                    ymin: 0.0,
2415                    xmax: 0.5,
2416                    ymax: 0.5,
2417                },
2418                score: 0.9,
2419                label: 0, // class 0
2420            },
2421            DetectBox {
2422                bbox: BoundingBox {
2423                    xmin: 0.1,
2424                    ymin: 0.1,
2425                    xmax: 0.6,
2426                    ymax: 0.6,
2427                },
2428                score: 0.8,
2429                label: 1, // class 1 - different class
2430            },
2431        ];
2432
2433        // Class-aware NMS should keep both boxes (different classes, IoU ~0.47 >
2434        // threshold 0.3)
2435        let result = nms_class_aware_float(0.3, None, boxes.clone());
2436        assert_eq!(
2437            result.len(),
2438            2,
2439            "Class-aware NMS should keep both boxes with different classes"
2440        );
2441
2442        // Now test with same class - should suppress one
2443        let same_class_boxes = vec![
2444            DetectBox {
2445                bbox: BoundingBox {
2446                    xmin: 0.0,
2447                    ymin: 0.0,
2448                    xmax: 0.5,
2449                    ymax: 0.5,
2450                },
2451                score: 0.9,
2452                label: 0,
2453            },
2454            DetectBox {
2455                bbox: BoundingBox {
2456                    xmin: 0.1,
2457                    ymin: 0.1,
2458                    xmax: 0.6,
2459                    ymax: 0.6,
2460                },
2461                score: 0.8,
2462                label: 0, // same class
2463            },
2464        ];
2465
2466        let result = nms_class_aware_float(0.3, None, same_class_boxes);
2467        assert_eq!(
2468            result.len(),
2469            1,
2470            "Class-aware NMS should suppress overlapping box with same class"
2471        );
2472        assert_eq!(result[0].label, 0);
2473        assert!((result[0].score - 0.9).abs() < 1e-6);
2474    }
2475
2476    #[test]
2477    fn test_class_agnostic_vs_aware_nms() {
2478        use crate::float::{nms_class_aware_float, nms_float};
2479
2480        // Two overlapping boxes with different classes
2481        let boxes = vec![
2482            DetectBox {
2483                bbox: BoundingBox {
2484                    xmin: 0.0,
2485                    ymin: 0.0,
2486                    xmax: 0.5,
2487                    ymax: 0.5,
2488                },
2489                score: 0.9,
2490                label: 0,
2491            },
2492            DetectBox {
2493                bbox: BoundingBox {
2494                    xmin: 0.1,
2495                    ymin: 0.1,
2496                    xmax: 0.6,
2497                    ymax: 0.6,
2498                },
2499                score: 0.8,
2500                label: 1,
2501            },
2502        ];
2503
2504        // Class-agnostic should suppress one (IoU ~0.47 > threshold 0.3)
2505        let agnostic_result = nms_float(0.3, None, boxes.clone());
2506        assert_eq!(
2507            agnostic_result.len(),
2508            1,
2509            "Class-agnostic NMS should suppress overlapping boxes"
2510        );
2511
2512        // Class-aware should keep both (different classes)
2513        let aware_result = nms_class_aware_float(0.3, None, boxes);
2514        assert_eq!(
2515            aware_result.len(),
2516            2,
2517            "Class-aware NMS should keep boxes with different classes"
2518        );
2519    }
2520
2521    #[test]
2522    fn test_class_aware_nms_int() {
2523        use crate::byte::nms_class_aware_int;
2524
2525        // Create two overlapping boxes with different classes
2526        let boxes = vec![
2527            DetectBoxQuantized {
2528                bbox: BoundingBox {
2529                    xmin: 0.0,
2530                    ymin: 0.0,
2531                    xmax: 0.5,
2532                    ymax: 0.5,
2533                },
2534                score: 200_u8,
2535                label: 0,
2536            },
2537            DetectBoxQuantized {
2538                bbox: BoundingBox {
2539                    xmin: 0.1,
2540                    ymin: 0.1,
2541                    xmax: 0.6,
2542                    ymax: 0.6,
2543                },
2544                score: 180_u8,
2545                label: 1, // different class
2546            },
2547        ];
2548
2549        // Should keep both (different classes)
2550        let result = nms_class_aware_int(0.5, None, boxes);
2551        assert_eq!(
2552            result.len(),
2553            2,
2554            "Class-aware NMS (int) should keep boxes with different classes"
2555        );
2556    }
2557
2558    #[test]
2559    fn test_nms_enum_default() {
2560        // Test that Nms enum has the correct default
2561        let default_nms: configs::Nms = Default::default();
2562        assert_eq!(default_nms, configs::Nms::ClassAgnostic);
2563    }
2564
2565    #[test]
2566    fn test_decoder_nms_mode() {
2567        // Test that decoder properly stores NMS mode
2568        let decoder = DecoderBuilder::default()
2569            .with_config_yolo_det(
2570                configs::Detection {
2571                    anchors: None,
2572                    decoder: DecoderType::Ultralytics,
2573                    quantization: None,
2574                    shape: vec![1, 84, 8400],
2575                    dshape: Vec::new(),
2576                    normalized: Some(true),
2577                },
2578                None,
2579            )
2580            .with_nms(Some(configs::Nms::ClassAware))
2581            .build()
2582            .unwrap();
2583
2584        assert_eq!(decoder.nms, Some(configs::Nms::ClassAware));
2585    }
2586
2587    #[test]
2588    fn test_decoder_nms_bypass() {
2589        // Test that decoder can be configured with nms=None (bypass)
2590        let decoder = DecoderBuilder::default()
2591            .with_config_yolo_det(
2592                configs::Detection {
2593                    anchors: None,
2594                    decoder: DecoderType::Ultralytics,
2595                    quantization: None,
2596                    shape: vec![1, 84, 8400],
2597                    dshape: Vec::new(),
2598                    normalized: Some(true),
2599                },
2600                None,
2601            )
2602            .with_nms(None)
2603            .build()
2604            .unwrap();
2605
2606        assert_eq!(decoder.nms, None);
2607    }
2608
2609    #[test]
2610    fn test_decoder_normalized_boxes_true() {
2611        // Test that normalized_boxes returns Some(true) when explicitly set
2612        let decoder = DecoderBuilder::default()
2613            .with_config_yolo_det(
2614                configs::Detection {
2615                    anchors: None,
2616                    decoder: DecoderType::Ultralytics,
2617                    quantization: None,
2618                    shape: vec![1, 84, 8400],
2619                    dshape: Vec::new(),
2620                    normalized: Some(true),
2621                },
2622                None,
2623            )
2624            .build()
2625            .unwrap();
2626
2627        assert_eq!(decoder.normalized_boxes(), Some(true));
2628    }
2629
2630    #[test]
2631    fn test_decoder_normalized_boxes_false() {
2632        // Test that normalized_boxes returns Some(false) when config specifies
2633        // unnormalized
2634        let decoder = DecoderBuilder::default()
2635            .with_config_yolo_det(
2636                configs::Detection {
2637                    anchors: None,
2638                    decoder: DecoderType::Ultralytics,
2639                    quantization: None,
2640                    shape: vec![1, 84, 8400],
2641                    dshape: Vec::new(),
2642                    normalized: Some(false),
2643                },
2644                None,
2645            )
2646            .build()
2647            .unwrap();
2648
2649        assert_eq!(decoder.normalized_boxes(), Some(false));
2650    }
2651
2652    #[test]
2653    fn test_decoder_normalized_boxes_unknown() {
2654        // Test that normalized_boxes returns None when not specified in config
2655        let decoder = DecoderBuilder::default()
2656            .with_config_yolo_det(
2657                configs::Detection {
2658                    anchors: None,
2659                    decoder: DecoderType::Ultralytics,
2660                    quantization: None,
2661                    shape: vec![1, 84, 8400],
2662                    dshape: Vec::new(),
2663                    normalized: None,
2664                },
2665                Some(DecoderVersion::Yolo11),
2666            )
2667            .build()
2668            .unwrap();
2669
2670        assert_eq!(decoder.normalized_boxes(), None);
2671    }
2672
2673    pub fn quantize_ndarray<T: PrimInt + 'static, D: Dimension, F: Float + AsPrimitive<T>>(
2674        input: ArrayView<F, D>,
2675        quant: Quantization,
2676    ) -> Array<T, D>
2677    where
2678        i32: num_traits::AsPrimitive<F>,
2679        f32: num_traits::AsPrimitive<F>,
2680    {
2681        let zero_point = quant.zero_point.as_();
2682        let div_scale = F::one() / quant.scale.as_();
2683        if zero_point != F::zero() {
2684            input.mapv(|d| (d * div_scale + zero_point).round().as_())
2685        } else {
2686            input.mapv(|d| (d * div_scale).round().as_())
2687        }
2688    }
2689
2690    fn real_data_expected_boxes() -> [DetectBox; 2] {
2691        [
2692            DetectBox {
2693                bbox: BoundingBox {
2694                    xmin: 0.08515105,
2695                    ymin: 0.7131401,
2696                    xmax: 0.29802868,
2697                    ymax: 0.8195788,
2698                },
2699                score: 0.91537374,
2700                label: 23,
2701            },
2702            DetectBox {
2703                bbox: BoundingBox {
2704                    xmin: 0.59605736,
2705                    ymin: 0.25545314,
2706                    xmax: 0.93666154,
2707                    ymax: 0.72378385,
2708                },
2709                score: 0.91537374,
2710                label: 23,
2711            },
2712        ]
2713    }
2714
2715    fn e2e_expected_boxes_quant() -> [DetectBox; 1] {
2716        [DetectBox {
2717            bbox: BoundingBox {
2718                xmin: 0.12549022,
2719                ymin: 0.12549022,
2720                xmax: 0.23529413,
2721                ymax: 0.23529413,
2722            },
2723            score: 0.98823535,
2724            label: 2,
2725        }]
2726    }
2727
2728    fn e2e_expected_boxes_float() -> [DetectBox; 1] {
2729        [DetectBox {
2730            bbox: BoundingBox {
2731                xmin: 0.1234,
2732                ymin: 0.1234,
2733                xmax: 0.2345,
2734                ymax: 0.2345,
2735            },
2736            score: 0.9876,
2737            label: 2,
2738        }]
2739    }
2740
2741    macro_rules! real_data_proto_test {
2742        ($name:ident, quantized, $layout:ident) => {
2743            #[test]
2744            fn $name() {
2745                let is_split = matches!(stringify!($layout), "split");
2746
2747                let score_threshold = 0.45;
2748                let iou_threshold = 0.45;
2749                let quant_boxes = (0.021287762_f32, 31_i32);
2750                let quant_protos = (0.02491162_f32, -117_i32);
2751
2752                let raw_boxes = edgefirst_bench::testdata::read("yolov8_boxes_116x8400.bin");
2753                let raw_boxes = unsafe {
2754                    std::slice::from_raw_parts(raw_boxes.as_ptr() as *const i8, raw_boxes.len())
2755                };
2756                let boxes_i8 =
2757                    ndarray::Array3::from_shape_vec((1, 116, 8400), raw_boxes.to_vec()).unwrap();
2758
2759                let raw_protos = edgefirst_bench::testdata::read("yolov8_protos_160x160x32.bin");
2760                let raw_protos = unsafe {
2761                    std::slice::from_raw_parts(raw_protos.as_ptr() as *const i8, raw_protos.len())
2762                };
2763                let protos_i8 =
2764                    ndarray::Array4::from_shape_vec((1, 160, 160, 32), raw_protos.to_vec())
2765                        .unwrap();
2766
2767                // Pre-split (unused for combined, but harmless)
2768                let mask_split = boxes_i8.slice(s![.., 84.., ..]).to_owned();
2769                let scores_split = boxes_i8.slice(s![.., 4..84, ..]).to_owned();
2770                let boxes_split = boxes_i8.slice(s![.., ..4, ..]).to_owned();
2771                let boxes_combined = boxes_i8;
2772
2773                let decoder = if is_split {
2774                    build_yolo_split_segdet_decoder(
2775                        score_threshold,
2776                        iou_threshold,
2777                        quant_boxes,
2778                        quant_protos,
2779                    )
2780                } else {
2781                    build_yolov8_seg_decoder(score_threshold, iou_threshold)
2782                };
2783
2784                let expected = real_data_expected_boxes();
2785                let mut output_boxes = Vec::with_capacity(50);
2786
2787                let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> = if is_split {
2788                    vec![
2789                        boxes_split.view().into(),
2790                        scores_split.view().into(),
2791                        mask_split.view().into(),
2792                        protos_i8.view().into(),
2793                    ]
2794                } else {
2795                    vec![boxes_combined.view().into(), protos_i8.view().into()]
2796                };
2797                decoder
2798                    .decode_quantized_proto(&inputs, &mut output_boxes)
2799                    .unwrap();
2800
2801                assert_eq!(output_boxes.len(), 2);
2802                assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
2803                assert!(output_boxes[1].equal_within_delta(&expected[1], 1.0 / 160.0));
2804            }
2805        };
2806        ($name:ident, float, $layout:ident) => {
2807            #[test]
2808            fn $name() {
2809                let is_split = matches!(stringify!($layout), "split");
2810
2811                let score_threshold = 0.45;
2812                let iou_threshold = 0.45;
2813                let quant_boxes = (0.021287762_f32, 31_i32);
2814                let quant_protos = (0.02491162_f32, -117_i32);
2815
2816                let raw_boxes = edgefirst_bench::testdata::read("yolov8_boxes_116x8400.bin");
2817                let raw_boxes = unsafe {
2818                    std::slice::from_raw_parts(raw_boxes.as_ptr() as *const i8, raw_boxes.len())
2819                };
2820                let boxes_i8 =
2821                    ndarray::Array3::from_shape_vec((1, 116, 8400), raw_boxes.to_vec()).unwrap();
2822                let boxes_f32: Array3<f32> =
2823                    dequantize_ndarray(boxes_i8.view(), quant_boxes.into());
2824
2825                let raw_protos = edgefirst_bench::testdata::read("yolov8_protos_160x160x32.bin");
2826                let raw_protos = unsafe {
2827                    std::slice::from_raw_parts(raw_protos.as_ptr() as *const i8, raw_protos.len())
2828                };
2829                let protos_i8 =
2830                    ndarray::Array4::from_shape_vec((1, 160, 160, 32), raw_protos.to_vec())
2831                        .unwrap();
2832                let protos_f32: Array4<f32> =
2833                    dequantize_ndarray(protos_i8.view(), quant_protos.into());
2834
2835                // Pre-split from dequantized data
2836                let mask_split = boxes_f32.slice(s![.., 84.., ..]).to_owned();
2837                let scores_split = boxes_f32.slice(s![.., 4..84, ..]).to_owned();
2838                let boxes_split = boxes_f32.slice(s![.., ..4, ..]).to_owned();
2839                let boxes_combined = boxes_f32;
2840
2841                let decoder = if is_split {
2842                    build_yolo_split_segdet_decoder(
2843                        score_threshold,
2844                        iou_threshold,
2845                        quant_boxes,
2846                        quant_protos,
2847                    )
2848                } else {
2849                    build_yolov8_seg_decoder(score_threshold, iou_threshold)
2850                };
2851
2852                let expected = real_data_expected_boxes();
2853                let mut output_boxes = Vec::with_capacity(50);
2854
2855                let inputs = if is_split {
2856                    vec![
2857                        boxes_split.view().into_dyn(),
2858                        scores_split.view().into_dyn(),
2859                        mask_split.view().into_dyn(),
2860                        protos_f32.view().into_dyn(),
2861                    ]
2862                } else {
2863                    vec![
2864                        boxes_combined.view().into_dyn(),
2865                        protos_f32.view().into_dyn(),
2866                    ]
2867                };
2868                decoder
2869                    .decode_float_proto(&inputs, &mut output_boxes)
2870                    .unwrap();
2871
2872                assert_eq!(output_boxes.len(), 2);
2873                assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
2874                assert!(output_boxes[1].equal_within_delta(&expected[1], 1.0 / 160.0));
2875            }
2876        };
2877    }
2878
2879    real_data_proto_test!(test_decoder_segdet_proto, quantized, combined);
2880    real_data_proto_test!(test_decoder_segdet_proto_float, float, combined);
2881    real_data_proto_test!(test_decoder_segdet_split_proto, quantized, split);
2882    real_data_proto_test!(test_decoder_segdet_split_proto_float, float, split);
2883
2884    const E2E_COMBINED_DET_CONFIG: &str = "
2885decoder_version: yolo26
2886outputs:
2887 - type: detection
2888   decoder: ultralytics
2889   quantization: [0.00784313725490196, 0]
2890   shape: [1, 10, 6]
2891   dshape:
2892    - [batch, 1]
2893    - [num_boxes, 10]
2894    - [num_features, 6]
2895   normalized: true
2896";
2897
2898    const E2E_COMBINED_SEGDET_CONFIG: &str = "
2899decoder_version: yolo26
2900outputs:
2901 - type: detection
2902   decoder: ultralytics
2903   quantization: [0.00784313725490196, 0]
2904   shape: [1, 10, 38]
2905   dshape:
2906    - [batch, 1]
2907    - [num_boxes, 10]
2908    - [num_features, 38]
2909   normalized: true
2910 - type: protos
2911   decoder: ultralytics
2912   quantization: [0.0039215686274509803921568627451, 128]
2913   shape: [1, 160, 160, 32]
2914   dshape:
2915    - [batch, 1]
2916    - [height, 160]
2917    - [width, 160]
2918    - [num_protos, 32]
2919";
2920
2921    const E2E_SPLIT_DET_CONFIG: &str = "
2922decoder_version: yolo26
2923outputs:
2924 - type: boxes
2925   decoder: ultralytics
2926   quantization: [0.00784313725490196, 0]
2927   shape: [1, 10, 4]
2928   dshape:
2929    - [batch, 1]
2930    - [num_boxes, 10]
2931    - [box_coords, 4]
2932   normalized: true
2933 - type: scores
2934   decoder: ultralytics
2935   quantization: [0.00784313725490196, 0]
2936   shape: [1, 10, 1]
2937   dshape:
2938    - [batch, 1]
2939    - [num_boxes, 10]
2940    - [num_classes, 1]
2941 - type: classes
2942   decoder: ultralytics
2943   quantization: [0.00784313725490196, 0]
2944   shape: [1, 10, 1]
2945   dshape:
2946    - [batch, 1]
2947    - [num_boxes, 10]
2948    - [num_classes, 1]
2949";
2950
2951    const E2E_SPLIT_SEGDET_CONFIG: &str = "
2952decoder_version: yolo26
2953outputs:
2954 - type: boxes
2955   decoder: ultralytics
2956   quantization: [0.00784313725490196, 0]
2957   shape: [1, 10, 4]
2958   dshape:
2959    - [batch, 1]
2960    - [num_boxes, 10]
2961    - [box_coords, 4]
2962   normalized: true
2963 - type: scores
2964   decoder: ultralytics
2965   quantization: [0.00784313725490196, 0]
2966   shape: [1, 10, 1]
2967   dshape:
2968    - [batch, 1]
2969    - [num_boxes, 10]
2970    - [num_classes, 1]
2971 - type: classes
2972   decoder: ultralytics
2973   quantization: [0.00784313725490196, 0]
2974   shape: [1, 10, 1]
2975   dshape:
2976    - [batch, 1]
2977    - [num_boxes, 10]
2978    - [num_classes, 1]
2979 - type: mask_coefficients
2980   decoder: ultralytics
2981   quantization: [0.00784313725490196, 0]
2982   shape: [1, 10, 32]
2983   dshape:
2984    - [batch, 1]
2985    - [num_boxes, 10]
2986    - [num_protos, 32]
2987 - type: protos
2988   decoder: ultralytics
2989   quantization: [0.0039215686274509803921568627451, 128]
2990   shape: [1, 160, 160, 32]
2991   dshape:
2992    - [batch, 1]
2993    - [height, 160]
2994    - [width, 160]
2995    - [num_protos, 32]
2996";
2997
2998    macro_rules! e2e_segdet_test {
2999        ($name:ident, quantized, $layout:ident, $output:ident) => {
3000            #[test]
3001            fn $name() {
3002                let is_split = matches!(stringify!($layout), "split");
3003                let is_proto = matches!(stringify!($output), "proto");
3004
3005                let score_threshold = 0.45;
3006                let iou_threshold = 0.45;
3007
3008                let mut boxes = Array2::zeros((10, 4));
3009                let mut scores = Array2::zeros((10, 1));
3010                let mut classes = Array2::zeros((10, 1));
3011                let mask = Array2::zeros((10, 32));
3012                let protos = Array3::<f64>::zeros((160, 160, 32));
3013                let protos = protos.insert_axis(Axis(0));
3014                let protos_quant = (1.0 / 255.0, 0.0);
3015                let protos: Array4<u8> = quantize_ndarray(protos.view(), protos_quant.into());
3016
3017                boxes
3018                    .slice_mut(s![0, ..])
3019                    .assign(&array![0.1234, 0.1234, 0.2345, 0.2345]);
3020                scores.slice_mut(s![0, ..]).assign(&array![0.9876]);
3021                classes.slice_mut(s![0, ..]).assign(&array![2.0]);
3022
3023                let detect_quant = (2.0 / 255.0, 0.0);
3024
3025                let decoder = if is_split {
3026                    DecoderBuilder::default()
3027                        .with_config_yaml_str(E2E_SPLIT_SEGDET_CONFIG.to_string())
3028                        .with_score_threshold(score_threshold)
3029                        .with_iou_threshold(iou_threshold)
3030                        .build()
3031                        .unwrap()
3032                } else {
3033                    DecoderBuilder::default()
3034                        .with_config_yaml_str(E2E_COMBINED_SEGDET_CONFIG.to_string())
3035                        .with_score_threshold(score_threshold)
3036                        .with_iou_threshold(iou_threshold)
3037                        .build()
3038                        .unwrap()
3039                };
3040
3041                let expected = e2e_expected_boxes_quant();
3042                let mut output_boxes = Vec::with_capacity(50);
3043
3044                if is_split {
3045                    let boxes = boxes.insert_axis(Axis(0));
3046                    let scores = scores.insert_axis(Axis(0));
3047                    let classes = classes.insert_axis(Axis(0));
3048                    let mask = mask.insert_axis(Axis(0));
3049
3050                    let boxes: Array3<u8> = quantize_ndarray(boxes.view(), detect_quant.into());
3051                    let scores: Array3<u8> = quantize_ndarray(scores.view(), detect_quant.into());
3052                    let classes: Array3<u8> = quantize_ndarray(classes.view(), detect_quant.into());
3053                    let mask: Array3<u8> = quantize_ndarray(mask.view(), detect_quant.into());
3054
3055                    if is_proto {
3056                        let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> = vec![
3057                            boxes.view().into(),
3058                            scores.view().into(),
3059                            classes.view().into(),
3060                            mask.view().into(),
3061                            protos.view().into(),
3062                        ];
3063                        decoder
3064                            .decode_quantized_proto(&inputs, &mut output_boxes)
3065                            .unwrap();
3066
3067                        assert_eq!(output_boxes.len(), 1);
3068                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
3069                    } else {
3070                        let mut output_masks = Vec::with_capacity(50);
3071                        let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> = vec![
3072                            boxes.view().into(),
3073                            scores.view().into(),
3074                            classes.view().into(),
3075                            mask.view().into(),
3076                            protos.view().into(),
3077                        ];
3078                        decoder
3079                            .decode_quantized(&inputs, &mut output_boxes, &mut output_masks)
3080                            .unwrap();
3081
3082                        assert_eq!(output_boxes.len(), 1);
3083                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
3084                    }
3085                } else {
3086                    // Combined layout
3087                    let detect = ndarray::concatenate![
3088                        Axis(1),
3089                        boxes.view(),
3090                        scores.view(),
3091                        classes.view(),
3092                        mask.view()
3093                    ];
3094                    let detect = detect.insert_axis(Axis(0));
3095                    assert_eq!(detect.shape(), &[1, 10, 38]);
3096                    let detect: Array3<u8> = quantize_ndarray(detect.view(), detect_quant.into());
3097
3098                    if is_proto {
3099                        let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> =
3100                            vec![detect.view().into(), protos.view().into()];
3101                        decoder
3102                            .decode_quantized_proto(&inputs, &mut output_boxes)
3103                            .unwrap();
3104
3105                        assert_eq!(output_boxes.len(), 1);
3106                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
3107                    } else {
3108                        let mut output_masks = Vec::with_capacity(50);
3109                        let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> =
3110                            vec![detect.view().into(), protos.view().into()];
3111                        decoder
3112                            .decode_quantized(&inputs, &mut output_boxes, &mut output_masks)
3113                            .unwrap();
3114
3115                        assert_eq!(output_boxes.len(), 1);
3116                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
3117                    }
3118                }
3119            }
3120        };
3121        ($name:ident, float, $layout:ident, $output:ident) => {
3122            #[test]
3123            fn $name() {
3124                let is_split = matches!(stringify!($layout), "split");
3125                let is_proto = matches!(stringify!($output), "proto");
3126
3127                let score_threshold = 0.45;
3128                let iou_threshold = 0.45;
3129
3130                let mut boxes = Array2::zeros((10, 4));
3131                let mut scores = Array2::zeros((10, 1));
3132                let mut classes = Array2::zeros((10, 1));
3133                let mask: Array2<f64> = Array2::zeros((10, 32));
3134                let protos = Array3::<f64>::zeros((160, 160, 32));
3135                let protos = protos.insert_axis(Axis(0));
3136
3137                boxes
3138                    .slice_mut(s![0, ..])
3139                    .assign(&array![0.1234, 0.1234, 0.2345, 0.2345]);
3140                scores.slice_mut(s![0, ..]).assign(&array![0.9876]);
3141                classes.slice_mut(s![0, ..]).assign(&array![2.0]);
3142
3143                let decoder = if is_split {
3144                    DecoderBuilder::default()
3145                        .with_config_yaml_str(E2E_SPLIT_SEGDET_CONFIG.to_string())
3146                        .with_score_threshold(score_threshold)
3147                        .with_iou_threshold(iou_threshold)
3148                        .build()
3149                        .unwrap()
3150                } else {
3151                    DecoderBuilder::default()
3152                        .with_config_yaml_str(E2E_COMBINED_SEGDET_CONFIG.to_string())
3153                        .with_score_threshold(score_threshold)
3154                        .with_iou_threshold(iou_threshold)
3155                        .build()
3156                        .unwrap()
3157                };
3158
3159                let expected = e2e_expected_boxes_float();
3160                let mut output_boxes = Vec::with_capacity(50);
3161
3162                if is_split {
3163                    let boxes = boxes.insert_axis(Axis(0));
3164                    let scores = scores.insert_axis(Axis(0));
3165                    let classes = classes.insert_axis(Axis(0));
3166                    let mask = mask.insert_axis(Axis(0));
3167
3168                    if is_proto {
3169                        let inputs = vec![
3170                            boxes.view().into_dyn(),
3171                            scores.view().into_dyn(),
3172                            classes.view().into_dyn(),
3173                            mask.view().into_dyn(),
3174                            protos.view().into_dyn(),
3175                        ];
3176                        decoder
3177                            .decode_float_proto(&inputs, &mut output_boxes)
3178                            .unwrap();
3179
3180                        assert_eq!(output_boxes.len(), 1);
3181                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
3182                    } else {
3183                        let mut output_masks = Vec::with_capacity(50);
3184                        let inputs = vec![
3185                            boxes.view().into_dyn(),
3186                            scores.view().into_dyn(),
3187                            classes.view().into_dyn(),
3188                            mask.view().into_dyn(),
3189                            protos.view().into_dyn(),
3190                        ];
3191                        decoder
3192                            .decode_float(&inputs, &mut output_boxes, &mut output_masks)
3193                            .unwrap();
3194
3195                        assert_eq!(output_boxes.len(), 1);
3196                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
3197                    }
3198                } else {
3199                    // Combined layout
3200                    let detect = ndarray::concatenate![
3201                        Axis(1),
3202                        boxes.view(),
3203                        scores.view(),
3204                        classes.view(),
3205                        mask.view()
3206                    ];
3207                    let detect = detect.insert_axis(Axis(0));
3208                    assert_eq!(detect.shape(), &[1, 10, 38]);
3209
3210                    if is_proto {
3211                        let inputs = vec![detect.view().into_dyn(), protos.view().into_dyn()];
3212                        decoder
3213                            .decode_float_proto(&inputs, &mut output_boxes)
3214                            .unwrap();
3215
3216                        assert_eq!(output_boxes.len(), 1);
3217                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
3218                    } else {
3219                        let mut output_masks = Vec::with_capacity(50);
3220                        let inputs = vec![detect.view().into_dyn(), protos.view().into_dyn()];
3221                        decoder
3222                            .decode_float(&inputs, &mut output_boxes, &mut output_masks)
3223                            .unwrap();
3224
3225                        assert_eq!(output_boxes.len(), 1);
3226                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
3227                    }
3228                }
3229            }
3230        };
3231    }
3232
3233    e2e_segdet_test!(test_decoder_end_to_end_segdet, quantized, combined, masks);
3234    e2e_segdet_test!(test_decoder_end_to_end_segdet_float, float, combined, masks);
3235    e2e_segdet_test!(
3236        test_decoder_end_to_end_segdet_proto,
3237        quantized,
3238        combined,
3239        proto
3240    );
3241    e2e_segdet_test!(
3242        test_decoder_end_to_end_segdet_proto_float,
3243        float,
3244        combined,
3245        proto
3246    );
3247    e2e_segdet_test!(
3248        test_decoder_end_to_end_segdet_split,
3249        quantized,
3250        split,
3251        masks
3252    );
3253    e2e_segdet_test!(
3254        test_decoder_end_to_end_segdet_split_float,
3255        float,
3256        split,
3257        masks
3258    );
3259    e2e_segdet_test!(
3260        test_decoder_end_to_end_segdet_split_proto,
3261        quantized,
3262        split,
3263        proto
3264    );
3265    e2e_segdet_test!(
3266        test_decoder_end_to_end_segdet_split_proto_float,
3267        float,
3268        split,
3269        proto
3270    );
3271
3272    macro_rules! e2e_det_test {
3273        ($name:ident, quantized, $layout:ident) => {
3274            #[test]
3275            fn $name() {
3276                let is_split = matches!(stringify!($layout), "split");
3277
3278                let score_threshold = 0.45;
3279                let iou_threshold = 0.45;
3280
3281                let mut boxes = Array3::zeros((1, 10, 4));
3282                let mut scores = Array3::zeros((1, 10, 1));
3283                let mut classes = Array3::zeros((1, 10, 1));
3284
3285                boxes
3286                    .slice_mut(s![0, 0, ..])
3287                    .assign(&array![0.1234, 0.1234, 0.2345, 0.2345]);
3288                scores.slice_mut(s![0, 0, ..]).assign(&array![0.9876]);
3289                classes.slice_mut(s![0, 0, ..]).assign(&array![2.0]);
3290
3291                let detect_quant = (2.0 / 255.0, 0_i32);
3292
3293                let decoder = if is_split {
3294                    DecoderBuilder::default()
3295                        .with_config_yaml_str(E2E_SPLIT_DET_CONFIG.to_string())
3296                        .with_score_threshold(score_threshold)
3297                        .with_iou_threshold(iou_threshold)
3298                        .build()
3299                        .unwrap()
3300                } else {
3301                    DecoderBuilder::default()
3302                        .with_config_yaml_str(E2E_COMBINED_DET_CONFIG.to_string())
3303                        .with_score_threshold(score_threshold)
3304                        .with_iou_threshold(iou_threshold)
3305                        .build()
3306                        .unwrap()
3307                };
3308
3309                let expected = e2e_expected_boxes_quant();
3310                let mut output_boxes = Vec::with_capacity(50);
3311
3312                if is_split {
3313                    let boxes: Array<u8, _> = quantize_ndarray(boxes.view(), detect_quant.into());
3314                    let scores: Array<u8, _> = quantize_ndarray(scores.view(), detect_quant.into());
3315                    let classes: Array<u8, _> =
3316                        quantize_ndarray(classes.view(), detect_quant.into());
3317                    let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> = vec![
3318                        boxes.view().into(),
3319                        scores.view().into(),
3320                        classes.view().into(),
3321                    ];
3322                    decoder
3323                        .decode_quantized(&inputs, &mut output_boxes, &mut Vec::new())
3324                        .unwrap();
3325                } else {
3326                    let detect =
3327                        ndarray::concatenate![Axis(2), boxes.view(), scores.view(), classes.view()];
3328                    assert_eq!(detect.shape(), &[1, 10, 6]);
3329                    let detect: Array3<u8> = quantize_ndarray(detect.view(), detect_quant.into());
3330                    let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> =
3331                        vec![detect.view().into()];
3332                    decoder
3333                        .decode_quantized(&inputs, &mut output_boxes, &mut Vec::new())
3334                        .unwrap();
3335                }
3336
3337                assert_eq!(output_boxes.len(), 1);
3338                assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
3339            }
3340        };
3341        ($name:ident, float, $layout:ident) => {
3342            #[test]
3343            fn $name() {
3344                let is_split = matches!(stringify!($layout), "split");
3345
3346                let score_threshold = 0.45;
3347                let iou_threshold = 0.45;
3348
3349                let mut boxes = Array3::zeros((1, 10, 4));
3350                let mut scores = Array3::zeros((1, 10, 1));
3351                let mut classes = Array3::zeros((1, 10, 1));
3352
3353                boxes
3354                    .slice_mut(s![0, 0, ..])
3355                    .assign(&array![0.1234, 0.1234, 0.2345, 0.2345]);
3356                scores.slice_mut(s![0, 0, ..]).assign(&array![0.9876]);
3357                classes.slice_mut(s![0, 0, ..]).assign(&array![2.0]);
3358
3359                let decoder = if is_split {
3360                    DecoderBuilder::default()
3361                        .with_config_yaml_str(E2E_SPLIT_DET_CONFIG.to_string())
3362                        .with_score_threshold(score_threshold)
3363                        .with_iou_threshold(iou_threshold)
3364                        .build()
3365                        .unwrap()
3366                } else {
3367                    DecoderBuilder::default()
3368                        .with_config_yaml_str(E2E_COMBINED_DET_CONFIG.to_string())
3369                        .with_score_threshold(score_threshold)
3370                        .with_iou_threshold(iou_threshold)
3371                        .build()
3372                        .unwrap()
3373                };
3374
3375                let expected = e2e_expected_boxes_float();
3376                let mut output_boxes = Vec::with_capacity(50);
3377
3378                if is_split {
3379                    let inputs = vec![
3380                        boxes.view().into_dyn(),
3381                        scores.view().into_dyn(),
3382                        classes.view().into_dyn(),
3383                    ];
3384                    decoder
3385                        .decode_float(&inputs, &mut output_boxes, &mut Vec::new())
3386                        .unwrap();
3387                } else {
3388                    let detect =
3389                        ndarray::concatenate![Axis(2), boxes.view(), scores.view(), classes.view()];
3390                    assert_eq!(detect.shape(), &[1, 10, 6]);
3391                    let inputs = vec![detect.view().into_dyn()];
3392                    decoder
3393                        .decode_float(&inputs, &mut output_boxes, &mut Vec::new())
3394                        .unwrap();
3395                }
3396
3397                assert_eq!(output_boxes.len(), 1);
3398                assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
3399            }
3400        };
3401    }
3402
3403    e2e_det_test!(test_decoder_end_to_end_combined_det, quantized, combined);
3404    e2e_det_test!(test_decoder_end_to_end_combined_det_float, float, combined);
3405    e2e_det_test!(test_decoder_end_to_end_split_det, quantized, split);
3406    e2e_det_test!(test_decoder_end_to_end_split_det_float, float, split);
3407
3408    #[test]
3409    fn test_decode_tensor() {
3410        let score_threshold = 0.45;
3411        let iou_threshold = 0.45;
3412
3413        let raw_boxes = edgefirst_bench::testdata::read("yolov8_boxes_116x8400.bin");
3414        let raw_boxes =
3415            unsafe { std::slice::from_raw_parts(raw_boxes.as_ptr() as *const i8, raw_boxes.len()) };
3416        let boxes_i8: Tensor<i8> = Tensor::new(&[1, 116, 8400], None, None).unwrap();
3417        boxes_i8
3418            .map()
3419            .unwrap()
3420            .as_mut_slice()
3421            .copy_from_slice(raw_boxes);
3422        let boxes_i8 = boxes_i8.into();
3423
3424        let raw_protos = edgefirst_bench::testdata::read("yolov8_protos_160x160x32.bin");
3425        let raw_protos = unsafe {
3426            std::slice::from_raw_parts(raw_protos.as_ptr() as *const i8, raw_protos.len())
3427        };
3428        let protos_i8: Tensor<i8> = Tensor::new(&[1, 160, 160, 32], None, None).unwrap();
3429        protos_i8
3430            .map()
3431            .unwrap()
3432            .as_mut_slice()
3433            .copy_from_slice(raw_protos);
3434        let protos_i8 = protos_i8.into();
3435
3436        let decoder = build_yolov8_seg_decoder(score_threshold, iou_threshold);
3437        let expected = real_data_expected_boxes();
3438        let mut output_boxes = Vec::with_capacity(50);
3439
3440        decoder
3441            .decode(&[&boxes_i8, &protos_i8], &mut output_boxes, &mut Vec::new())
3442            .unwrap();
3443
3444        assert_eq!(output_boxes.len(), 2);
3445        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
3446        assert!(output_boxes[1].equal_within_delta(&expected[1], 1.0 / 160.0));
3447    }
3448
3449    #[test]
3450    fn test_decode_tensor_f32() {
3451        let score_threshold = 0.45;
3452        let iou_threshold = 0.45;
3453
3454        let quant_boxes = (0.021287762_f32, 31_i32);
3455        let quant_protos = (0.02491162_f32, -117_i32);
3456        let raw_boxes = edgefirst_bench::testdata::read("yolov8_boxes_116x8400.bin");
3457        let raw_boxes =
3458            unsafe { std::slice::from_raw_parts(raw_boxes.as_ptr() as *const i8, raw_boxes.len()) };
3459        let mut raw_boxes_f32 = vec![0f32; raw_boxes.len()];
3460        dequantize_cpu(raw_boxes, quant_boxes.into(), &mut raw_boxes_f32);
3461        let boxes_f32: Tensor<f32> = Tensor::new(&[1, 116, 8400], None, None).unwrap();
3462        boxes_f32
3463            .map()
3464            .unwrap()
3465            .as_mut_slice()
3466            .copy_from_slice(&raw_boxes_f32);
3467        let boxes_f32 = boxes_f32.into();
3468
3469        let raw_protos = edgefirst_bench::testdata::read("yolov8_protos_160x160x32.bin");
3470        let raw_protos = unsafe {
3471            std::slice::from_raw_parts(raw_protos.as_ptr() as *const i8, raw_protos.len())
3472        };
3473        let mut raw_protos_f32 = vec![0f32; raw_protos.len()];
3474        dequantize_cpu(raw_protos, quant_protos.into(), &mut raw_protos_f32);
3475        let protos_f32: Tensor<f32> = Tensor::new(&[1, 160, 160, 32], None, None).unwrap();
3476        protos_f32
3477            .map()
3478            .unwrap()
3479            .as_mut_slice()
3480            .copy_from_slice(&raw_protos_f32);
3481        let protos_f32 = protos_f32.into();
3482
3483        let decoder = build_yolov8_seg_decoder(score_threshold, iou_threshold);
3484
3485        let expected = real_data_expected_boxes();
3486        let mut output_boxes = Vec::with_capacity(50);
3487
3488        decoder
3489            .decode(
3490                &[&boxes_f32, &protos_f32],
3491                &mut output_boxes,
3492                &mut Vec::new(),
3493            )
3494            .unwrap();
3495
3496        assert_eq!(output_boxes.len(), 2);
3497        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
3498        assert!(output_boxes[1].equal_within_delta(&expected[1], 1.0 / 160.0));
3499    }
3500
3501    #[test]
3502    fn test_decode_tensor_f64() {
3503        let score_threshold = 0.45;
3504        let iou_threshold = 0.45;
3505
3506        let quant_boxes = (0.021287762_f32, 31_i32);
3507        let quant_protos = (0.02491162_f32, -117_i32);
3508        let raw_boxes = edgefirst_bench::testdata::read("yolov8_boxes_116x8400.bin");
3509        let raw_boxes =
3510            unsafe { std::slice::from_raw_parts(raw_boxes.as_ptr() as *const i8, raw_boxes.len()) };
3511        let mut raw_boxes_f64 = vec![0f64; raw_boxes.len()];
3512        dequantize_cpu(raw_boxes, quant_boxes.into(), &mut raw_boxes_f64);
3513        let boxes_f64: Tensor<f64> = Tensor::new(&[1, 116, 8400], None, None).unwrap();
3514        boxes_f64
3515            .map()
3516            .unwrap()
3517            .as_mut_slice()
3518            .copy_from_slice(&raw_boxes_f64);
3519        let boxes_f64 = boxes_f64.into();
3520
3521        let raw_protos = edgefirst_bench::testdata::read("yolov8_protos_160x160x32.bin");
3522        let raw_protos = unsafe {
3523            std::slice::from_raw_parts(raw_protos.as_ptr() as *const i8, raw_protos.len())
3524        };
3525        let mut raw_protos_f64 = vec![0f64; raw_protos.len()];
3526        dequantize_cpu(raw_protos, quant_protos.into(), &mut raw_protos_f64);
3527        let protos_f64: Tensor<f64> = Tensor::new(&[1, 160, 160, 32], None, None).unwrap();
3528        protos_f64
3529            .map()
3530            .unwrap()
3531            .as_mut_slice()
3532            .copy_from_slice(&raw_protos_f64);
3533        let protos_f64 = protos_f64.into();
3534
3535        let decoder = build_yolov8_seg_decoder(score_threshold, iou_threshold);
3536
3537        let expected = real_data_expected_boxes();
3538        let mut output_boxes = Vec::with_capacity(50);
3539
3540        decoder
3541            .decode(
3542                &[&boxes_f64, &protos_f64],
3543                &mut output_boxes,
3544                &mut Vec::new(),
3545            )
3546            .unwrap();
3547
3548        assert_eq!(output_boxes.len(), 2);
3549        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
3550        assert!(output_boxes[1].equal_within_delta(&expected[1], 1.0 / 160.0));
3551    }
3552
3553    #[test]
3554    fn test_decode_tensor_proto() {
3555        let score_threshold = 0.45;
3556        let iou_threshold = 0.45;
3557
3558        let raw_boxes = edgefirst_bench::testdata::read("yolov8_boxes_116x8400.bin");
3559        let raw_boxes =
3560            unsafe { std::slice::from_raw_parts(raw_boxes.as_ptr() as *const i8, raw_boxes.len()) };
3561        let boxes_i8: Tensor<i8> = Tensor::new(&[1, 116, 8400], None, None).unwrap();
3562        boxes_i8
3563            .map()
3564            .unwrap()
3565            .as_mut_slice()
3566            .copy_from_slice(raw_boxes);
3567        let boxes_i8 = boxes_i8.into();
3568
3569        let raw_protos = edgefirst_bench::testdata::read("yolov8_protos_160x160x32.bin");
3570        let raw_protos = unsafe {
3571            std::slice::from_raw_parts(raw_protos.as_ptr() as *const i8, raw_protos.len())
3572        };
3573        let protos_i8: Tensor<i8> = Tensor::new(&[1, 160, 160, 32], None, None).unwrap();
3574        protos_i8
3575            .map()
3576            .unwrap()
3577            .as_mut_slice()
3578            .copy_from_slice(raw_protos);
3579        let protos_i8 = protos_i8.into();
3580
3581        let decoder = build_yolov8_seg_decoder(score_threshold, iou_threshold);
3582
3583        let expected = real_data_expected_boxes();
3584        let mut output_boxes = Vec::with_capacity(50);
3585
3586        let proto_data = decoder
3587            .decode_proto(&[&boxes_i8, &protos_i8], &mut output_boxes)
3588            .unwrap();
3589
3590        assert_eq!(output_boxes.len(), 2);
3591        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
3592        assert!(output_boxes[1].equal_within_delta(&expected[1], 1.0 / 160.0));
3593
3594        let proto_data = proto_data.expect("segmentation model should return ProtoData");
3595        let coeffs_shape = proto_data.mask_coefficients.shape();
3596        assert_eq!(
3597            coeffs_shape[0],
3598            output_boxes.len(),
3599            "mask_coefficients count must match detection count"
3600        );
3601        assert_eq!(
3602            coeffs_shape[1], 32,
3603            "each detection should have 32 mask coefficients"
3604        );
3605    }
3606
3607    // =========================================================================
3608    // Physical-order contract regression tests
3609    //
3610    // These cover the TFLite VX-delegate vertical-stripe mask bug and the
3611    // Ara-2-via-overlay anchor-first split-tensor bug. The core claim:
3612    // declaring shape+dshape in physical memory order produces correct
3613    // strides at wrap time, and `swap_axes_if_needed` permutes the stride
3614    // tuple into canonical logical order without touching bytes.
3615    // =========================================================================
3616
3617    /// TFLite YOLOv8-seg protos arrive as physically-NHWC bytes with
3618    /// NNStreamer dim string `"32:160:160:1"` (innermost-first). The
3619    /// overlay's 2026-04-22 workaround declares `[1, 160, 160, 32]` with
3620    /// dshape `[batch, height, width, num_protos]` — shape matches
3621    /// physical order, so no permutation is needed and the view is
3622    /// already in canonical HWC after the batch-dim slice. This test
3623    /// confirms that path matches the reference (direct-HWC) decode.
3624    #[test]
3625    fn test_physical_order_tflite_nhwc_protos() {
3626        let score_threshold = 0.45;
3627        let iou_threshold = 0.45;
3628
3629        // Load HWC protos and dequantize to f32.
3630        let protos_hwc = load_yolov8_protos().slice_move(s![0, .., .., ..]);
3631        let quant_protos = Quantization::new(0.02491161972284317, -117);
3632        let protos_f32_hwc = dequantize_ndarray::<_, _, f32>(protos_hwc.view(), quant_protos);
3633
3634        let boxes_2d = load_yolov8_boxes().slice_move(s![0, .., ..]);
3635        let quant_boxes = Quantization::new(0.021287761628627777, 31);
3636        let seg = dequantize_ndarray::<_, _, f32>(boxes_2d.view(), quant_boxes);
3637
3638        // Reference: direct call with canonical HWC protos.
3639        let mut ref_boxes: Vec<_> = Vec::with_capacity(10);
3640        let mut ref_masks: Vec<_> = Vec::with_capacity(10);
3641        decode_yolo_segdet_float(
3642            seg.view(),
3643            protos_f32_hwc.view(),
3644            score_threshold,
3645            iou_threshold,
3646            Some(configs::Nms::ClassAgnostic),
3647            &mut ref_boxes,
3648            &mut ref_masks,
3649        )
3650        .unwrap();
3651
3652        // Build the same protos as a 4D NHWC tensor and feed it through
3653        // the config-driven path with dshape in physical order.
3654        let protos_nhwc = protos_f32_hwc.clone().insert_axis(Axis(0)); // [1, 160, 160, 32]
3655        let seg_3d = seg.insert_axis(Axis(0)); // [1, 116, 8400]
3656
3657        let decoder = DecoderBuilder::default()
3658            .with_config_yolo_segdet(
3659                configs::Detection {
3660                    decoder: configs::DecoderType::Ultralytics,
3661                    quantization: None,
3662                    shape: vec![1, 116, 8400],
3663                    dshape: vec![
3664                        (DimName::Batch, 1),
3665                        (DimName::NumFeatures, 116),
3666                        (DimName::NumBoxes, 8400),
3667                    ],
3668                    normalized: Some(true),
3669                    anchors: None,
3670                },
3671                configs::Protos {
3672                    decoder: configs::DecoderType::Ultralytics,
3673                    quantization: None,
3674                    shape: vec![1, 160, 160, 32],
3675                    // Physical NHWC — matches the TFLite buffer layout.
3676                    dshape: vec![
3677                        (DimName::Batch, 1),
3678                        (DimName::Height, 160),
3679                        (DimName::Width, 160),
3680                        (DimName::NumProtos, 32),
3681                    ],
3682                },
3683                None,
3684            )
3685            .with_score_threshold(score_threshold)
3686            .with_iou_threshold(iou_threshold)
3687            .build()
3688            .expect("config with NHWC protos dshape must build");
3689
3690        let mut cfg_boxes = Vec::with_capacity(10);
3691        let mut cfg_masks = Vec::with_capacity(10);
3692        decoder
3693            .decode_float(
3694                &[seg_3d.view().into_dyn(), protos_nhwc.view().into_dyn()],
3695                &mut cfg_boxes,
3696                &mut cfg_masks,
3697            )
3698            .unwrap();
3699
3700        assert_eq!(cfg_boxes.len(), ref_boxes.len(), "box count mismatch");
3701        for (c, r) in cfg_boxes.iter().zip(&ref_boxes) {
3702            assert!(
3703                c.equal_within_delta(r, 0.01),
3704                "NHWC-declared box does not match reference: {c:?} vs {r:?}"
3705            );
3706        }
3707        for (cm, rm) in cfg_masks.iter().zip(&ref_masks) {
3708            let cm_arr = segmentation_to_mask(seg_view(cm).view()).unwrap();
3709            let rm_arr = segmentation_to_mask(seg_view(rm).view()).unwrap();
3710            assert_eq!(
3711                cm_arr, rm_arr,
3712                "NHWC-declared mask must match reference pixel-for-pixel"
3713            );
3714        }
3715    }
3716
3717    /// Ara-2 split-tensor boxes arrive physically anchor-first: NNStreamer
3718    /// dim string `"4:1:8400:1"` means innermost axis is 4 (coords), next
3719    /// is 1 (padding), next is 8400 (anchors), outermost is 1 (batch).
3720    /// The correct physical-order declaration is `[1, 8400, 1, 4]` with
3721    /// dshape `[batch, num_boxes, padding, box_coords]`. This test
3722    /// confirms that declaration produces the same decoded boxes as the
3723    /// canonical features-first TFLite declaration over the same logical
3724    /// data.
3725    #[test]
3726    fn test_physical_order_ara2_anchor_first_split_boxes() {
3727        use configs::{Boxes, Scores};
3728
3729        // Build synthetic boxes data in features-first canonical layout:
3730        // shape [1, 4, 8400] with one detection at anchor 42.
3731        const N: usize = 8400;
3732        let mut boxes_canonical = Array3::<f32>::zeros((1, 4, N));
3733        let target_anchor = 42usize;
3734        boxes_canonical[[0, 0, target_anchor]] = 0.4; // xc
3735        boxes_canonical[[0, 1, target_anchor]] = 0.5; // yc
3736        boxes_canonical[[0, 2, target_anchor]] = 0.2; // w
3737        boxes_canonical[[0, 3, target_anchor]] = 0.2; // h
3738
3739        // Scores: one class at 0.9 for the same anchor.
3740        let mut scores_canonical = Array3::<f32>::zeros((1, 80, N));
3741        scores_canonical[[0, 0, target_anchor]] = 0.9;
3742
3743        // Reference: canonical (features-first) shape declaration.
3744        let ref_decoder = DecoderBuilder::default()
3745            .with_config_yolo_split_det(
3746                Boxes {
3747                    decoder: configs::DecoderType::Ultralytics,
3748                    quantization: None,
3749                    shape: vec![1, 4, N],
3750                    dshape: vec![
3751                        (DimName::Batch, 1),
3752                        (DimName::BoxCoords, 4),
3753                        (DimName::NumBoxes, N),
3754                    ],
3755                    normalized: Some(true),
3756                },
3757                Scores {
3758                    decoder: configs::DecoderType::Ultralytics,
3759                    quantization: None,
3760                    shape: vec![1, 80, N],
3761                    dshape: vec![
3762                        (DimName::Batch, 1),
3763                        (DimName::NumClasses, 80),
3764                        (DimName::NumBoxes, N),
3765                    ],
3766                },
3767            )
3768            .with_score_threshold(0.5)
3769            .with_iou_threshold(0.5)
3770            .with_nms(Some(configs::Nms::ClassAgnostic))
3771            .build()
3772            .expect("reference canonical split decoder must build");
3773
3774        let mut ref_boxes = Vec::with_capacity(4);
3775        let mut ref_masks = Vec::with_capacity(0);
3776        ref_decoder
3777            .decode_float(
3778                &[
3779                    boxes_canonical.view().into_dyn(),
3780                    scores_canonical.view().into_dyn(),
3781                ],
3782                &mut ref_boxes,
3783                &mut ref_masks,
3784            )
3785            .unwrap();
3786        assert_eq!(ref_boxes.len(), 1, "reference should produce one box");
3787
3788        // Ara-2 physical layout: transpose axes 1↔2 so the innermost
3789        // axis is BoxCoords / NumClasses. Materialise to get a
3790        // C-contiguous Array3 with strides matching the physical order
3791        // the Ara-2 backend would write into the tensor.
3792        let boxes_ara2 = boxes_canonical.view().permuted_axes([0, 2, 1]).to_owned(); // [1, 8400, 4]
3793        let scores_ara2 = scores_canonical.view().permuted_axes([0, 2, 1]).to_owned(); // [1, 8400, 80]
3794
3795        let ara2_decoder = DecoderBuilder::default()
3796            .with_config_yolo_split_det(
3797                Boxes {
3798                    decoder: configs::DecoderType::Ultralytics,
3799                    quantization: None,
3800                    shape: vec![1, N, 4],
3801                    dshape: vec![
3802                        (DimName::Batch, 1),
3803                        (DimName::NumBoxes, N),
3804                        (DimName::BoxCoords, 4),
3805                    ],
3806                    normalized: Some(true),
3807                },
3808                Scores {
3809                    decoder: configs::DecoderType::Ultralytics,
3810                    quantization: None,
3811                    shape: vec![1, N, 80],
3812                    dshape: vec![
3813                        (DimName::Batch, 1),
3814                        (DimName::NumBoxes, N),
3815                        (DimName::NumClasses, 80),
3816                    ],
3817                },
3818            )
3819            .with_score_threshold(0.5)
3820            .with_iou_threshold(0.5)
3821            .with_nms(Some(configs::Nms::ClassAgnostic))
3822            .build()
3823            .expect("Ara-2 anchor-first decoder must build");
3824
3825        let mut ara2_boxes = Vec::with_capacity(4);
3826        let mut ara2_masks = Vec::with_capacity(0);
3827        ara2_decoder
3828            .decode_float(
3829                &[boxes_ara2.view().into_dyn(), scores_ara2.view().into_dyn()],
3830                &mut ara2_boxes,
3831                &mut ara2_masks,
3832            )
3833            .unwrap();
3834
3835        assert_eq!(
3836            ara2_boxes.len(),
3837            ref_boxes.len(),
3838            "Ara-2 anchor-first declaration must produce the same number \
3839             of boxes as the canonical features-first reference"
3840        );
3841        for (a, r) in ara2_boxes.iter().zip(&ref_boxes) {
3842            assert!(
3843                a.equal_within_delta(r, 1e-4),
3844                "Ara-2 box differs from reference: {a:?} vs {r:?}"
3845            );
3846        }
3847    }
3848
3849    /// Dshape whose per-axis sizes disagree with shape (the exact
3850    /// failure mode that caused the TFLite stripe bug — shape declared
3851    /// in one order, dshape in another) is rejected with a clear error.
3852    #[test]
3853    fn test_physical_order_rejects_shape_dshape_mismatch() {
3854        let result = DecoderBuilder::default()
3855            .with_config_yolo_segdet(
3856                configs::Detection {
3857                    decoder: configs::DecoderType::Ultralytics,
3858                    quantization: None,
3859                    shape: vec![1, 116, 8400],
3860                    dshape: vec![
3861                        (DimName::Batch, 1),
3862                        (DimName::NumFeatures, 116),
3863                        (DimName::NumBoxes, 8400),
3864                    ],
3865                    normalized: Some(true),
3866                    anchors: None,
3867                },
3868                configs::Protos {
3869                    decoder: configs::DecoderType::Ultralytics,
3870                    quantization: None,
3871                    // Shape in NCHW order...
3872                    shape: vec![1, 32, 160, 160],
3873                    // ...but dshape in NHWC order — the sizes don't line
3874                    // up positionally (dshape[1]=160 vs shape[1]=32).
3875                    dshape: vec![
3876                        (DimName::Batch, 1),
3877                        (DimName::Height, 160),
3878                        (DimName::Width, 160),
3879                        (DimName::NumProtos, 32),
3880                    ],
3881                },
3882                None,
3883            )
3884            .build();
3885
3886        match result {
3887            Err(DecoderError::InvalidConfig(msg)) => {
3888                assert!(
3889                    msg.contains("does not match shape"),
3890                    "expected shape/dshape size mismatch error, got: {msg}"
3891                );
3892            }
3893            other => panic!("expected InvalidConfig, got {other:?}"),
3894        }
3895    }
3896
3897    /// Duplicate dim name in dshape is rejected (two axes mapping to the
3898    /// same canonical slot would break the permutation).
3899    #[test]
3900    fn test_physical_order_rejects_duplicate_dshape_axis() {
3901        let result = DecoderBuilder::default()
3902            .with_config_yolo_split_det(
3903                configs::Boxes {
3904                    decoder: configs::DecoderType::Ultralytics,
3905                    quantization: None,
3906                    shape: vec![1, 4, 8400],
3907                    dshape: vec![
3908                        (DimName::Batch, 1),
3909                        (DimName::BoxCoords, 4),
3910                        (DimName::BoxCoords, 4), // duplicate (size matches shape[2]=8400? no)
3911                    ],
3912                    normalized: Some(true),
3913                },
3914                configs::Scores {
3915                    decoder: configs::DecoderType::Ultralytics,
3916                    quantization: None,
3917                    shape: vec![1, 80, 8400],
3918                    dshape: vec![
3919                        (DimName::Batch, 1),
3920                        (DimName::NumClasses, 80),
3921                        (DimName::NumBoxes, 8400),
3922                    ],
3923                },
3924            )
3925            .build();
3926
3927        // Shape[2] = 8400 ≠ dshape[2] size 4, so the positional-mismatch
3928        // check fires first — which is fine: any mis-shaped dshape
3929        // should be rejected. Check for either error as long as it's a
3930        // clear InvalidConfig.
3931        match result {
3932            Err(DecoderError::InvalidConfig(msg)) => {
3933                assert!(
3934                    msg.contains("appears at both index") || msg.contains("does not match shape"),
3935                    "expected positional or duplicate-axis error, got: {msg}"
3936                );
3937            }
3938            other => panic!("expected InvalidConfig, got {other:?}"),
3939        }
3940
3941        // Separate case: size-consistent duplicate to exercise the
3942        // duplicate-axis code path explicitly. Use `Batch` (no size
3943        // constraint, can repeat with size 1 against a shape that
3944        // carries two singleton dims).
3945        let result = DecoderBuilder::default()
3946            .with_config_yolo_split_det(
3947                configs::Boxes {
3948                    decoder: configs::DecoderType::Ultralytics,
3949                    quantization: None,
3950                    shape: vec![1, 1, 4, 8400],
3951                    dshape: vec![
3952                        (DimName::Batch, 1),
3953                        (DimName::Batch, 1), // duplicate, sizes match
3954                        (DimName::BoxCoords, 4),
3955                        (DimName::NumBoxes, 8400),
3956                    ],
3957                    normalized: Some(true),
3958                },
3959                configs::Scores {
3960                    decoder: configs::DecoderType::Ultralytics,
3961                    quantization: None,
3962                    shape: vec![1, 80, 8400],
3963                    dshape: vec![
3964                        (DimName::Batch, 1),
3965                        (DimName::NumClasses, 80),
3966                        (DimName::NumBoxes, 8400),
3967                    ],
3968                },
3969            )
3970            .build();
3971        match result {
3972            Err(DecoderError::InvalidConfig(msg)) => {
3973                assert!(
3974                    msg.contains("appears at both index"),
3975                    "expected duplicate-axis error, got: {msg}"
3976                );
3977            }
3978            other => panic!("expected InvalidConfig, got {other:?}"),
3979        }
3980    }
3981
3982    /// Canonical (dshape-omitted) high-level decode produces the same
3983    /// numeric result as the dshape-populated form. This closes a
3984    /// coverage gap flagged during review: dshape omission had been
3985    /// exercised only at the builder level, not through the full
3986    /// `decode_float` pipeline.
3987    #[test]
3988    fn test_physical_order_dshape_omitted_decodes_numerically() {
3989        let score_threshold = 0.45;
3990        let iou_threshold = 0.45;
3991
3992        let protos_hwc = load_yolov8_protos().slice_move(s![0, .., .., ..]);
3993        let quant_protos = Quantization::new(0.02491161972284317, -117);
3994        let protos_f32_hwc = dequantize_ndarray::<_, _, f32>(protos_hwc.view(), quant_protos);
3995
3996        let boxes_2d = load_yolov8_boxes().slice_move(s![0, .., ..]);
3997        let quant_boxes = Quantization::new(0.021287761628627777, 31);
3998        let seg = dequantize_ndarray::<_, _, f32>(boxes_2d.view(), quant_boxes);
3999
4000        let protos_nhwc = protos_f32_hwc.clone().insert_axis(Axis(0));
4001        let seg_3d = seg.insert_axis(Axis(0));
4002
4003        let build_decoder = |det_dshape: Vec<(DimName, usize)>,
4004                             proto_dshape: Vec<(DimName, usize)>| {
4005            DecoderBuilder::default()
4006                .with_config_yolo_segdet(
4007                    configs::Detection {
4008                        decoder: configs::DecoderType::Ultralytics,
4009                        quantization: None,
4010                        shape: vec![1, 116, 8400],
4011                        dshape: det_dshape,
4012                        normalized: Some(true),
4013                        anchors: None,
4014                    },
4015                    configs::Protos {
4016                        decoder: configs::DecoderType::Ultralytics,
4017                        quantization: None,
4018                        shape: vec![1, 160, 160, 32],
4019                        dshape: proto_dshape,
4020                    },
4021                    None,
4022                )
4023                .with_score_threshold(score_threshold)
4024                .with_iou_threshold(iou_threshold)
4025                .build()
4026                .unwrap()
4027        };
4028
4029        // Baseline: dshape populated in physical order.
4030        let dshaped = build_decoder(
4031            vec![
4032                (DimName::Batch, 1),
4033                (DimName::NumFeatures, 116),
4034                (DimName::NumBoxes, 8400),
4035            ],
4036            vec![
4037                (DimName::Batch, 1),
4038                (DimName::Height, 160),
4039                (DimName::Width, 160),
4040                (DimName::NumProtos, 32),
4041            ],
4042        );
4043        let mut dshaped_boxes = Vec::new();
4044        let mut dshaped_masks = Vec::new();
4045        dshaped
4046            .decode_float(
4047                &[seg_3d.view().into_dyn(), protos_nhwc.view().into_dyn()],
4048                &mut dshaped_boxes,
4049                &mut dshaped_masks,
4050            )
4051            .unwrap();
4052
4053        // Variant: dshape omitted — caller asserts shape is already in
4054        // the decoder's canonical order.
4055        let bare = build_decoder(vec![], vec![]);
4056        let mut bare_boxes = Vec::new();
4057        let mut bare_masks = Vec::new();
4058        bare.decode_float(
4059            &[seg_3d.view().into_dyn(), protos_nhwc.view().into_dyn()],
4060            &mut bare_boxes,
4061            &mut bare_masks,
4062        )
4063        .unwrap();
4064
4065        assert_eq!(bare_boxes.len(), dshaped_boxes.len());
4066        for (b, d) in bare_boxes.iter().zip(&dshaped_boxes) {
4067            assert!(
4068                b.equal_within_delta(d, 1e-4),
4069                "dshape-omitted box {b:?} differs from dshape-populated {d:?}"
4070            );
4071        }
4072        for (bm, dm) in bare_masks.iter().zip(&dshaped_masks) {
4073            let bm_arr = segmentation_to_mask(seg_view(bm).view()).unwrap();
4074            let dm_arr = segmentation_to_mask(seg_view(dm).view()).unwrap();
4075            assert_eq!(
4076                bm_arr, dm_arr,
4077                "dshape-omitted mask must match dshape-populated pixel-for-pixel"
4078            );
4079        }
4080    }
4081
4082    /// 4D anchor-first boxes schema with a trailing `padding` axis — the
4083    /// exact Ara-2 DVM shape (`"4:1:8400:1"` innermost-first = physical
4084    /// `[1, 8400, 1, 4]`). Exercises the schema path's
4085    /// `squeeze_padding_dims`: the caller declares a 4D physical-order
4086    /// layout including `padding`, HAL squeezes the size-1 axis during
4087    /// `to_legacy_config_outputs`, and the decoder sees the resulting
4088    /// 3D shape. Callers feed HAL a 3D squeezed view of the same bytes.
4089    /// Complements the 3D `test_physical_order_ara2_anchor_first_split_boxes`
4090    /// which exercises the programmatic-builder path.
4091    #[test]
4092    fn test_physical_order_ara2_4d_anchor_first_with_padding() {
4093        // Build synthetic data in anchor-first 3D layout (the shape HAL
4094        // actually operates on after squeezing). The 4D-with-padding
4095        // declaration in the schema is a producer-side convention; the
4096        // caller is expected to present HAL with a squeezed view.
4097        const N: usize = 8400;
4098        let mut boxes = Array3::<f32>::zeros((1, N, 4));
4099        let target = 42usize;
4100        boxes[[0, target, 0]] = 0.4;
4101        boxes[[0, target, 1]] = 0.5;
4102        boxes[[0, target, 2]] = 0.2;
4103        boxes[[0, target, 3]] = 0.2;
4104        let mut scores = Array3::<f32>::zeros((1, N, 80));
4105        scores[[0, target, 0]] = 0.9;
4106
4107        // Schema JSON declares shape+dshape in physical anchor-first
4108        // order including padding — mirrors `ara2_int8_edgefirst.json`'s
4109        // feature-first declaration but with the axes in physical
4110        // anchor-first order. `to_legacy_config_outputs` squeezes
4111        // padding before the decoder's rank-3 verification.
4112        let json = r#"{
4113          "schema_version": 2,
4114          "decoder_version": "yolov8",
4115          "nms": "class_agnostic",
4116          "outputs": [
4117            {"name": "boxes", "type": "boxes",
4118             "shape": [1, 8400, 1, 4],
4119             "dshape": [{"batch":1},{"num_boxes":8400},{"padding":1},{"box_coords":4}],
4120             "encoding": "direct",
4121             "decoder": "ultralytics",
4122             "normalized": true},
4123            {"name": "scores", "type": "scores",
4124             "shape": [1, 8400, 1, 80],
4125             "dshape": [{"batch":1},{"num_boxes":8400},{"padding":1},{"num_classes":80}],
4126             "decoder": "ultralytics",
4127             "score_format": "per_class"}
4128          ]
4129        }"#;
4130        let decoder = DecoderBuilder::default()
4131            .with_config_json_str(json.to_string())
4132            .with_score_threshold(0.5)
4133            .with_iou_threshold(0.5)
4134            .build()
4135            .expect("4D anchor-first schema should build via squeeze_padding_dims");
4136
4137        let mut out_boxes = Vec::with_capacity(4);
4138        let mut out_masks = Vec::with_capacity(0);
4139        decoder
4140            .decode_float(
4141                &[boxes.view().into_dyn(), scores.view().into_dyn()],
4142                &mut out_boxes,
4143                &mut out_masks,
4144            )
4145            .unwrap();
4146
4147        assert_eq!(
4148            out_boxes.len(),
4149            1,
4150            "4D anchor-first with padding should decode exactly one box from the seeded anchor"
4151        );
4152        let b = &out_boxes[0];
4153        // xywh(0.4, 0.5, 0.2, 0.2) → xyxy(0.3, 0.4, 0.5, 0.6)
4154        assert!((b.bbox.xmin - 0.3).abs() < 1e-3, "xmin wrong: {b:?}");
4155        assert!((b.bbox.ymin - 0.4).abs() < 1e-3, "ymin wrong: {b:?}");
4156        assert!((b.bbox.xmax - 0.5).abs() < 1e-3, "xmax wrong: {b:?}");
4157        assert!((b.bbox.ymax - 0.6).abs() < 1e-3, "ymax wrong: {b:?}");
4158        assert_eq!(b.label, 0);
4159        assert!(b.score > 0.85, "score {}: {b:?}", b.score);
4160    }
4161}
4162
4163#[cfg(feature = "tracker")]
4164#[cfg(test)]
4165#[cfg_attr(coverage_nightly, coverage(off))]
4166mod decoder_tracked_tests {
4167
4168    use crate::DetectBoxTestExt;
4169    use edgefirst_tracker::{ByteTrackBuilder, Tracker};
4170    use ndarray::{array, s, Array, Array2, Array3, Array4, ArrayView, Axis, Dimension};
4171    use num_traits::{AsPrimitive, Float, PrimInt};
4172    use rand::{RngExt, SeedableRng};
4173    use rand_distr::StandardNormal;
4174
4175    use crate::{
4176        configs::{self, DimName},
4177        dequantize_ndarray, BoundingBox, DecoderBuilder, DetectBox, Quantization,
4178    };
4179
4180    pub fn quantize_ndarray<T: PrimInt + 'static, D: Dimension, F: Float + AsPrimitive<T>>(
4181        input: ArrayView<F, D>,
4182        quant: Quantization,
4183    ) -> Array<T, D>
4184    where
4185        i32: num_traits::AsPrimitive<F>,
4186        f32: num_traits::AsPrimitive<F>,
4187    {
4188        let zero_point = quant.zero_point.as_();
4189        let div_scale = F::one() / quant.scale.as_();
4190        if zero_point != F::zero() {
4191            input.mapv(|d| (d * div_scale + zero_point).round().as_())
4192        } else {
4193            input.mapv(|d| (d * div_scale).round().as_())
4194        }
4195    }
4196
4197    #[test]
4198    fn test_decoder_tracked_random_jitter() {
4199        use crate::configs::{DecoderType, Nms};
4200        use crate::DecoderBuilder;
4201
4202        let score_threshold = 0.25;
4203        let iou_threshold = 0.1;
4204        let out = edgefirst_bench::testdata::read("yolov8s_80_classes.bin");
4205        let out = unsafe { std::slice::from_raw_parts(out.as_ptr() as *const i8, out.len()) };
4206        let out = Array3::from_shape_vec((1, 84, 8400), out.to_vec()).unwrap();
4207        let quant = (0.0040811873, -123).into();
4208
4209        let decoder = DecoderBuilder::default()
4210            .with_config_yolo_det(
4211                crate::configs::Detection {
4212                    decoder: DecoderType::Ultralytics,
4213                    shape: vec![1, 84, 8400],
4214                    anchors: None,
4215                    quantization: Some(quant),
4216                    dshape: vec![
4217                        (crate::configs::DimName::Batch, 1),
4218                        (crate::configs::DimName::NumFeatures, 84),
4219                        (crate::configs::DimName::NumBoxes, 8400),
4220                    ],
4221                    normalized: Some(true),
4222                },
4223                None,
4224            )
4225            .with_score_threshold(score_threshold)
4226            .with_iou_threshold(iou_threshold)
4227            .with_nms(Some(Nms::ClassAgnostic))
4228            .build()
4229            .unwrap();
4230        let mut rng = rand::rngs::StdRng::seed_from_u64(0xAB_BEEF); // fixed seed for reproducibility
4231
4232        let expected_boxes = [
4233            crate::DetectBox {
4234                bbox: crate::BoundingBox {
4235                    xmin: 0.5285137,
4236                    ymin: 0.05305544,
4237                    xmax: 0.87541467,
4238                    ymax: 0.9998909,
4239                },
4240                score: 0.5591227,
4241                label: 0,
4242            },
4243            crate::DetectBox {
4244                bbox: crate::BoundingBox {
4245                    xmin: 0.130598,
4246                    ymin: 0.43260583,
4247                    xmax: 0.35098213,
4248                    ymax: 0.9958097,
4249                },
4250                score: 0.33057618,
4251                label: 75,
4252            },
4253        ];
4254
4255        let mut tracker = ByteTrackBuilder::new()
4256            .track_update(0.1)
4257            .track_high_conf(0.3)
4258            .build();
4259
4260        let mut output_boxes = Vec::with_capacity(50);
4261        let mut output_masks = Vec::with_capacity(50);
4262        let mut output_tracks = Vec::with_capacity(50);
4263
4264        decoder
4265            .decode_tracked_quantized(
4266                &mut tracker,
4267                0,
4268                &[out.view().into()],
4269                &mut output_boxes,
4270                &mut output_masks,
4271                &mut output_tracks,
4272            )
4273            .unwrap();
4274
4275        assert_eq!(output_boxes.len(), 2);
4276        assert!(output_boxes[0].equal_within_delta(&expected_boxes[0], 1e-6));
4277        assert!(output_boxes[1].equal_within_delta(&expected_boxes[1], 1e-6));
4278
4279        let mut last_boxes = output_boxes.clone();
4280
4281        for i in 1..=100 {
4282            let mut out = out.clone();
4283            // introduce jitter into the XY coordinates to simulate movement and test tracking stability
4284            let mut x_values = out.slice_mut(s![0, 0, ..]);
4285            for x in x_values.iter_mut() {
4286                let r: f32 = rng.sample(StandardNormal);
4287                let r = r.clamp(-2.0, 2.0) / 2.0;
4288                *x = x.saturating_add((r * 1e-2 / quant.0) as i8);
4289            }
4290
4291            let mut y_values = out.slice_mut(s![0, 1, ..]);
4292            for y in y_values.iter_mut() {
4293                let r: f32 = rng.sample(StandardNormal);
4294                let r = r.clamp(-2.0, 2.0) / 2.0;
4295                *y = y.saturating_add((r * 1e-2 / quant.0) as i8);
4296            }
4297
4298            decoder
4299                .decode_tracked_quantized(
4300                    &mut tracker,
4301                    100_000_000 * i / 3, // simulate 33.333ms between frames
4302                    &[out.view().into()],
4303                    &mut output_boxes,
4304                    &mut output_masks,
4305                    &mut output_tracks,
4306                )
4307                .unwrap();
4308
4309            assert_eq!(output_boxes.len(), 2);
4310            assert!(output_boxes[0].equal_within_delta(&expected_boxes[0], 5e-3));
4311            assert!(output_boxes[1].equal_within_delta(&expected_boxes[1], 5e-3));
4312
4313            assert!(output_boxes[0].equal_within_delta(&last_boxes[0], 1e-3));
4314            assert!(output_boxes[1].equal_within_delta(&last_boxes[1], 1e-3));
4315            last_boxes = output_boxes.clone();
4316        }
4317    }
4318
4319    // ─── Shared helpers for tracked decoder tests ────────────────────
4320
4321    fn real_data_expected_boxes() -> [DetectBox; 2] {
4322        [
4323            DetectBox {
4324                bbox: BoundingBox {
4325                    xmin: 0.08515105,
4326                    ymin: 0.7131401,
4327                    xmax: 0.29802868,
4328                    ymax: 0.8195788,
4329                },
4330                score: 0.91537374,
4331                label: 23,
4332            },
4333            DetectBox {
4334                bbox: BoundingBox {
4335                    xmin: 0.59605736,
4336                    ymin: 0.25545314,
4337                    xmax: 0.93666154,
4338                    ymax: 0.72378385,
4339                },
4340                score: 0.91537374,
4341                label: 23,
4342            },
4343        ]
4344    }
4345
4346    fn e2e_expected_boxes_quant() -> [DetectBox; 1] {
4347        [DetectBox {
4348            bbox: BoundingBox {
4349                xmin: 0.12549022,
4350                ymin: 0.12549022,
4351                xmax: 0.23529413,
4352                ymax: 0.23529413,
4353            },
4354            score: 0.98823535,
4355            label: 2,
4356        }]
4357    }
4358
4359    fn e2e_expected_boxes_float() -> [DetectBox; 1] {
4360        [DetectBox {
4361            bbox: BoundingBox {
4362                xmin: 0.1234,
4363                ymin: 0.1234,
4364                xmax: 0.2345,
4365                ymax: 0.2345,
4366            },
4367            score: 0.9876,
4368            label: 2,
4369        }]
4370    }
4371
4372    fn build_yolo_split_segdet_decoder(
4373        score_threshold: f32,
4374        iou_threshold: f32,
4375        quant_boxes: (f32, i32),
4376        quant_protos: (f32, i32),
4377    ) -> crate::Decoder {
4378        DecoderBuilder::default()
4379            .with_config_yolo_split_segdet(
4380                configs::Boxes {
4381                    decoder: configs::DecoderType::Ultralytics,
4382                    quantization: Some(quant_boxes.into()),
4383                    shape: vec![1, 4, 8400],
4384                    dshape: vec![
4385                        (DimName::Batch, 1),
4386                        (DimName::BoxCoords, 4),
4387                        (DimName::NumBoxes, 8400),
4388                    ],
4389                    normalized: Some(true),
4390                },
4391                configs::Scores {
4392                    decoder: configs::DecoderType::Ultralytics,
4393                    quantization: Some(quant_boxes.into()),
4394                    shape: vec![1, 80, 8400],
4395                    dshape: vec![
4396                        (DimName::Batch, 1),
4397                        (DimName::NumClasses, 80),
4398                        (DimName::NumBoxes, 8400),
4399                    ],
4400                },
4401                configs::MaskCoefficients {
4402                    decoder: configs::DecoderType::Ultralytics,
4403                    quantization: Some(quant_boxes.into()),
4404                    shape: vec![1, 32, 8400],
4405                    dshape: vec![
4406                        (DimName::Batch, 1),
4407                        (DimName::NumProtos, 32),
4408                        (DimName::NumBoxes, 8400),
4409                    ],
4410                },
4411                configs::Protos {
4412                    decoder: configs::DecoderType::Ultralytics,
4413                    quantization: Some(quant_protos.into()),
4414                    shape: vec![1, 160, 160, 32],
4415                    dshape: vec![
4416                        (DimName::Batch, 1),
4417                        (DimName::Height, 160),
4418                        (DimName::Width, 160),
4419                        (DimName::NumProtos, 32),
4420                    ],
4421                },
4422            )
4423            .with_score_threshold(score_threshold)
4424            .with_iou_threshold(iou_threshold)
4425            .build()
4426            .unwrap()
4427    }
4428
4429    fn build_yolov8_seg_decoder(score_threshold: f32, iou_threshold: f32) -> crate::Decoder {
4430        let config_yaml = edgefirst_bench::testdata::read_to_string("yolov8_seg.yaml");
4431        DecoderBuilder::default()
4432            .with_config_yaml_str(config_yaml.to_string())
4433            .with_score_threshold(score_threshold)
4434            .with_iou_threshold(iou_threshold)
4435            .build()
4436            .unwrap()
4437    }
4438
4439    // ─── Real-data tracked test macro ───────────────────────────────
4440    //
4441    // Generates tests that load i8 binary test data from testdata/ and
4442    // exercise all (quant/float) × (combined/split) × (masks/proto)
4443    // decoder paths.
4444
4445    macro_rules! real_data_tracked_test {
4446        ($name:ident, quantized, $layout:ident, $output:ident) => {
4447            #[test]
4448            fn $name() {
4449                let is_split = matches!(stringify!($layout), "split");
4450                let is_proto = matches!(stringify!($output), "proto");
4451
4452                let score_threshold = 0.45;
4453                let iou_threshold = 0.45;
4454                let quant_boxes = (0.021287762_f32, 31_i32);
4455                let quant_protos = (0.02491162_f32, -117_i32);
4456
4457                let raw_boxes = edgefirst_bench::testdata::read("yolov8_boxes_116x8400.bin");
4458                let raw_boxes = unsafe {
4459                    std::slice::from_raw_parts(raw_boxes.as_ptr() as *const i8, raw_boxes.len())
4460                };
4461                let boxes_i8 =
4462                    ndarray::Array3::from_shape_vec((1, 116, 8400), raw_boxes.to_vec()).unwrap();
4463
4464                let raw_protos = edgefirst_bench::testdata::read("yolov8_protos_160x160x32.bin");
4465                let raw_protos = unsafe {
4466                    std::slice::from_raw_parts(raw_protos.as_ptr() as *const i8, raw_protos.len())
4467                };
4468                let protos_i8 =
4469                    ndarray::Array4::from_shape_vec((1, 160, 160, 32), raw_protos.to_vec())
4470                        .unwrap();
4471
4472                // Pre-split (unused for combined, but harmless)
4473                let mask_split = boxes_i8.slice(s![.., 84.., ..]).to_owned();
4474                let mut scores_split = boxes_i8.slice(s![.., 4..84, ..]).to_owned();
4475                let boxes_split = boxes_i8.slice(s![.., ..4, ..]).to_owned();
4476                let mut boxes_combined = boxes_i8;
4477
4478                let decoder = if is_split {
4479                    build_yolo_split_segdet_decoder(
4480                        score_threshold,
4481                        iou_threshold,
4482                        quant_boxes,
4483                        quant_protos,
4484                    )
4485                } else {
4486                    build_yolov8_seg_decoder(score_threshold, iou_threshold)
4487                };
4488
4489                let expected = real_data_expected_boxes();
4490                let mut tracker = ByteTrackBuilder::new()
4491                    .track_update(0.1)
4492                    .track_high_conf(0.7)
4493                    .build();
4494                let mut output_boxes = Vec::with_capacity(50);
4495                let mut output_tracks = Vec::with_capacity(50);
4496
4497                // Frame 1: decode
4498                if is_proto {
4499                    {
4500                        let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> = if is_split {
4501                            vec![
4502                                boxes_split.view().into(),
4503                                scores_split.view().into(),
4504                                mask_split.view().into(),
4505                                protos_i8.view().into(),
4506                            ]
4507                        } else {
4508                            vec![boxes_combined.view().into(), protos_i8.view().into()]
4509                        };
4510                        decoder
4511                            .decode_tracked_quantized_proto(
4512                                &mut tracker,
4513                                0,
4514                                &inputs,
4515                                &mut output_boxes,
4516                                &mut output_tracks,
4517                            )
4518                            .unwrap();
4519                    }
4520                    assert_eq!(output_boxes.len(), 2);
4521                    assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
4522                    assert!(output_boxes[1].equal_within_delta(&expected[1], 1.0 / 160.0));
4523
4524                    // Zero scores for frame 2
4525                    if is_split {
4526                        for score in scores_split.iter_mut() {
4527                            *score = i8::MIN;
4528                        }
4529                    } else {
4530                        for score in boxes_combined.slice_mut(s![0, 4..84, ..]).iter_mut() {
4531                            *score = i8::MIN;
4532                        }
4533                    }
4534
4535                    let proto_result = {
4536                        let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> = if is_split {
4537                            vec![
4538                                boxes_split.view().into(),
4539                                scores_split.view().into(),
4540                                mask_split.view().into(),
4541                                protos_i8.view().into(),
4542                            ]
4543                        } else {
4544                            vec![boxes_combined.view().into(), protos_i8.view().into()]
4545                        };
4546                        decoder
4547                            .decode_tracked_quantized_proto(
4548                                &mut tracker,
4549                                100_000_000 / 3,
4550                                &inputs,
4551                                &mut output_boxes,
4552                                &mut output_tracks,
4553                            )
4554                            .unwrap()
4555                    };
4556                    assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
4557                    assert!(output_boxes[1].equal_within_delta(&expected[1], 1e-6));
4558                    assert!(proto_result.is_some_and(|x| x.mask_coefficients.shape()[0] == 0));
4559                } else {
4560                    let mut output_masks = Vec::with_capacity(50);
4561                    {
4562                        let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> = if is_split {
4563                            vec![
4564                                boxes_split.view().into(),
4565                                scores_split.view().into(),
4566                                mask_split.view().into(),
4567                                protos_i8.view().into(),
4568                            ]
4569                        } else {
4570                            vec![boxes_combined.view().into(), protos_i8.view().into()]
4571                        };
4572                        decoder
4573                            .decode_tracked_quantized(
4574                                &mut tracker,
4575                                0,
4576                                &inputs,
4577                                &mut output_boxes,
4578                                &mut output_masks,
4579                                &mut output_tracks,
4580                            )
4581                            .unwrap();
4582                    }
4583                    assert_eq!(output_boxes.len(), 2);
4584                    assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
4585                    assert!(output_boxes[1].equal_within_delta(&expected[1], 1.0 / 160.0));
4586
4587                    if is_split {
4588                        for score in scores_split.iter_mut() {
4589                            *score = i8::MIN;
4590                        }
4591                    } else {
4592                        for score in boxes_combined.slice_mut(s![0, 4..84, ..]).iter_mut() {
4593                            *score = i8::MIN;
4594                        }
4595                    }
4596
4597                    {
4598                        let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> = if is_split {
4599                            vec![
4600                                boxes_split.view().into(),
4601                                scores_split.view().into(),
4602                                mask_split.view().into(),
4603                                protos_i8.view().into(),
4604                            ]
4605                        } else {
4606                            vec![boxes_combined.view().into(), protos_i8.view().into()]
4607                        };
4608                        decoder
4609                            .decode_tracked_quantized(
4610                                &mut tracker,
4611                                100_000_000 / 3,
4612                                &inputs,
4613                                &mut output_boxes,
4614                                &mut output_masks,
4615                                &mut output_tracks,
4616                            )
4617                            .unwrap();
4618                    }
4619                    assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
4620                    assert!(output_boxes[1].equal_within_delta(&expected[1], 1e-6));
4621                    assert!(output_masks.is_empty());
4622                }
4623            }
4624        };
4625        ($name:ident, float, $layout:ident, $output:ident) => {
4626            #[test]
4627            fn $name() {
4628                let is_split = matches!(stringify!($layout), "split");
4629                let is_proto = matches!(stringify!($output), "proto");
4630
4631                let score_threshold = 0.45;
4632                let iou_threshold = 0.45;
4633                let quant_boxes = (0.021287762_f32, 31_i32);
4634                let quant_protos = (0.02491162_f32, -117_i32);
4635
4636                let raw_boxes = edgefirst_bench::testdata::read("yolov8_boxes_116x8400.bin");
4637                let raw_boxes = unsafe {
4638                    std::slice::from_raw_parts(raw_boxes.as_ptr() as *const i8, raw_boxes.len())
4639                };
4640                let boxes_i8 =
4641                    ndarray::Array3::from_shape_vec((1, 116, 8400), raw_boxes.to_vec()).unwrap();
4642                let boxes_f32 = dequantize_ndarray(boxes_i8.view(), quant_boxes.into());
4643
4644                let raw_protos = edgefirst_bench::testdata::read("yolov8_protos_160x160x32.bin");
4645                let raw_protos = unsafe {
4646                    std::slice::from_raw_parts(raw_protos.as_ptr() as *const i8, raw_protos.len())
4647                };
4648                let protos_i8 =
4649                    ndarray::Array4::from_shape_vec((1, 160, 160, 32), raw_protos.to_vec())
4650                        .unwrap();
4651                let protos_f32 = dequantize_ndarray(protos_i8.view(), quant_protos.into());
4652
4653                // Pre-split from dequantized data
4654                let mask_split = boxes_f32.slice(s![.., 84.., ..]).to_owned();
4655                let mut scores_split = boxes_f32.slice(s![.., 4..84, ..]).to_owned();
4656                let boxes_split = boxes_f32.slice(s![.., ..4, ..]).to_owned();
4657                let mut boxes_combined = boxes_f32;
4658
4659                let decoder = if is_split {
4660                    build_yolo_split_segdet_decoder(
4661                        score_threshold,
4662                        iou_threshold,
4663                        quant_boxes,
4664                        quant_protos,
4665                    )
4666                } else {
4667                    build_yolov8_seg_decoder(score_threshold, iou_threshold)
4668                };
4669
4670                let expected = real_data_expected_boxes();
4671                let mut tracker = ByteTrackBuilder::new()
4672                    .track_update(0.1)
4673                    .track_high_conf(0.7)
4674                    .build();
4675                let mut output_boxes = Vec::with_capacity(50);
4676                let mut output_tracks = Vec::with_capacity(50);
4677
4678                if is_proto {
4679                    {
4680                        let inputs = if is_split {
4681                            vec![
4682                                boxes_split.view().into_dyn(),
4683                                scores_split.view().into_dyn(),
4684                                mask_split.view().into_dyn(),
4685                                protos_f32.view().into_dyn(),
4686                            ]
4687                        } else {
4688                            vec![
4689                                boxes_combined.view().into_dyn(),
4690                                protos_f32.view().into_dyn(),
4691                            ]
4692                        };
4693                        decoder
4694                            .decode_tracked_float_proto(
4695                                &mut tracker,
4696                                0,
4697                                &inputs,
4698                                &mut output_boxes,
4699                                &mut output_tracks,
4700                            )
4701                            .unwrap();
4702                    }
4703                    assert_eq!(output_boxes.len(), 2);
4704                    assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
4705                    assert!(output_boxes[1].equal_within_delta(&expected[1], 1.0 / 160.0));
4706
4707                    if is_split {
4708                        for score in scores_split.iter_mut() {
4709                            *score = 0.0;
4710                        }
4711                    } else {
4712                        for score in boxes_combined.slice_mut(s![0, 4..84, ..]).iter_mut() {
4713                            *score = 0.0;
4714                        }
4715                    }
4716
4717                    let proto_result = {
4718                        let inputs = if is_split {
4719                            vec![
4720                                boxes_split.view().into_dyn(),
4721                                scores_split.view().into_dyn(),
4722                                mask_split.view().into_dyn(),
4723                                protos_f32.view().into_dyn(),
4724                            ]
4725                        } else {
4726                            vec![
4727                                boxes_combined.view().into_dyn(),
4728                                protos_f32.view().into_dyn(),
4729                            ]
4730                        };
4731                        decoder
4732                            .decode_tracked_float_proto(
4733                                &mut tracker,
4734                                100_000_000 / 3,
4735                                &inputs,
4736                                &mut output_boxes,
4737                                &mut output_tracks,
4738                            )
4739                            .unwrap()
4740                    };
4741                    assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
4742                    assert!(output_boxes[1].equal_within_delta(&expected[1], 1e-6));
4743                    assert!(proto_result.is_some_and(|x| x.mask_coefficients.shape()[0] == 0));
4744                } else {
4745                    let mut output_masks = Vec::with_capacity(50);
4746                    {
4747                        let inputs = if is_split {
4748                            vec![
4749                                boxes_split.view().into_dyn(),
4750                                scores_split.view().into_dyn(),
4751                                mask_split.view().into_dyn(),
4752                                protos_f32.view().into_dyn(),
4753                            ]
4754                        } else {
4755                            vec![
4756                                boxes_combined.view().into_dyn(),
4757                                protos_f32.view().into_dyn(),
4758                            ]
4759                        };
4760                        decoder
4761                            .decode_tracked_float(
4762                                &mut tracker,
4763                                0,
4764                                &inputs,
4765                                &mut output_boxes,
4766                                &mut output_masks,
4767                                &mut output_tracks,
4768                            )
4769                            .unwrap();
4770                    }
4771                    assert_eq!(output_boxes.len(), 2);
4772                    assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
4773                    assert!(output_boxes[1].equal_within_delta(&expected[1], 1.0 / 160.0));
4774
4775                    if is_split {
4776                        for score in scores_split.iter_mut() {
4777                            *score = 0.0;
4778                        }
4779                    } else {
4780                        for score in boxes_combined.slice_mut(s![0, 4..84, ..]).iter_mut() {
4781                            *score = 0.0;
4782                        }
4783                    }
4784
4785                    {
4786                        let inputs = if is_split {
4787                            vec![
4788                                boxes_split.view().into_dyn(),
4789                                scores_split.view().into_dyn(),
4790                                mask_split.view().into_dyn(),
4791                                protos_f32.view().into_dyn(),
4792                            ]
4793                        } else {
4794                            vec![
4795                                boxes_combined.view().into_dyn(),
4796                                protos_f32.view().into_dyn(),
4797                            ]
4798                        };
4799                        decoder
4800                            .decode_tracked_float(
4801                                &mut tracker,
4802                                100_000_000 / 3,
4803                                &inputs,
4804                                &mut output_boxes,
4805                                &mut output_masks,
4806                                &mut output_tracks,
4807                            )
4808                            .unwrap();
4809                    }
4810                    assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
4811                    assert!(output_boxes[1].equal_within_delta(&expected[1], 1e-6));
4812                    assert!(output_masks.is_empty());
4813                }
4814            }
4815        };
4816    }
4817
4818    real_data_tracked_test!(test_decoder_tracked_segdet, quantized, combined, masks);
4819    real_data_tracked_test!(test_decoder_tracked_segdet_float, float, combined, masks);
4820    real_data_tracked_test!(
4821        test_decoder_tracked_segdet_proto,
4822        quantized,
4823        combined,
4824        proto
4825    );
4826    real_data_tracked_test!(
4827        test_decoder_tracked_segdet_proto_float,
4828        float,
4829        combined,
4830        proto
4831    );
4832    real_data_tracked_test!(test_decoder_tracked_segdet_split, quantized, split, masks);
4833    real_data_tracked_test!(test_decoder_tracked_segdet_split_float, float, split, masks);
4834    real_data_tracked_test!(
4835        test_decoder_tracked_segdet_split_proto,
4836        quantized,
4837        split,
4838        proto
4839    );
4840    real_data_tracked_test!(
4841        test_decoder_tracked_segdet_split_proto_float,
4842        float,
4843        split,
4844        proto
4845    );
4846
4847    // ─── End-to-end tracked test macro ──────────────────────────────
4848    //
4849    // Generates tests with synthetic data to exercise all tracked
4850    // decode paths without needing real model output files.
4851
4852    const E2E_COMBINED_CONFIG: &str = "
4853decoder_version: yolo26
4854outputs:
4855 - type: detection
4856   decoder: ultralytics
4857   quantization: [0.00784313725490196, 0]
4858   shape: [1, 10, 38]
4859   dshape:
4860    - [batch, 1]
4861    - [num_boxes, 10]
4862    - [num_features, 38]
4863   normalized: true
4864 - type: protos
4865   decoder: ultralytics
4866   quantization: [0.0039215686274509803921568627451, 128]
4867   shape: [1, 160, 160, 32]
4868   dshape:
4869    - [batch, 1]
4870    - [height, 160]
4871    - [width, 160]
4872    - [num_protos, 32]
4873";
4874
4875    const E2E_SPLIT_CONFIG: &str = "
4876decoder_version: yolo26
4877outputs:
4878 - type: boxes
4879   decoder: ultralytics
4880   quantization: [0.00784313725490196, 0]
4881   shape: [1, 10, 4]
4882   dshape:
4883    - [batch, 1]
4884    - [num_boxes, 10]
4885    - [box_coords, 4]
4886   normalized: true
4887 - type: scores
4888   decoder: ultralytics
4889   quantization: [0.00784313725490196, 0]
4890   shape: [1, 10, 1]
4891   dshape:
4892    - [batch, 1]
4893    - [num_boxes, 10]
4894    - [num_classes, 1]
4895 - type: classes
4896   decoder: ultralytics
4897   quantization: [0.00784313725490196, 0]
4898   shape: [1, 10, 1]
4899   dshape:
4900    - [batch, 1]
4901    - [num_boxes, 10]
4902    - [num_classes, 1]
4903 - type: mask_coefficients
4904   decoder: ultralytics
4905   quantization: [0.00784313725490196, 0]
4906   shape: [1, 10, 32]
4907   dshape:
4908    - [batch, 1]
4909    - [num_boxes, 10]
4910    - [num_protos, 32]
4911 - type: protos
4912   decoder: ultralytics
4913   quantization: [0.0039215686274509803921568627451, 128]
4914   shape: [1, 160, 160, 32]
4915   dshape:
4916    - [batch, 1]
4917    - [height, 160]
4918    - [width, 160]
4919    - [num_protos, 32]
4920";
4921
4922    macro_rules! e2e_tracked_test {
4923        ($name:ident, quantized, $layout:ident, $output:ident) => {
4924            #[test]
4925            fn $name() {
4926                let is_split = matches!(stringify!($layout), "split");
4927                let is_proto = matches!(stringify!($output), "proto");
4928
4929                let score_threshold = 0.45;
4930                let iou_threshold = 0.45;
4931
4932                let mut boxes = Array2::zeros((10, 4));
4933                let mut scores = Array2::zeros((10, 1));
4934                let mut classes = Array2::zeros((10, 1));
4935                let mask = Array2::zeros((10, 32));
4936                let protos = Array3::<f64>::zeros((160, 160, 32));
4937                let protos = protos.insert_axis(Axis(0));
4938                let protos_quant = (1.0 / 255.0, 0.0);
4939                let protos: Array4<u8> = quantize_ndarray(protos.view(), protos_quant.into());
4940
4941                boxes
4942                    .slice_mut(s![0, ..])
4943                    .assign(&array![0.1234, 0.1234, 0.2345, 0.2345]);
4944                scores.slice_mut(s![0, ..]).assign(&array![0.9876]);
4945                classes.slice_mut(s![0, ..]).assign(&array![2.0]);
4946
4947                let detect_quant = (2.0 / 255.0, 0.0);
4948
4949                let decoder = if is_split {
4950                    DecoderBuilder::default()
4951                        .with_config_yaml_str(E2E_SPLIT_CONFIG.to_string())
4952                        .with_score_threshold(score_threshold)
4953                        .with_iou_threshold(iou_threshold)
4954                        .build()
4955                        .unwrap()
4956                } else {
4957                    DecoderBuilder::default()
4958                        .with_config_yaml_str(E2E_COMBINED_CONFIG.to_string())
4959                        .with_score_threshold(score_threshold)
4960                        .with_iou_threshold(iou_threshold)
4961                        .build()
4962                        .unwrap()
4963                };
4964
4965                let expected = e2e_expected_boxes_quant();
4966                let mut tracker = ByteTrackBuilder::new()
4967                    .track_update(0.1)
4968                    .track_high_conf(0.7)
4969                    .build();
4970                let mut output_boxes = Vec::with_capacity(50);
4971                let mut output_tracks = Vec::with_capacity(50);
4972
4973                if is_split {
4974                    let boxes = boxes.insert_axis(Axis(0));
4975                    let scores = scores.insert_axis(Axis(0));
4976                    let classes = classes.insert_axis(Axis(0));
4977                    let mask = mask.insert_axis(Axis(0));
4978
4979                    let boxes: Array3<u8> = quantize_ndarray(boxes.view(), detect_quant.into());
4980                    let mut scores: Array3<u8> =
4981                        quantize_ndarray(scores.view(), detect_quant.into());
4982                    let classes: Array3<u8> = quantize_ndarray(classes.view(), detect_quant.into());
4983                    let mask: Array3<u8> = quantize_ndarray(mask.view(), detect_quant.into());
4984
4985                    if is_proto {
4986                        {
4987                            let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> = vec![
4988                                boxes.view().into(),
4989                                scores.view().into(),
4990                                classes.view().into(),
4991                                mask.view().into(),
4992                                protos.view().into(),
4993                            ];
4994                            decoder
4995                                .decode_tracked_quantized_proto(
4996                                    &mut tracker,
4997                                    0,
4998                                    &inputs,
4999                                    &mut output_boxes,
5000                                    &mut output_tracks,
5001                                )
5002                                .unwrap();
5003                        }
5004                        assert_eq!(output_boxes.len(), 1);
5005                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5006
5007                        for score in scores.slice_mut(s![.., .., ..]).iter_mut() {
5008                            *score = u8::MIN;
5009                        }
5010                        let proto_result = {
5011                            let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> = vec![
5012                                boxes.view().into(),
5013                                scores.view().into(),
5014                                classes.view().into(),
5015                                mask.view().into(),
5016                                protos.view().into(),
5017                            ];
5018                            decoder
5019                                .decode_tracked_quantized_proto(
5020                                    &mut tracker,
5021                                    100_000_000 / 3,
5022                                    &inputs,
5023                                    &mut output_boxes,
5024                                    &mut output_tracks,
5025                                )
5026                                .unwrap()
5027                        };
5028                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5029                        assert!(proto_result.is_some_and(|x| x.mask_coefficients.shape()[0] == 0));
5030                    } else {
5031                        let mut output_masks = Vec::with_capacity(50);
5032                        {
5033                            let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> = vec![
5034                                boxes.view().into(),
5035                                scores.view().into(),
5036                                classes.view().into(),
5037                                mask.view().into(),
5038                                protos.view().into(),
5039                            ];
5040                            decoder
5041                                .decode_tracked_quantized(
5042                                    &mut tracker,
5043                                    0,
5044                                    &inputs,
5045                                    &mut output_boxes,
5046                                    &mut output_masks,
5047                                    &mut output_tracks,
5048                                )
5049                                .unwrap();
5050                        }
5051                        assert_eq!(output_boxes.len(), 1);
5052                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5053
5054                        for score in scores.slice_mut(s![.., .., ..]).iter_mut() {
5055                            *score = u8::MIN;
5056                        }
5057                        {
5058                            let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> = vec![
5059                                boxes.view().into(),
5060                                scores.view().into(),
5061                                classes.view().into(),
5062                                mask.view().into(),
5063                                protos.view().into(),
5064                            ];
5065                            decoder
5066                                .decode_tracked_quantized(
5067                                    &mut tracker,
5068                                    100_000_000 / 3,
5069                                    &inputs,
5070                                    &mut output_boxes,
5071                                    &mut output_masks,
5072                                    &mut output_tracks,
5073                                )
5074                                .unwrap();
5075                        }
5076                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5077                        assert!(output_masks.is_empty());
5078                    }
5079                } else {
5080                    // Combined layout
5081                    let detect = ndarray::concatenate![
5082                        Axis(1),
5083                        boxes.view(),
5084                        scores.view(),
5085                        classes.view(),
5086                        mask.view()
5087                    ];
5088                    let detect = detect.insert_axis(Axis(0));
5089                    assert_eq!(detect.shape(), &[1, 10, 38]);
5090                    let mut detect: Array3<u8> =
5091                        quantize_ndarray(detect.view(), detect_quant.into());
5092
5093                    if is_proto {
5094                        {
5095                            let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> =
5096                                vec![detect.view().into(), protos.view().into()];
5097                            decoder
5098                                .decode_tracked_quantized_proto(
5099                                    &mut tracker,
5100                                    0,
5101                                    &inputs,
5102                                    &mut output_boxes,
5103                                    &mut output_tracks,
5104                                )
5105                                .unwrap();
5106                        }
5107                        assert_eq!(output_boxes.len(), 1);
5108                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5109
5110                        for score in detect.slice_mut(s![.., .., 4]).iter_mut() {
5111                            *score = u8::MIN;
5112                        }
5113                        let proto_result = {
5114                            let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> =
5115                                vec![detect.view().into(), protos.view().into()];
5116                            decoder
5117                                .decode_tracked_quantized_proto(
5118                                    &mut tracker,
5119                                    100_000_000 / 3,
5120                                    &inputs,
5121                                    &mut output_boxes,
5122                                    &mut output_tracks,
5123                                )
5124                                .unwrap()
5125                        };
5126                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5127                        assert!(proto_result.is_some_and(|x| x.mask_coefficients.shape()[0] == 0));
5128                    } else {
5129                        let mut output_masks = Vec::with_capacity(50);
5130                        {
5131                            let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> =
5132                                vec![detect.view().into(), protos.view().into()];
5133                            decoder
5134                                .decode_tracked_quantized(
5135                                    &mut tracker,
5136                                    0,
5137                                    &inputs,
5138                                    &mut output_boxes,
5139                                    &mut output_masks,
5140                                    &mut output_tracks,
5141                                )
5142                                .unwrap();
5143                        }
5144                        assert_eq!(output_boxes.len(), 1);
5145                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5146
5147                        for score in detect.slice_mut(s![.., .., 4]).iter_mut() {
5148                            *score = u8::MIN;
5149                        }
5150                        {
5151                            let inputs: Vec<crate::decoder::ArrayViewDQuantized<'_>> =
5152                                vec![detect.view().into(), protos.view().into()];
5153                            decoder
5154                                .decode_tracked_quantized(
5155                                    &mut tracker,
5156                                    100_000_000 / 3,
5157                                    &inputs,
5158                                    &mut output_boxes,
5159                                    &mut output_masks,
5160                                    &mut output_tracks,
5161                                )
5162                                .unwrap();
5163                        }
5164                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5165                        assert!(output_masks.is_empty());
5166                    }
5167                }
5168            }
5169        };
5170        ($name:ident, float, $layout:ident, $output:ident) => {
5171            #[test]
5172            fn $name() {
5173                let is_split = matches!(stringify!($layout), "split");
5174                let is_proto = matches!(stringify!($output), "proto");
5175
5176                let score_threshold = 0.45;
5177                let iou_threshold = 0.45;
5178
5179                let mut boxes = Array2::zeros((10, 4));
5180                let mut scores = Array2::zeros((10, 1));
5181                let mut classes = Array2::zeros((10, 1));
5182                let mask: Array2<f64> = Array2::zeros((10, 32));
5183                let protos = Array3::<f64>::zeros((160, 160, 32));
5184                let protos = protos.insert_axis(Axis(0));
5185
5186                boxes
5187                    .slice_mut(s![0, ..])
5188                    .assign(&array![0.1234, 0.1234, 0.2345, 0.2345]);
5189                scores.slice_mut(s![0, ..]).assign(&array![0.9876]);
5190                classes.slice_mut(s![0, ..]).assign(&array![2.0]);
5191
5192                let decoder = if is_split {
5193                    DecoderBuilder::default()
5194                        .with_config_yaml_str(E2E_SPLIT_CONFIG.to_string())
5195                        .with_score_threshold(score_threshold)
5196                        .with_iou_threshold(iou_threshold)
5197                        .build()
5198                        .unwrap()
5199                } else {
5200                    DecoderBuilder::default()
5201                        .with_config_yaml_str(E2E_COMBINED_CONFIG.to_string())
5202                        .with_score_threshold(score_threshold)
5203                        .with_iou_threshold(iou_threshold)
5204                        .build()
5205                        .unwrap()
5206                };
5207
5208                let expected = e2e_expected_boxes_float();
5209                let mut tracker = ByteTrackBuilder::new()
5210                    .track_update(0.1)
5211                    .track_high_conf(0.7)
5212                    .build();
5213                let mut output_boxes = Vec::with_capacity(50);
5214                let mut output_tracks = Vec::with_capacity(50);
5215
5216                if is_split {
5217                    let boxes = boxes.insert_axis(Axis(0));
5218                    let mut scores = scores.insert_axis(Axis(0));
5219                    let classes = classes.insert_axis(Axis(0));
5220                    let mask = mask.insert_axis(Axis(0));
5221
5222                    if is_proto {
5223                        {
5224                            let inputs = vec![
5225                                boxes.view().into_dyn(),
5226                                scores.view().into_dyn(),
5227                                classes.view().into_dyn(),
5228                                mask.view().into_dyn(),
5229                                protos.view().into_dyn(),
5230                            ];
5231                            decoder
5232                                .decode_tracked_float_proto(
5233                                    &mut tracker,
5234                                    0,
5235                                    &inputs,
5236                                    &mut output_boxes,
5237                                    &mut output_tracks,
5238                                )
5239                                .unwrap();
5240                        }
5241                        assert_eq!(output_boxes.len(), 1);
5242                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5243
5244                        for score in scores.slice_mut(s![.., .., ..]).iter_mut() {
5245                            *score = 0.0;
5246                        }
5247                        let proto_result = {
5248                            let inputs = vec![
5249                                boxes.view().into_dyn(),
5250                                scores.view().into_dyn(),
5251                                classes.view().into_dyn(),
5252                                mask.view().into_dyn(),
5253                                protos.view().into_dyn(),
5254                            ];
5255                            decoder
5256                                .decode_tracked_float_proto(
5257                                    &mut tracker,
5258                                    100_000_000 / 3,
5259                                    &inputs,
5260                                    &mut output_boxes,
5261                                    &mut output_tracks,
5262                                )
5263                                .unwrap()
5264                        };
5265                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5266                        assert!(proto_result.is_some_and(|x| x.mask_coefficients.shape()[0] == 0));
5267                    } else {
5268                        let mut output_masks = Vec::with_capacity(50);
5269                        {
5270                            let inputs = vec![
5271                                boxes.view().into_dyn(),
5272                                scores.view().into_dyn(),
5273                                classes.view().into_dyn(),
5274                                mask.view().into_dyn(),
5275                                protos.view().into_dyn(),
5276                            ];
5277                            decoder
5278                                .decode_tracked_float(
5279                                    &mut tracker,
5280                                    0,
5281                                    &inputs,
5282                                    &mut output_boxes,
5283                                    &mut output_masks,
5284                                    &mut output_tracks,
5285                                )
5286                                .unwrap();
5287                        }
5288                        assert_eq!(output_boxes.len(), 1);
5289                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5290
5291                        for score in scores.slice_mut(s![.., .., ..]).iter_mut() {
5292                            *score = 0.0;
5293                        }
5294                        {
5295                            let inputs = vec![
5296                                boxes.view().into_dyn(),
5297                                scores.view().into_dyn(),
5298                                classes.view().into_dyn(),
5299                                mask.view().into_dyn(),
5300                                protos.view().into_dyn(),
5301                            ];
5302                            decoder
5303                                .decode_tracked_float(
5304                                    &mut tracker,
5305                                    100_000_000 / 3,
5306                                    &inputs,
5307                                    &mut output_boxes,
5308                                    &mut output_masks,
5309                                    &mut output_tracks,
5310                                )
5311                                .unwrap();
5312                        }
5313                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5314                        assert!(output_masks.is_empty());
5315                    }
5316                } else {
5317                    // Combined layout
5318                    let detect = ndarray::concatenate![
5319                        Axis(1),
5320                        boxes.view(),
5321                        scores.view(),
5322                        classes.view(),
5323                        mask.view()
5324                    ];
5325                    let mut detect = detect.insert_axis(Axis(0));
5326                    assert_eq!(detect.shape(), &[1, 10, 38]);
5327
5328                    if is_proto {
5329                        {
5330                            let inputs = vec![detect.view().into_dyn(), protos.view().into_dyn()];
5331                            decoder
5332                                .decode_tracked_float_proto(
5333                                    &mut tracker,
5334                                    0,
5335                                    &inputs,
5336                                    &mut output_boxes,
5337                                    &mut output_tracks,
5338                                )
5339                                .unwrap();
5340                        }
5341                        assert_eq!(output_boxes.len(), 1);
5342                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5343
5344                        for score in detect.slice_mut(s![.., .., 4]).iter_mut() {
5345                            *score = 0.0;
5346                        }
5347                        let proto_result = {
5348                            let inputs = vec![detect.view().into_dyn(), protos.view().into_dyn()];
5349                            decoder
5350                                .decode_tracked_float_proto(
5351                                    &mut tracker,
5352                                    100_000_000 / 3,
5353                                    &inputs,
5354                                    &mut output_boxes,
5355                                    &mut output_tracks,
5356                                )
5357                                .unwrap()
5358                        };
5359                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5360                        assert!(proto_result.is_some_and(|x| x.mask_coefficients.shape()[0] == 0));
5361                    } else {
5362                        let mut output_masks = Vec::with_capacity(50);
5363                        {
5364                            let inputs = vec![detect.view().into_dyn(), protos.view().into_dyn()];
5365                            decoder
5366                                .decode_tracked_float(
5367                                    &mut tracker,
5368                                    0,
5369                                    &inputs,
5370                                    &mut output_boxes,
5371                                    &mut output_masks,
5372                                    &mut output_tracks,
5373                                )
5374                                .unwrap();
5375                        }
5376                        assert_eq!(output_boxes.len(), 1);
5377                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5378
5379                        for score in detect.slice_mut(s![.., .., 4]).iter_mut() {
5380                            *score = 0.0;
5381                        }
5382                        {
5383                            let inputs = vec![detect.view().into_dyn(), protos.view().into_dyn()];
5384                            decoder
5385                                .decode_tracked_float(
5386                                    &mut tracker,
5387                                    100_000_000 / 3,
5388                                    &inputs,
5389                                    &mut output_boxes,
5390                                    &mut output_masks,
5391                                    &mut output_tracks,
5392                                )
5393                                .unwrap();
5394                        }
5395                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5396                        assert!(output_masks.is_empty());
5397                    }
5398                }
5399            }
5400        };
5401    }
5402
5403    e2e_tracked_test!(
5404        test_decoder_tracked_end_to_end_segdet,
5405        quantized,
5406        combined,
5407        masks
5408    );
5409    e2e_tracked_test!(
5410        test_decoder_tracked_end_to_end_segdet_float,
5411        float,
5412        combined,
5413        masks
5414    );
5415    e2e_tracked_test!(
5416        test_decoder_tracked_end_to_end_segdet_proto,
5417        quantized,
5418        combined,
5419        proto
5420    );
5421    e2e_tracked_test!(
5422        test_decoder_tracked_end_to_end_segdet_proto_float,
5423        float,
5424        combined,
5425        proto
5426    );
5427    e2e_tracked_test!(
5428        test_decoder_tracked_end_to_end_segdet_split,
5429        quantized,
5430        split,
5431        masks
5432    );
5433    e2e_tracked_test!(
5434        test_decoder_tracked_end_to_end_segdet_split_float,
5435        float,
5436        split,
5437        masks
5438    );
5439    e2e_tracked_test!(
5440        test_decoder_tracked_end_to_end_segdet_split_proto,
5441        quantized,
5442        split,
5443        proto
5444    );
5445    e2e_tracked_test!(
5446        test_decoder_tracked_end_to_end_segdet_split_proto_float,
5447        float,
5448        split,
5449        proto
5450    );
5451
5452    // ─── End-to-end tracked TensorDyn test macro ────────────────────
5453    //
5454    // Same as e2e_tracked_test but wraps data in TensorDyn and exercises
5455    // the public decode_tracked / decode_proto_tracked API.
5456
5457    macro_rules! e2e_tracked_tensor_test {
5458        ($name:ident, quantized, $layout:ident, $output:ident) => {
5459            #[test]
5460            fn $name() {
5461                use edgefirst_tensor::{Tensor, TensorMapTrait, TensorTrait};
5462
5463                let is_split = matches!(stringify!($layout), "split");
5464                let is_proto = matches!(stringify!($output), "proto");
5465
5466                let score_threshold = 0.45;
5467                let iou_threshold = 0.45;
5468
5469                let mut boxes = Array2::zeros((10, 4));
5470                let mut scores = Array2::zeros((10, 1));
5471                let mut classes = Array2::zeros((10, 1));
5472                let mask = Array2::zeros((10, 32));
5473                let protos_f64 = Array3::<f64>::zeros((160, 160, 32));
5474                let protos_f64 = protos_f64.insert_axis(Axis(0));
5475                let protos_quant = (1.0 / 255.0, 0.0);
5476                let protos_u8: Array4<u8> =
5477                    quantize_ndarray(protos_f64.view(), protos_quant.into());
5478
5479                boxes
5480                    .slice_mut(s![0, ..])
5481                    .assign(&array![0.1234, 0.1234, 0.2345, 0.2345]);
5482                scores.slice_mut(s![0, ..]).assign(&array![0.9876]);
5483                classes.slice_mut(s![0, ..]).assign(&array![2.0]);
5484
5485                let detect_quant = (2.0 / 255.0, 0.0);
5486
5487                let decoder = if is_split {
5488                    DecoderBuilder::default()
5489                        .with_config_yaml_str(E2E_SPLIT_CONFIG.to_string())
5490                        .with_score_threshold(score_threshold)
5491                        .with_iou_threshold(iou_threshold)
5492                        .build()
5493                        .unwrap()
5494                } else {
5495                    DecoderBuilder::default()
5496                        .with_config_yaml_str(E2E_COMBINED_CONFIG.to_string())
5497                        .with_score_threshold(score_threshold)
5498                        .with_iou_threshold(iou_threshold)
5499                        .build()
5500                        .unwrap()
5501                };
5502
5503                // Helper to wrap a u8 slice into a TensorDyn
5504                let make_u8_tensor =
5505                    |shape: &[usize], data: &[u8]| -> edgefirst_tensor::TensorDyn {
5506                        let t = Tensor::<u8>::new(shape, None, None).unwrap();
5507                        t.map().unwrap().as_mut_slice()[..data.len()].copy_from_slice(data);
5508                        t.into()
5509                    };
5510
5511                let expected = e2e_expected_boxes_quant();
5512                let mut tracker = ByteTrackBuilder::new()
5513                    .track_update(0.1)
5514                    .track_high_conf(0.7)
5515                    .build();
5516                let mut output_boxes = Vec::with_capacity(50);
5517                let mut output_tracks = Vec::with_capacity(50);
5518
5519                let protos_td = make_u8_tensor(protos_u8.shape(), protos_u8.as_slice().unwrap());
5520
5521                if is_split {
5522                    let boxes = boxes.insert_axis(Axis(0));
5523                    let scores = scores.insert_axis(Axis(0));
5524                    let classes = classes.insert_axis(Axis(0));
5525                    let mask = mask.insert_axis(Axis(0));
5526
5527                    let boxes_q: Array3<u8> = quantize_ndarray(boxes.view(), detect_quant.into());
5528                    let mut scores_q: Array3<u8> =
5529                        quantize_ndarray(scores.view(), detect_quant.into());
5530                    let classes_q: Array3<u8> =
5531                        quantize_ndarray(classes.view(), detect_quant.into());
5532                    let mask_q: Array3<u8> = quantize_ndarray(mask.view(), detect_quant.into());
5533
5534                    let boxes_td = make_u8_tensor(boxes_q.shape(), boxes_q.as_slice().unwrap());
5535                    let classes_td =
5536                        make_u8_tensor(classes_q.shape(), classes_q.as_slice().unwrap());
5537                    let mask_td = make_u8_tensor(mask_q.shape(), mask_q.as_slice().unwrap());
5538
5539                    if is_proto {
5540                        let scores_td =
5541                            make_u8_tensor(scores_q.shape(), scores_q.as_slice().unwrap());
5542                        decoder
5543                            .decode_proto_tracked(
5544                                &mut tracker,
5545                                0,
5546                                &[&boxes_td, &scores_td, &classes_td, &mask_td, &protos_td],
5547                                &mut output_boxes,
5548                                &mut output_tracks,
5549                            )
5550                            .unwrap();
5551
5552                        assert_eq!(output_boxes.len(), 1);
5553                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5554
5555                        for score in scores_q.slice_mut(s![.., .., ..]).iter_mut() {
5556                            *score = u8::MIN;
5557                        }
5558                        let scores_td =
5559                            make_u8_tensor(scores_q.shape(), scores_q.as_slice().unwrap());
5560                        let proto_result = decoder
5561                            .decode_proto_tracked(
5562                                &mut tracker,
5563                                100_000_000 / 3,
5564                                &[&boxes_td, &scores_td, &classes_td, &mask_td, &protos_td],
5565                                &mut output_boxes,
5566                                &mut output_tracks,
5567                            )
5568                            .unwrap();
5569                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5570                        assert!(proto_result.is_some_and(|x| x.mask_coefficients.shape()[0] == 0));
5571                    } else {
5572                        let scores_td =
5573                            make_u8_tensor(scores_q.shape(), scores_q.as_slice().unwrap());
5574                        let mut output_masks = Vec::with_capacity(50);
5575                        decoder
5576                            .decode_tracked(
5577                                &mut tracker,
5578                                0,
5579                                &[&boxes_td, &scores_td, &classes_td, &mask_td, &protos_td],
5580                                &mut output_boxes,
5581                                &mut output_masks,
5582                                &mut output_tracks,
5583                            )
5584                            .unwrap();
5585
5586                        assert_eq!(output_boxes.len(), 1);
5587                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5588
5589                        for score in scores_q.slice_mut(s![.., .., ..]).iter_mut() {
5590                            *score = u8::MIN;
5591                        }
5592                        let scores_td =
5593                            make_u8_tensor(scores_q.shape(), scores_q.as_slice().unwrap());
5594                        decoder
5595                            .decode_tracked(
5596                                &mut tracker,
5597                                100_000_000 / 3,
5598                                &[&boxes_td, &scores_td, &classes_td, &mask_td, &protos_td],
5599                                &mut output_boxes,
5600                                &mut output_masks,
5601                                &mut output_tracks,
5602                            )
5603                            .unwrap();
5604                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5605                        assert!(output_masks.is_empty());
5606                    }
5607                } else {
5608                    // Combined layout
5609                    let detect = ndarray::concatenate![
5610                        Axis(1),
5611                        boxes.view(),
5612                        scores.view(),
5613                        classes.view(),
5614                        mask.view()
5615                    ];
5616                    let detect = detect.insert_axis(Axis(0));
5617                    assert_eq!(detect.shape(), &[1, 10, 38]);
5618                    // Ensure contiguous layout after concatenation for as_slice()
5619                    let detect =
5620                        Array3::from_shape_vec(detect.raw_dim(), detect.iter().copied().collect())
5621                            .unwrap();
5622                    let mut detect_q: Array3<u8> =
5623                        quantize_ndarray(detect.view(), detect_quant.into());
5624
5625                    if is_proto {
5626                        let detect_td =
5627                            make_u8_tensor(detect_q.shape(), detect_q.as_slice().unwrap());
5628                        decoder
5629                            .decode_proto_tracked(
5630                                &mut tracker,
5631                                0,
5632                                &[&detect_td, &protos_td],
5633                                &mut output_boxes,
5634                                &mut output_tracks,
5635                            )
5636                            .unwrap();
5637
5638                        assert_eq!(output_boxes.len(), 1);
5639                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5640
5641                        for score in detect_q.slice_mut(s![.., .., 4]).iter_mut() {
5642                            *score = u8::MIN;
5643                        }
5644                        let detect_td =
5645                            make_u8_tensor(detect_q.shape(), detect_q.as_slice().unwrap());
5646                        let proto_result = decoder
5647                            .decode_proto_tracked(
5648                                &mut tracker,
5649                                100_000_000 / 3,
5650                                &[&detect_td, &protos_td],
5651                                &mut output_boxes,
5652                                &mut output_tracks,
5653                            )
5654                            .unwrap();
5655                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5656                        assert!(proto_result.is_some_and(|x| x.mask_coefficients.shape()[0] == 0));
5657                    } else {
5658                        let detect_td =
5659                            make_u8_tensor(detect_q.shape(), detect_q.as_slice().unwrap());
5660                        let mut output_masks = Vec::with_capacity(50);
5661                        decoder
5662                            .decode_tracked(
5663                                &mut tracker,
5664                                0,
5665                                &[&detect_td, &protos_td],
5666                                &mut output_boxes,
5667                                &mut output_masks,
5668                                &mut output_tracks,
5669                            )
5670                            .unwrap();
5671
5672                        assert_eq!(output_boxes.len(), 1);
5673                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5674
5675                        for score in detect_q.slice_mut(s![.., .., 4]).iter_mut() {
5676                            *score = u8::MIN;
5677                        }
5678                        let detect_td =
5679                            make_u8_tensor(detect_q.shape(), detect_q.as_slice().unwrap());
5680                        decoder
5681                            .decode_tracked(
5682                                &mut tracker,
5683                                100_000_000 / 3,
5684                                &[&detect_td, &protos_td],
5685                                &mut output_boxes,
5686                                &mut output_masks,
5687                                &mut output_tracks,
5688                            )
5689                            .unwrap();
5690                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5691                        assert!(output_masks.is_empty());
5692                    }
5693                }
5694            }
5695        };
5696        ($name:ident, float, $layout:ident, $output:ident) => {
5697            #[test]
5698            fn $name() {
5699                use edgefirst_tensor::{Tensor, TensorMapTrait, TensorTrait};
5700
5701                let is_split = matches!(stringify!($layout), "split");
5702                let is_proto = matches!(stringify!($output), "proto");
5703
5704                let score_threshold = 0.45;
5705                let iou_threshold = 0.45;
5706
5707                let mut boxes = Array2::zeros((10, 4));
5708                let mut scores = Array2::zeros((10, 1));
5709                let mut classes = Array2::zeros((10, 1));
5710                let mask: Array2<f64> = Array2::zeros((10, 32));
5711                let protos = Array3::<f64>::zeros((160, 160, 32));
5712                let protos = protos.insert_axis(Axis(0));
5713
5714                boxes
5715                    .slice_mut(s![0, ..])
5716                    .assign(&array![0.1234, 0.1234, 0.2345, 0.2345]);
5717                scores.slice_mut(s![0, ..]).assign(&array![0.9876]);
5718                classes.slice_mut(s![0, ..]).assign(&array![2.0]);
5719
5720                let decoder = if is_split {
5721                    DecoderBuilder::default()
5722                        .with_config_yaml_str(E2E_SPLIT_CONFIG.to_string())
5723                        .with_score_threshold(score_threshold)
5724                        .with_iou_threshold(iou_threshold)
5725                        .build()
5726                        .unwrap()
5727                } else {
5728                    DecoderBuilder::default()
5729                        .with_config_yaml_str(E2E_COMBINED_CONFIG.to_string())
5730                        .with_score_threshold(score_threshold)
5731                        .with_iou_threshold(iou_threshold)
5732                        .build()
5733                        .unwrap()
5734                };
5735
5736                // Helper to wrap an f64 slice into a TensorDyn
5737                let make_f64_tensor =
5738                    |shape: &[usize], data: &[f64]| -> edgefirst_tensor::TensorDyn {
5739                        let t = Tensor::<f64>::new(shape, None, None).unwrap();
5740                        t.map().unwrap().as_mut_slice()[..data.len()].copy_from_slice(data);
5741                        t.into()
5742                    };
5743
5744                let expected = e2e_expected_boxes_float();
5745                let mut tracker = ByteTrackBuilder::new()
5746                    .track_update(0.1)
5747                    .track_high_conf(0.7)
5748                    .build();
5749                let mut output_boxes = Vec::with_capacity(50);
5750                let mut output_tracks = Vec::with_capacity(50);
5751
5752                let protos_td = make_f64_tensor(protos.shape(), protos.as_slice().unwrap());
5753
5754                if is_split {
5755                    let boxes = boxes.insert_axis(Axis(0));
5756                    let mut scores = scores.insert_axis(Axis(0));
5757                    let classes = classes.insert_axis(Axis(0));
5758                    let mask = mask.insert_axis(Axis(0));
5759
5760                    let boxes_td = make_f64_tensor(boxes.shape(), boxes.as_slice().unwrap());
5761                    let classes_td = make_f64_tensor(classes.shape(), classes.as_slice().unwrap());
5762                    let mask_td = make_f64_tensor(mask.shape(), mask.as_slice().unwrap());
5763
5764                    if is_proto {
5765                        let scores_td = make_f64_tensor(scores.shape(), scores.as_slice().unwrap());
5766                        decoder
5767                            .decode_proto_tracked(
5768                                &mut tracker,
5769                                0,
5770                                &[&boxes_td, &scores_td, &classes_td, &mask_td, &protos_td],
5771                                &mut output_boxes,
5772                                &mut output_tracks,
5773                            )
5774                            .unwrap();
5775
5776                        assert_eq!(output_boxes.len(), 1);
5777                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5778
5779                        for score in scores.slice_mut(s![.., .., ..]).iter_mut() {
5780                            *score = 0.0;
5781                        }
5782                        let scores_td = make_f64_tensor(scores.shape(), scores.as_slice().unwrap());
5783                        let proto_result = decoder
5784                            .decode_proto_tracked(
5785                                &mut tracker,
5786                                100_000_000 / 3,
5787                                &[&boxes_td, &scores_td, &classes_td, &mask_td, &protos_td],
5788                                &mut output_boxes,
5789                                &mut output_tracks,
5790                            )
5791                            .unwrap();
5792                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5793                        assert!(proto_result.is_some_and(|x| x.mask_coefficients.shape()[0] == 0));
5794                    } else {
5795                        let scores_td = make_f64_tensor(scores.shape(), scores.as_slice().unwrap());
5796                        let mut output_masks = Vec::with_capacity(50);
5797                        decoder
5798                            .decode_tracked(
5799                                &mut tracker,
5800                                0,
5801                                &[&boxes_td, &scores_td, &classes_td, &mask_td, &protos_td],
5802                                &mut output_boxes,
5803                                &mut output_masks,
5804                                &mut output_tracks,
5805                            )
5806                            .unwrap();
5807
5808                        assert_eq!(output_boxes.len(), 1);
5809                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5810
5811                        for score in scores.slice_mut(s![.., .., ..]).iter_mut() {
5812                            *score = 0.0;
5813                        }
5814                        let scores_td = make_f64_tensor(scores.shape(), scores.as_slice().unwrap());
5815                        decoder
5816                            .decode_tracked(
5817                                &mut tracker,
5818                                100_000_000 / 3,
5819                                &[&boxes_td, &scores_td, &classes_td, &mask_td, &protos_td],
5820                                &mut output_boxes,
5821                                &mut output_masks,
5822                                &mut output_tracks,
5823                            )
5824                            .unwrap();
5825                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5826                        assert!(output_masks.is_empty());
5827                    }
5828                } else {
5829                    // Combined layout
5830                    let detect = ndarray::concatenate![
5831                        Axis(1),
5832                        boxes.view(),
5833                        scores.view(),
5834                        classes.view(),
5835                        mask.view()
5836                    ];
5837                    let detect = detect.insert_axis(Axis(0));
5838                    assert_eq!(detect.shape(), &[1, 10, 38]);
5839                    // Ensure contiguous layout after concatenation for as_slice()
5840                    let mut detect =
5841                        Array3::from_shape_vec(detect.raw_dim(), detect.iter().copied().collect())
5842                            .unwrap();
5843
5844                    if is_proto {
5845                        let detect_td = make_f64_tensor(detect.shape(), detect.as_slice().unwrap());
5846                        decoder
5847                            .decode_proto_tracked(
5848                                &mut tracker,
5849                                0,
5850                                &[&detect_td, &protos_td],
5851                                &mut output_boxes,
5852                                &mut output_tracks,
5853                            )
5854                            .unwrap();
5855
5856                        assert_eq!(output_boxes.len(), 1);
5857                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5858
5859                        for score in detect.slice_mut(s![.., .., 4]).iter_mut() {
5860                            *score = 0.0;
5861                        }
5862                        let detect_td = make_f64_tensor(detect.shape(), detect.as_slice().unwrap());
5863                        let proto_result = decoder
5864                            .decode_proto_tracked(
5865                                &mut tracker,
5866                                100_000_000 / 3,
5867                                &[&detect_td, &protos_td],
5868                                &mut output_boxes,
5869                                &mut output_tracks,
5870                            )
5871                            .unwrap();
5872                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5873                        assert!(proto_result.is_some_and(|x| x.mask_coefficients.shape()[0] == 0));
5874                    } else {
5875                        let detect_td = make_f64_tensor(detect.shape(), detect.as_slice().unwrap());
5876                        let mut output_masks = Vec::with_capacity(50);
5877                        decoder
5878                            .decode_tracked(
5879                                &mut tracker,
5880                                0,
5881                                &[&detect_td, &protos_td],
5882                                &mut output_boxes,
5883                                &mut output_masks,
5884                                &mut output_tracks,
5885                            )
5886                            .unwrap();
5887
5888                        assert_eq!(output_boxes.len(), 1);
5889                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1.0 / 160.0));
5890
5891                        for score in detect.slice_mut(s![.., .., 4]).iter_mut() {
5892                            *score = 0.0;
5893                        }
5894                        let detect_td = make_f64_tensor(detect.shape(), detect.as_slice().unwrap());
5895                        decoder
5896                            .decode_tracked(
5897                                &mut tracker,
5898                                100_000_000 / 3,
5899                                &[&detect_td, &protos_td],
5900                                &mut output_boxes,
5901                                &mut output_masks,
5902                                &mut output_tracks,
5903                            )
5904                            .unwrap();
5905                        assert!(output_boxes[0].equal_within_delta(&expected[0], 1e-6));
5906                        assert!(output_masks.is_empty());
5907                    }
5908                }
5909            }
5910        };
5911    }
5912
5913    e2e_tracked_tensor_test!(
5914        test_decoder_tracked_tensor_end_to_end_segdet,
5915        quantized,
5916        combined,
5917        masks
5918    );
5919    e2e_tracked_tensor_test!(
5920        test_decoder_tracked_tensor_end_to_end_segdet_float,
5921        float,
5922        combined,
5923        masks
5924    );
5925    e2e_tracked_tensor_test!(
5926        test_decoder_tracked_tensor_end_to_end_segdet_proto,
5927        quantized,
5928        combined,
5929        proto
5930    );
5931    e2e_tracked_tensor_test!(
5932        test_decoder_tracked_tensor_end_to_end_segdet_proto_float,
5933        float,
5934        combined,
5935        proto
5936    );
5937    e2e_tracked_tensor_test!(
5938        test_decoder_tracked_tensor_end_to_end_segdet_split,
5939        quantized,
5940        split,
5941        masks
5942    );
5943    e2e_tracked_tensor_test!(
5944        test_decoder_tracked_tensor_end_to_end_segdet_split_float,
5945        float,
5946        split,
5947        masks
5948    );
5949    e2e_tracked_tensor_test!(
5950        test_decoder_tracked_tensor_end_to_end_segdet_split_proto,
5951        quantized,
5952        split,
5953        proto
5954    );
5955    e2e_tracked_tensor_test!(
5956        test_decoder_tracked_tensor_end_to_end_segdet_split_proto_float,
5957        float,
5958        split,
5959        proto
5960    );
5961
5962    #[test]
5963    fn test_decoder_tracked_linear_motion() {
5964        use crate::configs::{DecoderType, Nms};
5965        use crate::DecoderBuilder;
5966
5967        let score_threshold = 0.25;
5968        let iou_threshold = 0.1;
5969        let out = edgefirst_bench::testdata::read("yolov8s_80_classes.bin");
5970        let out = unsafe { std::slice::from_raw_parts(out.as_ptr() as *const i8, out.len()) };
5971        let mut out = Array3::from_shape_vec((1, 84, 8400), out.to_vec()).unwrap();
5972        let quant = (0.0040811873, -123).into();
5973
5974        let decoder = DecoderBuilder::default()
5975            .with_config_yolo_det(
5976                crate::configs::Detection {
5977                    decoder: DecoderType::Ultralytics,
5978                    shape: vec![1, 84, 8400],
5979                    anchors: None,
5980                    quantization: Some(quant),
5981                    dshape: vec![
5982                        (crate::configs::DimName::Batch, 1),
5983                        (crate::configs::DimName::NumFeatures, 84),
5984                        (crate::configs::DimName::NumBoxes, 8400),
5985                    ],
5986                    normalized: Some(true),
5987                },
5988                None,
5989            )
5990            .with_score_threshold(score_threshold)
5991            .with_iou_threshold(iou_threshold)
5992            .with_nms(Some(Nms::ClassAgnostic))
5993            .build()
5994            .unwrap();
5995
5996        let mut expected_boxes = [
5997            DetectBox {
5998                bbox: BoundingBox {
5999                    xmin: 0.5285137,
6000                    ymin: 0.05305544,
6001                    xmax: 0.87541467,
6002                    ymax: 0.9998909,
6003                },
6004                score: 0.5591227,
6005                label: 0,
6006            },
6007            DetectBox {
6008                bbox: BoundingBox {
6009                    xmin: 0.130598,
6010                    ymin: 0.43260583,
6011                    xmax: 0.35098213,
6012                    ymax: 0.9958097,
6013                },
6014                score: 0.33057618,
6015                label: 75,
6016            },
6017        ];
6018
6019        let mut tracker = ByteTrackBuilder::new()
6020            .track_update(0.1)
6021            .track_high_conf(0.3)
6022            .build();
6023
6024        let mut output_boxes = Vec::with_capacity(50);
6025        let mut output_masks = Vec::with_capacity(50);
6026        let mut output_tracks = Vec::with_capacity(50);
6027
6028        decoder
6029            .decode_tracked_quantized(
6030                &mut tracker,
6031                0,
6032                &[out.view().into()],
6033                &mut output_boxes,
6034                &mut output_masks,
6035                &mut output_tracks,
6036            )
6037            .unwrap();
6038
6039        assert_eq!(output_boxes.len(), 2);
6040        assert!(output_boxes[0].equal_within_delta(&expected_boxes[0], 1e-6));
6041        assert!(output_boxes[1].equal_within_delta(&expected_boxes[1], 1e-6));
6042
6043        for i in 1..=100 {
6044            let mut out = out.clone();
6045            // introduce linear movement into the XY coordinates
6046            let mut x_values = out.slice_mut(s![0, 0, ..]);
6047            for x in x_values.iter_mut() {
6048                *x = x.saturating_add((i as f32 * 1e-3 / quant.0).round() as i8);
6049            }
6050
6051            decoder
6052                .decode_tracked_quantized(
6053                    &mut tracker,
6054                    100_000_000 * i / 3, // simulate 33.333ms between frames
6055                    &[out.view().into()],
6056                    &mut output_boxes,
6057                    &mut output_masks,
6058                    &mut output_tracks,
6059                )
6060                .unwrap();
6061
6062            assert_eq!(output_boxes.len(), 2);
6063        }
6064        let tracks = tracker.get_active_tracks();
6065        let predicted_boxes: Vec<_> = tracks
6066            .iter()
6067            .map(|track| {
6068                let mut l = track.last_box;
6069                l.bbox = track.info.tracked_location.into();
6070                l
6071            })
6072            .collect();
6073        expected_boxes[0].bbox.xmin += 0.1; // compensate for linear movement
6074        expected_boxes[0].bbox.xmax += 0.1;
6075        expected_boxes[1].bbox.xmin += 0.1;
6076        expected_boxes[1].bbox.xmax += 0.1;
6077
6078        assert!(predicted_boxes[0].equal_within_delta(&expected_boxes[0], 1e-3));
6079        assert!(predicted_boxes[1].equal_within_delta(&expected_boxes[1], 1e-3));
6080
6081        // give the decoder a final frame with no detections to ensure tracks are properly predicting forward when detection is missing
6082        let mut scores_values = out.slice_mut(s![0, 4.., ..]);
6083        for score in scores_values.iter_mut() {
6084            *score = i8::MIN; // set all scores to minimum to simulate no detections
6085        }
6086        decoder
6087            .decode_tracked_quantized(
6088                &mut tracker,
6089                100_000_000 * 101 / 3,
6090                &[out.view().into()],
6091                &mut output_boxes,
6092                &mut output_masks,
6093                &mut output_tracks,
6094            )
6095            .unwrap();
6096        expected_boxes[0].bbox.xmin += 0.001; // compensate for expected movement
6097        expected_boxes[0].bbox.xmax += 0.001;
6098        expected_boxes[1].bbox.xmin += 0.001;
6099        expected_boxes[1].bbox.xmax += 0.001;
6100
6101        assert!(output_boxes[0].equal_within_delta(&expected_boxes[0], 1e-3));
6102        assert!(output_boxes[1].equal_within_delta(&expected_boxes[1], 1e-3));
6103    }
6104
6105    #[test]
6106    fn test_decoder_tracked_end_to_end_float() {
6107        let score_threshold = 0.45;
6108        let iou_threshold = 0.45;
6109
6110        let mut boxes = Array2::zeros((10, 4));
6111        let mut scores = Array2::zeros((10, 1));
6112        let mut classes = Array2::zeros((10, 1));
6113
6114        boxes
6115            .slice_mut(s![0, ..,])
6116            .assign(&array![0.1234, 0.1234, 0.2345, 0.2345]);
6117        scores.slice_mut(s![0, ..]).assign(&array![0.9876]);
6118        classes.slice_mut(s![0, ..]).assign(&array![2.0]);
6119
6120        let detect = ndarray::concatenate![Axis(1), boxes.view(), scores.view(), classes.view(),];
6121        let mut detect = detect.insert_axis(Axis(0));
6122        assert_eq!(detect.shape(), &[1, 10, 6]);
6123        let config = "
6124decoder_version: yolo26
6125outputs:
6126 - type: detection
6127   decoder: ultralytics
6128   quantization: [0.00784313725490196, 0]
6129   shape: [1, 10, 6]
6130   dshape:
6131    - [batch, 1]
6132    - [num_boxes, 10]
6133    - [num_features, 6]
6134   normalized: true
6135";
6136
6137        let decoder = DecoderBuilder::default()
6138            .with_config_yaml_str(config.to_string())
6139            .with_score_threshold(score_threshold)
6140            .with_iou_threshold(iou_threshold)
6141            .build()
6142            .unwrap();
6143
6144        let expected_boxes = [DetectBox {
6145            bbox: BoundingBox {
6146                xmin: 0.1234,
6147                ymin: 0.1234,
6148                xmax: 0.2345,
6149                ymax: 0.2345,
6150            },
6151            score: 0.9876,
6152            label: 2,
6153        }];
6154
6155        let mut tracker = ByteTrackBuilder::new()
6156            .track_update(0.1)
6157            .track_high_conf(0.7)
6158            .build();
6159
6160        let mut output_boxes = Vec::with_capacity(50);
6161        let mut output_masks = Vec::with_capacity(50);
6162        let mut output_tracks = Vec::with_capacity(50);
6163
6164        decoder
6165            .decode_tracked_float(
6166                &mut tracker,
6167                0,
6168                &[detect.view().into_dyn()],
6169                &mut output_boxes,
6170                &mut output_masks,
6171                &mut output_tracks,
6172            )
6173            .unwrap();
6174
6175        assert_eq!(output_boxes.len(), 1);
6176        assert!(output_boxes[0].equal_within_delta(&expected_boxes[0], 1e-6));
6177
6178        // give the decoder a final frame with no detections to ensure tracks are properly predicting forward when detection is missing
6179
6180        for score in detect.slice_mut(s![.., .., 4]).iter_mut() {
6181            *score = 0.0; // set all scores to minimum to simulate no detections
6182        }
6183
6184        decoder
6185            .decode_tracked_float(
6186                &mut tracker,
6187                100_000_000 / 3,
6188                &[detect.view().into_dyn()],
6189                &mut output_boxes,
6190                &mut output_masks,
6191                &mut output_tracks,
6192            )
6193            .unwrap();
6194        assert!(output_boxes[0].equal_within_delta(&expected_boxes[0], 1e-6));
6195    }
6196}