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