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