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