Skip to main content

edgefirst_decoder/decoder/
mod.rs

1// SPDX-FileCopyrightText: Copyright 2025 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4use ndarray::{ArrayView, ArrayViewD, Dimension};
5use num_traits::{AsPrimitive, Float};
6
7use crate::{DecoderError, DetectBox, ProtoData, Segmentation};
8
9pub mod config;
10pub mod configs;
11
12use configs::ModelType;
13
14#[derive(Debug)]
15pub struct Decoder {
16    model_type: ModelType,
17    pub iou_threshold: f32,
18    pub score_threshold: f32,
19    /// NMS mode (always a concrete variant after build — `Nms::Auto` is
20    /// resolved during `DecoderBuilder::build()` and never stored here):
21    /// - `Some(ClassAgnostic)` — class-agnostic NMS
22    /// - `Some(ClassAware)` — class-aware NMS
23    /// - `None` — NMS bypassed (end-to-end models)
24    pub nms: Option<configs::Nms>,
25    /// Maximum number of candidate boxes fed into NMS after score filtering.
26    /// Reduces O(N²) NMS cost when many low-confidence proposals pass the
27    /// threshold (common during COCO mAP evaluation with threshold ≈ 0.001).
28    /// Candidates are ranked by score; only the top `pre_nms_top_k` proceed
29    /// to NMS.  Default: 300.  Ignored when `nms` is `None`.
30    ///
31    /// # ⚠️ Validation vs Deployment
32    ///
33    /// The default of 300 is tuned for **deployment** (score threshold ≥ 0.25)
34    /// where few anchors pass the score filter, making top-K a no-op in
35    /// practice while bounding worst-case NMS cost.
36    ///
37    /// For **mAP evaluation** (score threshold ≈ 0.001), most of the 8 400
38    /// YOLO anchors pass the score filter. At `pre_nms_top_k = 300`, roughly
39    /// 74 % of candidates that would survive NMS are discarded *before* NMS
40    /// runs, causing **~9 pp box mAP loss** — a measurement artifact, not a
41    /// model quality issue.
42    ///
43    /// | Use case | `pre_nms_top_k` | `score_threshold` |
44    /// |----------|----------------:|------------------:|
45    /// | Deployment | 300 (default) | ≥ 0.25 |
46    /// | COCO mAP evaluation | 8 400 (all anchors) | 0.001 |
47    /// | Unbounded | 0 (no limit) | any |
48    ///
49    /// Post-processing latency scales with the number of candidates entering
50    /// NMS. At deployment thresholds the candidate count is already small, so
51    /// raising `pre_nms_top_k` has negligible cost. At validation thresholds
52    /// the increase is measurable but necessary for correct recall.
53    pub pre_nms_top_k: usize,
54    /// Maximum number of detections returned after NMS. Matches the
55    /// Ultralytics `max_det` parameter.  Default: 300.
56    ///
57    /// This bound applies uniformly across all segmentation and detection
58    /// decode paths reached via [`Decoder::decode`] / [`Decoder::decode_proto`].
59    /// The output `Vec`'s capacity is only an allocation hint; the post-NMS
60    /// detection count is bounded solely by `max_det` (EDGEAI-1302).
61    pub max_det: usize,
62    /// Whether decoded boxes are in normalized [0,1] coordinates.
63    /// - `Some(true)`: Coordinates in [0,1] range
64    /// - `Some(false)`: Pixel coordinates
65    /// - `None`: Unknown, caller must infer (e.g., check if any coordinate >
66    ///   1.0)
67    normalized: Option<bool>,
68    /// Model input spatial dimensions `(width, height)`, captured from
69    /// the schema's `input.shape` / `input.dshape` at builder time.
70    /// Required to honour `normalized: false`: pixel-space box coords
71    /// emitted by the model are divided by these dimensions before NMS
72    /// so the post-NMS bbox is in `[0, 1]`. `None` when no schema input
73    /// spec is available — the legacy >2.0 reject in `protobox` then
74    /// preserves the previous safety net (EDGEAI-1303).
75    input_dims: Option<(usize, usize)>,
76    /// Schema v2 merge program. Present when the decoder was built from
77    /// a [`crate::schema::SchemaV2`] whose logical outputs carry
78    /// physical children. Absent for flat configurations (v1 and
79    /// flat-v2).
80    pub(crate) decode_program: Option<merge::DecodeProgram>,
81    /// When `true`, emit one detection per (anchor, class) pair for every
82    /// class whose score meets the threshold — matching Ultralytics `val`
83    /// multi-label decode.  **OFF by default.**
84    ///
85    /// # Deployment safety
86    /// Multi-label must never be enabled on the tracked/deployment path:
87    /// the ByteTrack IoU-only tracker spawns one track per detection, so
88    /// duplicate boxes for the same anchor at different labels produce
89    /// phantom tracks.  `decode_tracked_*` entry points assert this is off.
90    ///
91    /// This field is intentionally builder-only (not schema/config-driven)
92    /// so a deployed `edgefirst.json` can never accidentally enable it.
93    pub(crate) multi_label: bool,
94    /// Per-scale fast path. Constructed at build time from a schema-v2
95    /// document with per-scale children. Wrapped in `Mutex` because
96    /// `Decoder::decode_proto` and `Decoder::decode` are `&self` but
97    /// the per-scale buffers are mutated per-frame.
98    pub(crate) per_scale: Option<std::sync::Mutex<crate::per_scale::PerScaleDecoder>>,
99}
100
101impl PartialEq for Decoder {
102    fn eq(&self, other: &Self) -> bool {
103        // DecodeProgram and PerScaleDecoder have non-comparable embedded
104        // data; compare by the config-derived fields only.
105        self.model_type == other.model_type
106            && self.iou_threshold == other.iou_threshold
107            && self.score_threshold == other.score_threshold
108            && self.nms == other.nms
109            && self.pre_nms_top_k == other.pre_nms_top_k
110            && self.max_det == other.max_det
111            && self.normalized == other.normalized
112            && self.input_dims == other.input_dims
113            && self.multi_label == other.multi_label
114            && self.decode_program.is_some() == other.decode_program.is_some()
115            && self.per_scale.is_some() == other.per_scale.is_some()
116    }
117}
118
119impl Clone for Decoder {
120    /// Cloning a `Decoder` preserves the legacy decode path
121    /// (`decode_program`) but drops the per-scale fast path:
122    /// `PerScaleDecoder` owns mutable per-frame scratch buffers and is
123    /// not `Clone`. Decoders built from a per-scale schema should be
124    /// rebuilt via [`DecoderBuilder`] rather than cloned to preserve the
125    /// fast path; cloning is intended for tests and rare configs.
126    fn clone(&self) -> Self {
127        Self {
128            model_type: self.model_type.clone(),
129            iou_threshold: self.iou_threshold,
130            score_threshold: self.score_threshold,
131            nms: self.nms,
132            pre_nms_top_k: self.pre_nms_top_k,
133            max_det: self.max_det,
134            normalized: self.normalized,
135            input_dims: self.input_dims,
136            multi_label: self.multi_label,
137            decode_program: self.decode_program.clone(),
138            per_scale: None,
139        }
140    }
141}
142
143#[derive(Debug)]
144pub(crate) enum ArrayViewDQuantized<'a> {
145    UInt8(ArrayViewD<'a, u8>),
146    Int8(ArrayViewD<'a, i8>),
147    UInt16(ArrayViewD<'a, u16>),
148    Int16(ArrayViewD<'a, i16>),
149    UInt32(ArrayViewD<'a, u32>),
150    Int32(ArrayViewD<'a, i32>),
151}
152
153impl<'a, D> From<ArrayView<'a, u8, D>> for ArrayViewDQuantized<'a>
154where
155    D: Dimension,
156{
157    fn from(arr: ArrayView<'a, u8, D>) -> Self {
158        Self::UInt8(arr.into_dyn())
159    }
160}
161
162impl<'a, D> From<ArrayView<'a, i8, D>> for ArrayViewDQuantized<'a>
163where
164    D: Dimension,
165{
166    fn from(arr: ArrayView<'a, i8, D>) -> Self {
167        Self::Int8(arr.into_dyn())
168    }
169}
170
171impl<'a, D> From<ArrayView<'a, u16, D>> for ArrayViewDQuantized<'a>
172where
173    D: Dimension,
174{
175    fn from(arr: ArrayView<'a, u16, D>) -> Self {
176        Self::UInt16(arr.into_dyn())
177    }
178}
179
180impl<'a, D> From<ArrayView<'a, i16, D>> for ArrayViewDQuantized<'a>
181where
182    D: Dimension,
183{
184    fn from(arr: ArrayView<'a, i16, D>) -> Self {
185        Self::Int16(arr.into_dyn())
186    }
187}
188
189impl<'a, D> From<ArrayView<'a, u32, D>> for ArrayViewDQuantized<'a>
190where
191    D: Dimension,
192{
193    fn from(arr: ArrayView<'a, u32, D>) -> Self {
194        Self::UInt32(arr.into_dyn())
195    }
196}
197
198impl<'a, D> From<ArrayView<'a, i32, D>> for ArrayViewDQuantized<'a>
199where
200    D: Dimension,
201{
202    fn from(arr: ArrayView<'a, i32, D>) -> Self {
203        Self::Int32(arr.into_dyn())
204    }
205}
206
207impl<'a> ArrayViewDQuantized<'a> {
208    /// Returns the shape of the underlying array.
209    pub(crate) fn shape(&self) -> &[usize] {
210        match self {
211            ArrayViewDQuantized::UInt8(a) => a.shape(),
212            ArrayViewDQuantized::Int8(a) => a.shape(),
213            ArrayViewDQuantized::UInt16(a) => a.shape(),
214            ArrayViewDQuantized::Int16(a) => a.shape(),
215            ArrayViewDQuantized::UInt32(a) => a.shape(),
216            ArrayViewDQuantized::Int32(a) => a.shape(),
217        }
218    }
219}
220
221/// WARNING: Do NOT nest `with_quantized!` calls. Each level multiplies
222/// monomorphized code paths by 6 (one per integer variant), so nesting
223/// N levels deep produces 6^N instantiations.
224///
225/// Instead, dequantize each tensor sequentially with `dequant_3d!`/`dequant_4d!`
226/// (6*N paths) or split into independent phases that each nest at most 2 levels.
227macro_rules! with_quantized {
228    ($x:expr, $var:ident, $body:expr) => {
229        match $x {
230            ArrayViewDQuantized::UInt8(x) => {
231                let $var = x;
232                $body
233            }
234            ArrayViewDQuantized::Int8(x) => {
235                let $var = x;
236                $body
237            }
238            ArrayViewDQuantized::UInt16(x) => {
239                let $var = x;
240                $body
241            }
242            ArrayViewDQuantized::Int16(x) => {
243                let $var = x;
244                $body
245            }
246            ArrayViewDQuantized::UInt32(x) => {
247                let $var = x;
248                $body
249            }
250            ArrayViewDQuantized::Int32(x) => {
251                let $var = x;
252                $body
253            }
254        }
255    };
256}
257
258mod builder;
259mod helpers;
260mod merge;
261mod per_scale_bridge;
262mod postprocess;
263mod tensor_bridge;
264mod tests;
265
266pub use builder::DecoderBuilder;
267pub use config::{ConfigOutput, ConfigOutputRef, ConfigOutputs};
268
269impl Decoder {
270    /// Static label identifying which dispatch path `decode` / `decode_proto`
271    /// will take, used as a tracing-span attribute. Lets profiling tools
272    /// distinguish `per_scale` (the fast path), `decode_program` (schema-v2
273    /// merge), and `legacy` (config-driven) without requiring callers to
274    /// inspect the model.
275    fn decode_path_label(&self) -> &'static str {
276        if self.per_scale.is_some() {
277            "per_scale"
278        } else if self.decode_program.is_some() {
279            "decode_program"
280        } else {
281            "legacy"
282        }
283    }
284
285    /// This function returns the parsed model type of the decoder.
286    ///
287    /// # Examples
288    ///
289    /// ```rust
290    /// # use edgefirst_decoder::{DecoderBuilder, DecoderResult, configs::ModelType};
291    /// # fn main() -> DecoderResult<()> {
292    /// #    let config_yaml = edgefirst_bench::testdata::read_to_string("modelpack_split.yaml").to_string();
293    ///     let decoder = DecoderBuilder::default()
294    ///         .with_config_yaml_str(config_yaml)
295    ///         .build()?;
296    ///     assert!(matches!(
297    ///         decoder.model_type(),
298    ///         ModelType::ModelPackDetSplit { .. }
299    ///     ));
300    /// #    Ok(())
301    /// # }
302    /// ```
303    pub fn model_type(&self) -> &ModelType {
304        &self.model_type
305    }
306
307    /// Returns the coordinate format of the boxes the decoder emits to
308    /// the caller.
309    ///
310    /// - `Some(true)`: Boxes are in normalized `[0, 1]` coordinates
311    /// - `Some(false)`: Boxes are in pixel coordinates relative to the
312    ///   model input
313    /// - `None`: Unknown, caller must infer (e.g., check if any coordinate
314    ///   > 1.0)
315    ///
316    /// This describes the **post-decode** coordinate space, not the raw
317    /// schema annotation. The decoder applies EDGEAI-1303 normalization
318    /// (dividing bbox channels by `(input_w, input_h)`) on a per-path
319    /// basis, not unconditionally. Four paths are known to invoke the
320    /// helper uniformly across all of their entry points (`decode`,
321    /// `decode_proto`, and — where applicable — `decode_tracked` and
322    /// `decode_tracked_proto`):
323    ///
324    /// 1. The **per-scale fast path** (DFL/LTRB → dist2bbox → sigmoid),
325    ///    which emits pixel-space boxes by design and always normalizes
326    ///    before returning.
327    /// 2. [`ModelType::YoloSegDet`](crate::ModelType::YoloSegDet), whose
328    ///    quantized and float, tracked and untracked, masks and proto
329    ///    variants each call the helper after NMS.
330    /// 3. [`ModelType::YoloSplitSegDet`](crate::ModelType::YoloSplitSegDet),
331    ///    aligned across `decode`, `decode_proto`, `decode_tracked`,
332    ///    and `decode_tracked_proto` for both quantized and float
333    ///    variants.
334    /// 4. [`ModelType::YoloSegDet2Way`](crate::ModelType::YoloSegDet2Way),
335    ///    aligned across the same four entry points and both element
336    ///    type variants.
337    ///
338    /// When any of those paths is active and the schema declares
339    /// `normalized: false` with valid [`input_dims`](Self::input_dims),
340    /// this accessor reports `Some(true)` to match what the caller
341    /// actually receives.
342    ///
343    /// The remaining model types still surface the raw schema flag
344    /// because their post-decode contract differs:
345    /// [`ModelType::YoloDet`](crate::ModelType::YoloDet) and
346    /// [`ModelType::YoloSplitDet`](crate::ModelType::YoloSplitDet)
347    /// (detection-only, no protobox crop coupling), the
348    /// `YoloEndToEnd*` family (model embeds its own NMS and emits its
349    /// own coordinate space), and the `ModelPack*` family (separate
350    /// model conventions). For those, this accessor returns
351    /// `self.normalized` verbatim and leaves it to the caller to
352    /// handle pixel-space output explicitly (e.g. divide by
353    /// `input_dims()` themselves).
354    ///
355    /// # Examples
356    ///
357    /// ```rust
358    /// # use edgefirst_decoder::{DecoderBuilder, DecoderResult};
359    /// # fn main() -> DecoderResult<()> {
360    /// #    let config_yaml = edgefirst_bench::testdata::read_to_string("modelpack_split.yaml").to_string();
361    ///     let decoder = DecoderBuilder::default()
362    ///         .with_config_yaml_str(config_yaml)
363    ///         .build()?;
364    ///     // Config doesn't specify normalized, so it's None
365    ///     assert!(decoder.normalized_boxes().is_none());
366    /// #    Ok(())
367    /// # }
368    /// ```
369    pub fn normalized_boxes(&self) -> Option<bool> {
370        // Four paths invoke `yolo::maybe_normalize_boxes_in_place`
371        // uniformly across every entry point that can reach them:
372        //   - the per-scale fast path (always normalizes by design),
373        //   - `ModelType::YoloSegDet` (helper fires in
374        //     `decode`/`decode_proto` via `yolo::impl_yolo_segdet_*` and
375        //     in `decode_tracked`/`decode_tracked_proto` via the
376        //     `process_tracked_yolo_segmentation!` macro and
377        //     `process_tracked_yolo_segdet_float`),
378        //   - `ModelType::YoloSplitSegDet` (helper fires in
379        //     `decode_yolo_split_segdet_*`, `impl_yolo_split_segdet_*`,
380        //     `process_tracked_yolo_segmentation_split!`, and
381        //     `process_tracked_yolo_segdet_split_float`), and
382        //   - `ModelType::YoloSegDet2Way` (helper fires in
383        //     `decode_yolo_segdet_2way_*`, the float decode routes
384        //     through `impl_yolo_split_segdet_float*`,
385        //     `process_tracked_yolo_segmentation_2way!`, and the
386        //     inline tracked-2way float helpers).
387        // For those, `normalized == Some(false)` with valid `input_dims`
388        // upgrades to a post-decode `Some(true)`. Other paths invoke
389        // the helper inconsistently across `ModelType` variants and
390        // tracked/proto entry points — surface the raw schema flag
391        // there and let callers handle pixel-space output explicitly.
392        if self.per_scale.is_some() || self.legacy_path_normalizes_uniformly() {
393            match (self.normalized, self.input_dims) {
394                (Some(true), _) => Some(true),
395                (Some(false), Some((w, h))) if w != 0 && h != 0 => Some(true),
396                (Some(false), _) => Some(false),
397                (None, _) => None,
398            }
399        } else {
400            self.normalized
401        }
402    }
403
404    /// Returns true for legacy `ModelType` dispatch paths that are known
405    /// to call `yolo::maybe_normalize_boxes_in_place` on every entry
406    /// point (`decode`, `decode_proto`, `decode_tracked`,
407    /// `decode_tracked_proto`, both quantized and float variants).
408    ///
409    /// Used by [`normalized_boxes`](Self::normalized_boxes) to gate the
410    /// pixel→normalized upgrade for non-per-scale model types whose
411    /// post-decode contract matches the per-scale path. Extend this
412    /// list as additional `ModelType` variants are brought into
413    /// uniform-normalization alignment.
414    fn legacy_path_normalizes_uniformly(&self) -> bool {
415        matches!(
416            self.model_type,
417            ModelType::YoloSegDet { .. }
418                | ModelType::YoloSplitSegDet { .. }
419                | ModelType::YoloSegDet2Way { .. }
420        )
421    }
422
423    /// Model input dimensions `(width, height)` captured from the
424    /// schema's `input.shape` / `input.dshape`, or `None` when the
425    /// schema did not declare an input spec (e.g. flat YAML configs
426    /// or `DecoderBuilder::add_output(...)` programmatic builds).
427    ///
428    /// Drives EDGEAI-1303 normalization on the paths that invoke the
429    /// helper uniformly: when the schema declares pixel-space outputs
430    /// and `input_dims()` is `Some((w, h))`, the per-scale bridge and
431    /// the `ModelType::YoloSegDet`, `ModelType::YoloSplitSegDet`, and
432    /// `ModelType::YoloSegDet2Way` dispatch paths divide post-NMS
433    /// bbox coordinates by `(w, h)` so they enter the canonical
434    /// `[0, 1]` range before mask cropping / tracker dispatch, and
435    /// [`normalized_boxes`](Self::normalized_boxes) reports
436    /// `Some(true)` to match. The remaining legacy `ModelType`
437    /// dispatch paths (detection-only `YoloDet`/`YoloSplitDet`, the
438    /// `YoloEndToEnd*` family, and the `ModelPack*` family) do not
439    /// apply this division — see
440    /// [`normalized_boxes`](Self::normalized_boxes) for the per-path
441    /// contract. The legacy `protobox` `> 2.0` reject acts as a safety
442    /// net for paths that emit pixel-space coordinates.
443    ///
444    /// # Examples
445    ///
446    /// ```rust
447    /// # use edgefirst_decoder::{schema::SchemaV2, DecoderBuilder, DecoderResult};
448    /// # fn main() -> DecoderResult<()> {
449    ///     let json = r#"{
450    ///         "schema_version": 2,
451    ///         "nms": "class_agnostic",
452    ///         "input": {
453    ///             "shape": [1, 640, 640, 3],
454    ///             "dshape": [{"batch": 1}, {"height": 640}, {"width": 640}, {"num_features": 3}]
455    ///         },
456    ///         "outputs": [{
457    ///             "name": "out", "type": "detection",
458    ///             "shape": [1, 38, 256],
459    ///             "dshape": [{"batch": 1}, {"num_features": 38}, {"num_boxes": 256}],
460    ///             "decoder": "ultralytics", "encoding": "direct", "normalized": false
461    ///         }]
462    ///     }"#;
463    ///     let schema: SchemaV2 = serde_json::from_str(json).unwrap();
464    ///     let decoder = DecoderBuilder::default().with_schema(schema).build()?;
465    ///     assert_eq!(decoder.input_dims(), Some((640, 640)));
466    /// #   Ok(())
467    /// # }
468    /// ```
469    pub fn input_dims(&self) -> Option<(usize, usize)> {
470        self.input_dims
471    }
472
473    /// Decode quantized model outputs into detection boxes and segmentation
474    /// masks. The quantized outputs can be of u8, i8, u16, i16, u32, or i32
475    /// types. Clears the provided output vectors before populating them.
476    pub(crate) fn decode_quantized(
477        &self,
478        outputs: &[ArrayViewDQuantized],
479        output_boxes: &mut Vec<DetectBox>,
480        output_masks: &mut Vec<Segmentation>,
481    ) -> Result<(), DecoderError> {
482        output_boxes.clear();
483        output_masks.clear();
484        match &self.model_type {
485            ModelType::ModelPackSegDet {
486                boxes,
487                scores,
488                segmentation,
489            } => {
490                self.decode_modelpack_det_quantized(outputs, boxes, scores, output_boxes)?;
491                self.decode_modelpack_seg_quantized(outputs, segmentation, output_masks)
492            }
493            ModelType::ModelPackSegDetSplit {
494                detection,
495                segmentation,
496            } => {
497                self.decode_modelpack_det_split_quantized(outputs, detection, output_boxes)?;
498                self.decode_modelpack_seg_quantized(outputs, segmentation, output_masks)
499            }
500            ModelType::ModelPackDet { boxes, scores } => {
501                self.decode_modelpack_det_quantized(outputs, boxes, scores, output_boxes)
502            }
503            ModelType::ModelPackDetSplit { detection } => {
504                self.decode_modelpack_det_split_quantized(outputs, detection, output_boxes)
505            }
506            ModelType::ModelPackSeg { segmentation } => {
507                self.decode_modelpack_seg_quantized(outputs, segmentation, output_masks)
508            }
509            ModelType::YoloDet { boxes } => {
510                self.decode_yolo_det_quantized(outputs, boxes, output_boxes)
511            }
512            ModelType::YoloSegDet { boxes, protos } => self.decode_yolo_segdet_quantized(
513                outputs,
514                boxes,
515                protos,
516                output_boxes,
517                output_masks,
518            ),
519            ModelType::YoloSplitDet { boxes, scores } => {
520                self.decode_yolo_split_det_quantized(outputs, boxes, scores, output_boxes)
521            }
522            ModelType::YoloSplitSegDet {
523                boxes,
524                scores,
525                mask_coeff,
526                protos,
527            } => self.decode_yolo_split_segdet_quantized(
528                outputs,
529                boxes,
530                scores,
531                mask_coeff,
532                protos,
533                output_boxes,
534                output_masks,
535            ),
536            ModelType::YoloSegDet2Way {
537                boxes,
538                mask_coeff,
539                protos,
540            } => self.decode_yolo_segdet_2way_quantized(
541                outputs,
542                boxes,
543                mask_coeff,
544                protos,
545                output_boxes,
546                output_masks,
547            ),
548            ModelType::YoloEndToEndDet { boxes } => {
549                self.decode_yolo_end_to_end_det_quantized(outputs, boxes, output_boxes)
550            }
551            ModelType::YoloEndToEndSegDet { boxes, protos } => self
552                .decode_yolo_end_to_end_segdet_quantized(
553                    outputs,
554                    boxes,
555                    protos,
556                    output_boxes,
557                    output_masks,
558                ),
559            ModelType::YoloSplitEndToEndDet {
560                boxes,
561                scores,
562                classes,
563            } => self.decode_yolo_split_end_to_end_det_quantized(
564                outputs,
565                boxes,
566                scores,
567                classes,
568                output_boxes,
569            ),
570            ModelType::YoloSplitEndToEndSegDet {
571                boxes,
572                scores,
573                classes,
574                mask_coeff,
575                protos,
576            } => self.decode_yolo_split_end_to_end_segdet_quantized(
577                outputs,
578                boxes,
579                scores,
580                classes,
581                mask_coeff,
582                protos,
583                output_boxes,
584                output_masks,
585            ),
586            ModelType::PerScale => Err(DecoderError::Internal(
587                "per-scale path must be intercepted before ModelType dispatch".into(),
588            )),
589        }
590    }
591
592    /// Decode floating point model outputs into detection boxes and
593    /// segmentation masks. Clears the provided output vectors before
594    /// populating them.
595    pub(crate) fn decode_float<T>(
596        &self,
597        outputs: &[ArrayViewD<T>],
598        output_boxes: &mut Vec<DetectBox>,
599        output_masks: &mut Vec<Segmentation>,
600    ) -> Result<(), DecoderError>
601    where
602        T: Float + AsPrimitive<f32> + AsPrimitive<u8> + Send + Sync + 'static,
603        f32: AsPrimitive<T>,
604    {
605        output_boxes.clear();
606        output_masks.clear();
607        match &self.model_type {
608            ModelType::ModelPackSegDet {
609                boxes,
610                scores,
611                segmentation,
612            } => {
613                self.decode_modelpack_det_float(outputs, boxes, scores, output_boxes)?;
614                self.decode_modelpack_seg_float(outputs, segmentation, output_masks)?;
615            }
616            ModelType::ModelPackSegDetSplit {
617                detection,
618                segmentation,
619            } => {
620                self.decode_modelpack_det_split_float(outputs, detection, output_boxes)?;
621                self.decode_modelpack_seg_float(outputs, segmentation, output_masks)?;
622            }
623            ModelType::ModelPackDet { boxes, scores } => {
624                self.decode_modelpack_det_float(outputs, boxes, scores, output_boxes)?;
625            }
626            ModelType::ModelPackDetSplit { detection } => {
627                self.decode_modelpack_det_split_float(outputs, detection, output_boxes)?;
628            }
629            ModelType::ModelPackSeg { segmentation } => {
630                self.decode_modelpack_seg_float(outputs, segmentation, output_masks)?;
631            }
632            ModelType::YoloDet { boxes } => {
633                self.decode_yolo_det_float(outputs, boxes, output_boxes)?;
634            }
635            ModelType::YoloSegDet { boxes, protos } => {
636                self.decode_yolo_segdet_float(outputs, boxes, protos, output_boxes, output_masks)?;
637            }
638            ModelType::YoloSplitDet { boxes, scores } => {
639                self.decode_yolo_split_det_float(outputs, boxes, scores, output_boxes)?;
640            }
641            ModelType::YoloSplitSegDet {
642                boxes,
643                scores,
644                mask_coeff,
645                protos,
646            } => {
647                self.decode_yolo_split_segdet_float(
648                    outputs,
649                    boxes,
650                    scores,
651                    mask_coeff,
652                    protos,
653                    output_boxes,
654                    output_masks,
655                )?;
656            }
657            ModelType::YoloSegDet2Way {
658                boxes,
659                mask_coeff,
660                protos,
661            } => {
662                self.decode_yolo_segdet_2way_float(
663                    outputs,
664                    boxes,
665                    mask_coeff,
666                    protos,
667                    output_boxes,
668                    output_masks,
669                )?;
670            }
671            ModelType::YoloEndToEndDet { boxes } => {
672                self.decode_yolo_end_to_end_det_float(outputs, boxes, output_boxes)?;
673            }
674            ModelType::YoloEndToEndSegDet { boxes, protos } => {
675                self.decode_yolo_end_to_end_segdet_float(
676                    outputs,
677                    boxes,
678                    protos,
679                    output_boxes,
680                    output_masks,
681                )?;
682            }
683            ModelType::YoloSplitEndToEndDet {
684                boxes,
685                scores,
686                classes,
687            } => {
688                self.decode_yolo_split_end_to_end_det_float(
689                    outputs,
690                    boxes,
691                    scores,
692                    classes,
693                    output_boxes,
694                )?;
695            }
696            ModelType::YoloSplitEndToEndSegDet {
697                boxes,
698                scores,
699                classes,
700                mask_coeff,
701                protos,
702            } => {
703                self.decode_yolo_split_end_to_end_segdet_float(
704                    outputs,
705                    boxes,
706                    scores,
707                    classes,
708                    mask_coeff,
709                    protos,
710                    output_boxes,
711                    output_masks,
712                )?;
713            }
714            ModelType::PerScale => {
715                return Err(DecoderError::Internal(
716                    "per-scale path must be intercepted before ModelType dispatch".into(),
717                ));
718            }
719        }
720        Ok(())
721    }
722
723    /// Decodes quantized model outputs into detection boxes, returning raw
724    /// `ProtoData` for segmentation models instead of materialized masks.
725    ///
726    /// Returns `Ok(None)` for detection-only and ModelPack models (detections
727    /// are still decoded into `output_boxes`). Returns `Ok(Some(ProtoData))`
728    /// for YOLO segmentation models.
729    pub(crate) fn decode_quantized_proto(
730        &self,
731        outputs: &[ArrayViewDQuantized],
732        output_boxes: &mut Vec<DetectBox>,
733    ) -> Result<Option<ProtoData>, DecoderError> {
734        output_boxes.clear();
735        match &self.model_type {
736            // Detection-only variants: decode boxes, return None for proto data.
737            ModelType::ModelPackDet { boxes, scores } => {
738                self.decode_modelpack_det_quantized(outputs, boxes, scores, output_boxes)?;
739                Ok(None)
740            }
741            ModelType::ModelPackDetSplit { detection } => {
742                self.decode_modelpack_det_split_quantized(outputs, detection, output_boxes)?;
743                Ok(None)
744            }
745            ModelType::YoloDet { boxes } => {
746                self.decode_yolo_det_quantized(outputs, boxes, output_boxes)?;
747                Ok(None)
748            }
749            ModelType::YoloSplitDet { boxes, scores } => {
750                self.decode_yolo_split_det_quantized(outputs, boxes, scores, output_boxes)?;
751                Ok(None)
752            }
753            ModelType::YoloEndToEndDet { boxes } => {
754                self.decode_yolo_end_to_end_det_quantized(outputs, boxes, output_boxes)?;
755                Ok(None)
756            }
757            ModelType::YoloSplitEndToEndDet {
758                boxes,
759                scores,
760                classes,
761            } => {
762                self.decode_yolo_split_end_to_end_det_quantized(
763                    outputs,
764                    boxes,
765                    scores,
766                    classes,
767                    output_boxes,
768                )?;
769                Ok(None)
770            }
771            // ModelPack seg/segdet variants have no YOLO proto data.
772            ModelType::ModelPackSegDet { boxes, scores, .. } => {
773                self.decode_modelpack_det_quantized(outputs, boxes, scores, output_boxes)?;
774                Ok(None)
775            }
776            ModelType::ModelPackSegDetSplit { detection, .. } => {
777                self.decode_modelpack_det_split_quantized(outputs, detection, output_boxes)?;
778                Ok(None)
779            }
780            ModelType::ModelPackSeg { .. } => Ok(None),
781
782            ModelType::YoloSegDet { boxes, protos } => {
783                let proto =
784                    self.decode_yolo_segdet_quantized_proto(outputs, boxes, protos, output_boxes)?;
785                Ok(Some(proto))
786            }
787            ModelType::YoloSplitSegDet {
788                boxes,
789                scores,
790                mask_coeff,
791                protos,
792            } => {
793                let proto = self.decode_yolo_split_segdet_quantized_proto(
794                    outputs,
795                    boxes,
796                    scores,
797                    mask_coeff,
798                    protos,
799                    output_boxes,
800                )?;
801                Ok(Some(proto))
802            }
803            ModelType::YoloSegDet2Way {
804                boxes,
805                mask_coeff,
806                protos,
807            } => {
808                let proto = self.decode_yolo_segdet_2way_quantized_proto(
809                    outputs,
810                    boxes,
811                    mask_coeff,
812                    protos,
813                    output_boxes,
814                )?;
815                Ok(Some(proto))
816            }
817            ModelType::YoloEndToEndSegDet { boxes, protos } => {
818                let proto = self.decode_yolo_end_to_end_segdet_quantized_proto(
819                    outputs,
820                    boxes,
821                    protos,
822                    output_boxes,
823                )?;
824                Ok(Some(proto))
825            }
826            ModelType::YoloSplitEndToEndSegDet {
827                boxes,
828                scores,
829                classes,
830                mask_coeff,
831                protos,
832            } => {
833                let proto = self.decode_yolo_split_end_to_end_segdet_quantized_proto(
834                    outputs,
835                    boxes,
836                    scores,
837                    classes,
838                    mask_coeff,
839                    protos,
840                    output_boxes,
841                )?;
842                Ok(Some(proto))
843            }
844            ModelType::PerScale => Err(DecoderError::Internal(
845                "per-scale path must be intercepted before ModelType dispatch".into(),
846            )),
847        }
848    }
849
850    /// Decodes floating-point model outputs into detection boxes, returning
851    /// raw `ProtoData` for segmentation models instead of materialized masks.
852    ///
853    /// Returns `Ok(None)` for detection-only and ModelPack models (detections
854    /// are still decoded into `output_boxes`). Returns `Ok(Some(ProtoData))`
855    /// for YOLO segmentation models.
856    pub(crate) fn decode_float_proto<T>(
857        &self,
858        outputs: &[ArrayViewD<T>],
859        output_boxes: &mut Vec<DetectBox>,
860    ) -> Result<Option<ProtoData>, DecoderError>
861    where
862        T: Float + AsPrimitive<f32> + AsPrimitive<u8> + Send + Sync + crate::yolo::FloatProtoElem,
863        f32: AsPrimitive<T>,
864    {
865        output_boxes.clear();
866        match &self.model_type {
867            // Detection-only variants: decode boxes, return None for proto data.
868            ModelType::ModelPackDet { boxes, scores } => {
869                self.decode_modelpack_det_float(outputs, boxes, scores, output_boxes)?;
870                Ok(None)
871            }
872            ModelType::ModelPackDetSplit { detection } => {
873                self.decode_modelpack_det_split_float(outputs, detection, output_boxes)?;
874                Ok(None)
875            }
876            ModelType::YoloDet { boxes } => {
877                self.decode_yolo_det_float(outputs, boxes, output_boxes)?;
878                Ok(None)
879            }
880            ModelType::YoloSplitDet { boxes, scores } => {
881                self.decode_yolo_split_det_float(outputs, boxes, scores, output_boxes)?;
882                Ok(None)
883            }
884            ModelType::YoloEndToEndDet { boxes } => {
885                self.decode_yolo_end_to_end_det_float(outputs, boxes, output_boxes)?;
886                Ok(None)
887            }
888            ModelType::YoloSplitEndToEndDet {
889                boxes,
890                scores,
891                classes,
892            } => {
893                self.decode_yolo_split_end_to_end_det_float(
894                    outputs,
895                    boxes,
896                    scores,
897                    classes,
898                    output_boxes,
899                )?;
900                Ok(None)
901            }
902            // ModelPack seg/segdet variants have no YOLO proto data.
903            ModelType::ModelPackSegDet { boxes, scores, .. } => {
904                self.decode_modelpack_det_float(outputs, boxes, scores, output_boxes)?;
905                Ok(None)
906            }
907            ModelType::ModelPackSegDetSplit { detection, .. } => {
908                self.decode_modelpack_det_split_float(outputs, detection, output_boxes)?;
909                Ok(None)
910            }
911            ModelType::ModelPackSeg { .. } => Ok(None),
912
913            ModelType::YoloSegDet { boxes, protos } => {
914                let proto =
915                    self.decode_yolo_segdet_float_proto(outputs, boxes, protos, output_boxes)?;
916                Ok(Some(proto))
917            }
918            ModelType::YoloSplitSegDet {
919                boxes,
920                scores,
921                mask_coeff,
922                protos,
923            } => {
924                let proto = self.decode_yolo_split_segdet_float_proto(
925                    outputs,
926                    boxes,
927                    scores,
928                    mask_coeff,
929                    protos,
930                    output_boxes,
931                )?;
932                Ok(Some(proto))
933            }
934            ModelType::YoloSegDet2Way {
935                boxes,
936                mask_coeff,
937                protos,
938            } => {
939                let proto = self.decode_yolo_segdet_2way_float_proto(
940                    outputs,
941                    boxes,
942                    mask_coeff,
943                    protos,
944                    output_boxes,
945                )?;
946                Ok(Some(proto))
947            }
948            ModelType::YoloEndToEndSegDet { boxes, protos } => {
949                let proto = self.decode_yolo_end_to_end_segdet_float_proto(
950                    outputs,
951                    boxes,
952                    protos,
953                    output_boxes,
954                )?;
955                Ok(Some(proto))
956            }
957            ModelType::YoloSplitEndToEndSegDet {
958                boxes,
959                scores,
960                classes,
961                mask_coeff,
962                protos,
963            } => {
964                let proto = self.decode_yolo_split_end_to_end_segdet_float_proto(
965                    outputs,
966                    boxes,
967                    scores,
968                    classes,
969                    mask_coeff,
970                    protos,
971                    output_boxes,
972                )?;
973                Ok(Some(proto))
974            }
975            ModelType::PerScale => Err(DecoderError::Internal(
976                "per-scale path must be intercepted before ModelType dispatch".into(),
977            )),
978        }
979    }
980
981    // ========================================================================
982    // TensorDyn-based public API
983    // ========================================================================
984
985    /// Decode model outputs into detection boxes and segmentation masks.
986    ///
987    /// This is the primary decode API. Accepts `TensorDyn` outputs directly
988    /// from model inference. Automatically dispatches to quantized or float
989    /// paths based on the tensor dtype.
990    ///
991    /// # Arguments
992    ///
993    /// * `outputs` - Tensor outputs from model inference
994    /// * `output_boxes` - Destination for decoded detection boxes (cleared first)
995    /// * `output_masks` - Destination for decoded segmentation masks (cleared first)
996    ///
997    /// # `output_boxes` / `output_masks` capacity
998    ///
999    /// The capacity of the supplied `Vec`s is **only** an allocation hint —
1000    /// it is **not** a cap on the number of detections returned. The
1001    /// post-NMS detection count is bounded by [`Decoder::max_det`] (set
1002    /// via [`DecoderBuilder::with_max_det`], default `300`). Passing
1003    /// `Vec::new()` (capacity 0) returns up to `max_det` detections;
1004    /// pre-allocating with [`Vec::with_capacity`] only avoids the
1005    /// reallocation when the decoder grows the buffer.
1006    ///
1007    /// # Errors
1008    ///
1009    /// Returns `DecoderError` if tensor mapping fails, dtypes are unsupported,
1010    /// or the outputs don't match the decoder's model configuration.
1011    pub fn decode(
1012        &self,
1013        outputs: &[&edgefirst_tensor::TensorDyn],
1014        output_boxes: &mut Vec<DetectBox>,
1015        output_masks: &mut Vec<Segmentation>,
1016    ) -> Result<(), DecoderError> {
1017        let path = self.decode_path_label();
1018        let _span = tracing::trace_span!("decoder.decode", path = path, n_outputs = outputs.len())
1019            .entered();
1020        // Per-scale fast path — selected at builder time when the schema
1021        // declares per-scale children with DFL or LTRB encoding.
1022        if let Some(per_scale_mutex) = &self.per_scale {
1023            let mut ps = per_scale_mutex
1024                .lock()
1025                .map_err(|e| DecoderError::Internal(format!("per_scale mutex poisoned: {e}")))?;
1026            let decoded = ps.run(outputs)?;
1027            return per_scale_bridge::per_scale_to_masks(
1028                &decoded,
1029                output_boxes,
1030                output_masks,
1031                self.iou_threshold,
1032                self.score_threshold,
1033                self.nms,
1034                self.pre_nms_top_k,
1035                self.max_det,
1036                self.normalized,
1037                self.input_dims,
1038                self.multi_label,
1039            );
1040        }
1041
1042        // Schema v2 merge path: dequantize physical children into
1043        // logical float32 tensors, then feed through the float dispatch.
1044        if let Some(program) = &self.decode_program {
1045            let merged = program.execute(outputs)?;
1046            let views: Vec<_> = merged.iter().map(|a| a.view()).collect();
1047            return self.decode_float(&views, output_boxes, output_masks);
1048        }
1049
1050        let mapped = tensor_bridge::map_tensors(outputs)?;
1051        match &mapped {
1052            tensor_bridge::MappedOutputs::Quantized(maps) => {
1053                let views = tensor_bridge::quantized_views(maps)?;
1054                self.decode_quantized(&views, output_boxes, output_masks)
1055            }
1056            tensor_bridge::MappedOutputs::Float16(maps) => {
1057                let views = tensor_bridge::f16_views(maps)?;
1058                self.decode_float(&views, output_boxes, output_masks)
1059            }
1060            tensor_bridge::MappedOutputs::Float32(maps) => {
1061                let views = tensor_bridge::f32_views(maps)?;
1062                self.decode_float(&views, output_boxes, output_masks)
1063            }
1064            tensor_bridge::MappedOutputs::Float64(maps) => {
1065                let views = tensor_bridge::f64_views(maps)?;
1066                self.decode_float(&views, output_boxes, output_masks)
1067            }
1068        }
1069    }
1070
1071    /// Decode model outputs into detection boxes, returning raw proto data
1072    /// for segmentation models instead of materialized masks.
1073    ///
1074    /// Accepts `TensorDyn` outputs directly from model inference.
1075    /// Detections are always decoded into `output_boxes` regardless of model type.
1076    /// Returns `Ok(None)` for detection-only and ModelPack models.
1077    /// Returns `Ok(Some(ProtoData))` for YOLO segmentation models.
1078    ///
1079    /// # Arguments
1080    ///
1081    /// * `outputs` - Tensor outputs from model inference
1082    /// * `output_boxes` - Destination for decoded detection boxes (cleared first)
1083    ///
1084    /// # `output_boxes` capacity
1085    ///
1086    /// The capacity of `output_boxes` is **only** an allocation hint — it
1087    /// is **not** a cap on the number of detections returned. The
1088    /// post-NMS detection count is bounded by [`Decoder::max_det`] (set
1089    /// via [`DecoderBuilder::with_max_det`], default `300`). Passing
1090    /// `Vec::new()` (capacity 0) returns up to `max_det` detections.
1091    ///
1092    /// # Errors
1093    ///
1094    /// Returns `DecoderError` if tensor mapping fails, dtypes are unsupported,
1095    /// or the outputs don't match the decoder's model configuration.
1096    pub fn decode_proto(
1097        &self,
1098        outputs: &[&edgefirst_tensor::TensorDyn],
1099        output_boxes: &mut Vec<DetectBox>,
1100    ) -> Result<Option<ProtoData>, DecoderError> {
1101        let path = self.decode_path_label();
1102        let _span = tracing::trace_span!(
1103            "decoder.decode_proto",
1104            path = path,
1105            n_outputs = outputs.len()
1106        )
1107        .entered();
1108        // Per-scale fast path — selected at builder time when the schema
1109        // declares per-scale children with DFL or LTRB encoding.
1110        if let Some(per_scale_mutex) = &self.per_scale {
1111            let mut ps = per_scale_mutex
1112                .lock()
1113                .map_err(|e| DecoderError::Internal(format!("per_scale mutex poisoned: {e}")))?;
1114            let decoded = ps.run(outputs)?;
1115            return per_scale_bridge::per_scale_to_proto_data(
1116                &decoded,
1117                output_boxes,
1118                self.iou_threshold,
1119                self.score_threshold,
1120                self.nms,
1121                self.pre_nms_top_k,
1122                self.max_det,
1123                self.normalized,
1124                self.input_dims,
1125                self.multi_label,
1126            );
1127        }
1128
1129        // Schema v2 merge path: dequantize physical children into
1130        // logical float32 tensors, then feed through the float dispatch.
1131        if let Some(program) = &self.decode_program {
1132            let merged = program.execute(outputs)?;
1133            let views: Vec<_> = merged.iter().map(|a| a.view()).collect();
1134            return self.decode_float_proto(&views, output_boxes);
1135        }
1136
1137        let mapped = tensor_bridge::map_tensors(outputs)?;
1138        let result = match &mapped {
1139            tensor_bridge::MappedOutputs::Quantized(maps) => {
1140                let views = tensor_bridge::quantized_views(maps)?;
1141                self.decode_quantized_proto(&views, output_boxes)
1142            }
1143            tensor_bridge::MappedOutputs::Float16(maps) => {
1144                let views = tensor_bridge::f16_views(maps)?;
1145                self.decode_float_proto(&views, output_boxes)
1146            }
1147            tensor_bridge::MappedOutputs::Float32(maps) => {
1148                let views = tensor_bridge::f32_views(maps)?;
1149                self.decode_float_proto(&views, output_boxes)
1150            }
1151            tensor_bridge::MappedOutputs::Float64(maps) => {
1152                let views = tensor_bridge::f64_views(maps)?;
1153                self.decode_float_proto(&views, output_boxes)
1154            }
1155        };
1156        result
1157    }
1158
1159    /// Run the per-scale pipeline and return pre-NMS buffers as owned f32.
1160    ///
1161    /// Test-only entry point used by the parity-fixture tests to compare
1162    /// HAL stage output against the NumPy reference's stage output
1163    /// without NMS ordering noise. Returns an error if the decoder
1164    /// isn't configured for per-scale decoding.
1165    #[doc(hidden)]
1166    pub fn _testing_run_per_scale_pre_nms(
1167        &self,
1168        outputs: &[&edgefirst_tensor::TensorDyn],
1169    ) -> Result<crate::per_scale::PreNmsCapture, crate::error::DecoderError> {
1170        let mutex = self.per_scale.as_ref().ok_or_else(|| {
1171            crate::error::DecoderError::Internal("decoder not configured for per-scale".into())
1172        })?;
1173        let mut ps = mutex.lock().map_err(|e| {
1174            crate::error::DecoderError::Internal(format!("per_scale mutex poisoned: {e}"))
1175        })?;
1176        // Drop the borrowed view immediately so we can reborrow buffers below.
1177        {
1178            ps.run(outputs)?;
1179        }
1180        let total_anchors = ps.plan.total_anchors;
1181        let num_classes = ps.plan.num_classes;
1182        let num_mc = ps.plan.num_mask_coefs;
1183        Ok(ps
1184            .buffers
1185            .snapshot_owned_f32(total_anchors, num_classes, num_mc))
1186    }
1187}
1188
1189#[cfg(feature = "tracker")]
1190pub use edgefirst_tracker::TrackInfo;
1191
1192#[cfg(feature = "tracker")]
1193pub use edgefirst_tracker::Tracker;
1194
1195#[cfg(feature = "tracker")]
1196impl Decoder {
1197    /// Decode quantized model outputs into detection boxes and segmentation
1198    /// masks with tracking. Clears the provided output vectors before
1199    /// populating them.
1200    pub(crate) fn decode_tracked_quantized<TR: edgefirst_tracker::Tracker<DetectBox>>(
1201        &self,
1202        tracker: &mut TR,
1203        timestamp: u64,
1204        outputs: &[ArrayViewDQuantized],
1205        output_boxes: &mut Vec<DetectBox>,
1206        output_masks: &mut Vec<Segmentation>,
1207        output_tracks: &mut Vec<edgefirst_tracker::TrackInfo>,
1208    ) -> Result<(), DecoderError> {
1209        // multi-label duplicates would spawn phantom tracks in ByteTrack
1210        // (tracker matches IoU-only; two boxes for the same anchor at different
1211        // labels become two tracks). Deployment must always stay on argmax.
1212        debug_assert!(
1213            !self.multi_label,
1214            "multi_label must be off on the tracked/deployment path"
1215        );
1216        output_boxes.clear();
1217        output_masks.clear();
1218        output_tracks.clear();
1219
1220        // yolo segdet variants require special handling to separate boxes that come from decoding vs active tracks.
1221        // Only boxes that come from decoding can be used for proto/mask generation.
1222        match &self.model_type {
1223            ModelType::YoloSegDet { boxes, protos } => self.decode_tracked_yolo_segdet_quantized(
1224                tracker,
1225                timestamp,
1226                outputs,
1227                boxes,
1228                protos,
1229                output_boxes,
1230                output_masks,
1231                output_tracks,
1232            ),
1233            ModelType::YoloSplitSegDet {
1234                boxes,
1235                scores,
1236                mask_coeff,
1237                protos,
1238            } => self.decode_tracked_yolo_split_segdet_quantized(
1239                tracker,
1240                timestamp,
1241                outputs,
1242                boxes,
1243                scores,
1244                mask_coeff,
1245                protos,
1246                output_boxes,
1247                output_masks,
1248                output_tracks,
1249            ),
1250            ModelType::YoloEndToEndSegDet { boxes, protos } => self
1251                .decode_tracked_yolo_end_to_end_segdet_quantized(
1252                    tracker,
1253                    timestamp,
1254                    outputs,
1255                    boxes,
1256                    protos,
1257                    output_boxes,
1258                    output_masks,
1259                    output_tracks,
1260                ),
1261            ModelType::YoloSplitEndToEndSegDet {
1262                boxes,
1263                scores,
1264                classes,
1265                mask_coeff,
1266                protos,
1267            } => self.decode_tracked_yolo_split_end_to_end_segdet_quantized(
1268                tracker,
1269                timestamp,
1270                outputs,
1271                boxes,
1272                scores,
1273                classes,
1274                mask_coeff,
1275                protos,
1276                output_boxes,
1277                output_masks,
1278                output_tracks,
1279            ),
1280            ModelType::YoloSegDet2Way {
1281                boxes,
1282                mask_coeff,
1283                protos,
1284            } => self.decode_tracked_yolo_segdet_2way_quantized(
1285                tracker,
1286                timestamp,
1287                outputs,
1288                boxes,
1289                mask_coeff,
1290                protos,
1291                output_boxes,
1292                output_masks,
1293                output_tracks,
1294            ),
1295            _ => {
1296                self.decode_quantized(outputs, output_boxes, output_masks)?;
1297                Self::update_tracker(tracker, timestamp, output_boxes, output_tracks);
1298                Ok(())
1299            }
1300        }
1301    }
1302
1303    /// This function decodes floating point model outputs into detection boxes
1304    /// and segmentation masks. Up to `output_boxes.capacity()` boxes and
1305    /// masks will be decoded. The function clears the provided output
1306    /// vectors before populating them with the decoded results.
1307    ///
1308    /// This function returns an `Error` if the provided outputs don't
1309    /// match the configuration provided by the user when building the decoder.
1310    ///
1311    /// Any quantization information in the configuration will be ignored.
1312    pub(crate) fn decode_tracked_float<TR: edgefirst_tracker::Tracker<DetectBox>, T>(
1313        &self,
1314        tracker: &mut TR,
1315        timestamp: u64,
1316        outputs: &[ArrayViewD<T>],
1317        output_boxes: &mut Vec<DetectBox>,
1318        output_masks: &mut Vec<Segmentation>,
1319        output_tracks: &mut Vec<edgefirst_tracker::TrackInfo>,
1320    ) -> Result<(), DecoderError>
1321    where
1322        T: Float + AsPrimitive<f32> + AsPrimitive<u8> + Send + Sync + 'static,
1323        f32: AsPrimitive<T>,
1324    {
1325        debug_assert!(
1326            !self.multi_label,
1327            "multi_label must be off on the tracked/deployment path"
1328        );
1329        output_boxes.clear();
1330        output_masks.clear();
1331        output_tracks.clear();
1332        match &self.model_type {
1333            ModelType::YoloSegDet { boxes, protos } => {
1334                self.decode_tracked_yolo_segdet_float(
1335                    tracker,
1336                    timestamp,
1337                    outputs,
1338                    boxes,
1339                    protos,
1340                    output_boxes,
1341                    output_masks,
1342                    output_tracks,
1343                )?;
1344            }
1345            ModelType::YoloSplitSegDet {
1346                boxes,
1347                scores,
1348                mask_coeff,
1349                protos,
1350            } => {
1351                self.decode_tracked_yolo_split_segdet_float(
1352                    tracker,
1353                    timestamp,
1354                    outputs,
1355                    boxes,
1356                    scores,
1357                    mask_coeff,
1358                    protos,
1359                    output_boxes,
1360                    output_masks,
1361                    output_tracks,
1362                )?;
1363            }
1364            ModelType::YoloEndToEndSegDet { boxes, protos } => {
1365                self.decode_tracked_yolo_end_to_end_segdet_float(
1366                    tracker,
1367                    timestamp,
1368                    outputs,
1369                    boxes,
1370                    protos,
1371                    output_boxes,
1372                    output_masks,
1373                    output_tracks,
1374                )?;
1375            }
1376            ModelType::YoloSplitEndToEndSegDet {
1377                boxes,
1378                scores,
1379                classes,
1380                mask_coeff,
1381                protos,
1382            } => {
1383                self.decode_tracked_yolo_split_end_to_end_segdet_float(
1384                    tracker,
1385                    timestamp,
1386                    outputs,
1387                    boxes,
1388                    scores,
1389                    classes,
1390                    mask_coeff,
1391                    protos,
1392                    output_boxes,
1393                    output_masks,
1394                    output_tracks,
1395                )?;
1396            }
1397            ModelType::YoloSegDet2Way {
1398                boxes,
1399                mask_coeff,
1400                protos,
1401            } => {
1402                self.decode_tracked_yolo_segdet_2way_float(
1403                    tracker,
1404                    timestamp,
1405                    outputs,
1406                    boxes,
1407                    mask_coeff,
1408                    protos,
1409                    output_boxes,
1410                    output_masks,
1411                    output_tracks,
1412                )?;
1413            }
1414            _ => {
1415                self.decode_float(outputs, output_boxes, output_masks)?;
1416                Self::update_tracker(tracker, timestamp, output_boxes, output_tracks);
1417            }
1418        }
1419        Ok(())
1420    }
1421
1422    /// Decodes quantized model outputs into detection boxes, returning raw
1423    /// `ProtoData` for segmentation models instead of materialized masks.
1424    ///
1425    /// Returns `Ok(None)` for detection-only and ModelPack models (use
1426    /// `decode_quantized` for those). Returns `Ok(Some(ProtoData))` for
1427    /// YOLO segmentation models.
1428    pub(crate) fn decode_tracked_quantized_proto<TR: edgefirst_tracker::Tracker<DetectBox>>(
1429        &self,
1430        tracker: &mut TR,
1431        timestamp: u64,
1432        outputs: &[ArrayViewDQuantized],
1433        output_boxes: &mut Vec<DetectBox>,
1434        output_tracks: &mut Vec<edgefirst_tracker::TrackInfo>,
1435    ) -> Result<Option<ProtoData>, DecoderError> {
1436        debug_assert!(
1437            !self.multi_label,
1438            "multi_label must be off on the tracked/deployment path"
1439        );
1440        output_boxes.clear();
1441        output_tracks.clear();
1442        match &self.model_type {
1443            ModelType::YoloSegDet { boxes, protos } => {
1444                let proto = self.decode_tracked_yolo_segdet_quantized_proto(
1445                    tracker,
1446                    timestamp,
1447                    outputs,
1448                    boxes,
1449                    protos,
1450                    output_boxes,
1451                    output_tracks,
1452                )?;
1453                Ok(Some(proto))
1454            }
1455            ModelType::YoloSplitSegDet {
1456                boxes,
1457                scores,
1458                mask_coeff,
1459                protos,
1460            } => {
1461                let proto = self.decode_tracked_yolo_split_segdet_quantized_proto(
1462                    tracker,
1463                    timestamp,
1464                    outputs,
1465                    boxes,
1466                    scores,
1467                    mask_coeff,
1468                    protos,
1469                    output_boxes,
1470                    output_tracks,
1471                )?;
1472                Ok(Some(proto))
1473            }
1474            ModelType::YoloSegDet2Way {
1475                boxes,
1476                mask_coeff,
1477                protos,
1478            } => {
1479                let proto = self.decode_tracked_yolo_segdet_2way_quantized_proto(
1480                    tracker,
1481                    timestamp,
1482                    outputs,
1483                    boxes,
1484                    mask_coeff,
1485                    protos,
1486                    output_boxes,
1487                    output_tracks,
1488                )?;
1489                Ok(Some(proto))
1490            }
1491            ModelType::YoloEndToEndSegDet { boxes, protos } => {
1492                let proto = self.decode_tracked_yolo_end_to_end_segdet_quantized_proto(
1493                    tracker,
1494                    timestamp,
1495                    outputs,
1496                    boxes,
1497                    protos,
1498                    output_boxes,
1499                    output_tracks,
1500                )?;
1501                Ok(Some(proto))
1502            }
1503            ModelType::YoloSplitEndToEndSegDet {
1504                boxes,
1505                scores,
1506                classes,
1507                mask_coeff,
1508                protos,
1509            } => {
1510                let proto = self.decode_tracked_yolo_split_end_to_end_segdet_quantized_proto(
1511                    tracker,
1512                    timestamp,
1513                    outputs,
1514                    boxes,
1515                    scores,
1516                    classes,
1517                    mask_coeff,
1518                    protos,
1519                    output_boxes,
1520                    output_tracks,
1521                )?;
1522                Ok(Some(proto))
1523            }
1524            // Non-seg variants: decode boxes via the non-proto path, then track.
1525            _ => {
1526                let mut masks = Vec::new();
1527                self.decode_quantized(outputs, output_boxes, &mut masks)?;
1528                Self::update_tracker(tracker, timestamp, output_boxes, output_tracks);
1529                Ok(None)
1530            }
1531        }
1532    }
1533
1534    /// Decodes floating-point model outputs into detection boxes, returning
1535    /// raw `ProtoData` for segmentation models instead of materialized masks.
1536    ///
1537    /// Detections are always decoded into `output_boxes` regardless of model type.
1538    /// Returns `Ok(None)` for detection-only and ModelPack models. Returns
1539    /// `Ok(Some(ProtoData))` for YOLO segmentation models.
1540    pub(crate) fn decode_tracked_float_proto<TR: edgefirst_tracker::Tracker<DetectBox>, T>(
1541        &self,
1542        tracker: &mut TR,
1543        timestamp: u64,
1544        outputs: &[ArrayViewD<T>],
1545        output_boxes: &mut Vec<DetectBox>,
1546        output_tracks: &mut Vec<edgefirst_tracker::TrackInfo>,
1547    ) -> Result<Option<ProtoData>, DecoderError>
1548    where
1549        T: Float + AsPrimitive<f32> + AsPrimitive<u8> + Send + Sync + crate::yolo::FloatProtoElem,
1550        f32: AsPrimitive<T>,
1551    {
1552        output_boxes.clear();
1553        output_tracks.clear();
1554        match &self.model_type {
1555            ModelType::YoloSegDet { boxes, protos } => {
1556                let proto = self.decode_tracked_yolo_segdet_float_proto(
1557                    tracker,
1558                    timestamp,
1559                    outputs,
1560                    boxes,
1561                    protos,
1562                    output_boxes,
1563                    output_tracks,
1564                )?;
1565                Ok(Some(proto))
1566            }
1567            ModelType::YoloSplitSegDet {
1568                boxes,
1569                scores,
1570                mask_coeff,
1571                protos,
1572            } => {
1573                let proto = self.decode_tracked_yolo_split_segdet_float_proto(
1574                    tracker,
1575                    timestamp,
1576                    outputs,
1577                    boxes,
1578                    scores,
1579                    mask_coeff,
1580                    protos,
1581                    output_boxes,
1582                    output_tracks,
1583                )?;
1584                Ok(Some(proto))
1585            }
1586            ModelType::YoloSegDet2Way {
1587                boxes,
1588                mask_coeff,
1589                protos,
1590            } => {
1591                let proto = self.decode_tracked_yolo_segdet_2way_float_proto(
1592                    tracker,
1593                    timestamp,
1594                    outputs,
1595                    boxes,
1596                    mask_coeff,
1597                    protos,
1598                    output_boxes,
1599                    output_tracks,
1600                )?;
1601                Ok(Some(proto))
1602            }
1603            ModelType::YoloEndToEndSegDet { boxes, protos } => {
1604                let proto = self.decode_tracked_yolo_end_to_end_segdet_float_proto(
1605                    tracker,
1606                    timestamp,
1607                    outputs,
1608                    boxes,
1609                    protos,
1610                    output_boxes,
1611                    output_tracks,
1612                )?;
1613                Ok(Some(proto))
1614            }
1615            ModelType::YoloSplitEndToEndSegDet {
1616                boxes,
1617                scores,
1618                classes,
1619                mask_coeff,
1620                protos,
1621            } => {
1622                let proto = self.decode_tracked_yolo_split_end_to_end_segdet_float_proto(
1623                    tracker,
1624                    timestamp,
1625                    outputs,
1626                    boxes,
1627                    scores,
1628                    classes,
1629                    mask_coeff,
1630                    protos,
1631                    output_boxes,
1632                    output_tracks,
1633                )?;
1634                Ok(Some(proto))
1635            }
1636            // Non-seg variants: decode boxes via the non-proto path, then track.
1637            _ => {
1638                let mut masks = Vec::new();
1639                self.decode_float(outputs, output_boxes, &mut masks)?;
1640                Self::update_tracker(tracker, timestamp, output_boxes, output_tracks);
1641                Ok(None)
1642            }
1643        }
1644    }
1645
1646    // ========================================================================
1647    // TensorDyn-based tracked public API
1648    // ========================================================================
1649
1650    /// Decode model outputs with tracking.
1651    ///
1652    /// Accepts `TensorDyn` outputs directly from model inference. Automatically
1653    /// dispatches to quantized or float paths based on the tensor dtype, then
1654    /// updates the tracker with the decoded boxes.
1655    ///
1656    /// # Arguments
1657    ///
1658    /// * `tracker` - The tracker instance to update
1659    /// * `timestamp` - Current frame timestamp
1660    /// * `outputs` - Tensor outputs from model inference
1661    /// * `output_boxes` - Destination for decoded detection boxes (cleared first)
1662    /// * `output_masks` - Destination for decoded segmentation masks (cleared first)
1663    /// * `output_tracks` - Destination for track info (cleared first)
1664    ///
1665    /// # Errors
1666    ///
1667    /// Returns `DecoderError` if tensor mapping fails, dtypes are unsupported,
1668    /// or the outputs don't match the decoder's model configuration.
1669    pub fn decode_tracked<TR: edgefirst_tracker::Tracker<DetectBox>>(
1670        &self,
1671        tracker: &mut TR,
1672        timestamp: u64,
1673        outputs: &[&edgefirst_tensor::TensorDyn],
1674        output_boxes: &mut Vec<DetectBox>,
1675        output_masks: &mut Vec<Segmentation>,
1676        output_tracks: &mut Vec<edgefirst_tracker::TrackInfo>,
1677    ) -> Result<(), DecoderError> {
1678        // Per-scale fast path: route via the basic decode then update the
1679        // tracker. The current implementation keeps the tracker integration simple; per-frame
1680        // decoupling between detection and tracking is preserved.
1681        if self.per_scale.is_some() {
1682            output_tracks.clear();
1683            self.decode(outputs, output_boxes, output_masks)?;
1684            Self::update_tracker(tracker, timestamp, output_boxes, output_tracks);
1685            return Ok(());
1686        }
1687
1688        let mapped = tensor_bridge::map_tensors(outputs)?;
1689        match &mapped {
1690            tensor_bridge::MappedOutputs::Quantized(maps) => {
1691                let views = tensor_bridge::quantized_views(maps)?;
1692                self.decode_tracked_quantized(
1693                    tracker,
1694                    timestamp,
1695                    &views,
1696                    output_boxes,
1697                    output_masks,
1698                    output_tracks,
1699                )
1700            }
1701            tensor_bridge::MappedOutputs::Float16(maps) => {
1702                let views = tensor_bridge::f16_views(maps)?;
1703                self.decode_tracked_float(
1704                    tracker,
1705                    timestamp,
1706                    &views,
1707                    output_boxes,
1708                    output_masks,
1709                    output_tracks,
1710                )
1711            }
1712            tensor_bridge::MappedOutputs::Float32(maps) => {
1713                let views = tensor_bridge::f32_views(maps)?;
1714                self.decode_tracked_float(
1715                    tracker,
1716                    timestamp,
1717                    &views,
1718                    output_boxes,
1719                    output_masks,
1720                    output_tracks,
1721                )
1722            }
1723            tensor_bridge::MappedOutputs::Float64(maps) => {
1724                let views = tensor_bridge::f64_views(maps)?;
1725                self.decode_tracked_float(
1726                    tracker,
1727                    timestamp,
1728                    &views,
1729                    output_boxes,
1730                    output_masks,
1731                    output_tracks,
1732                )
1733            }
1734        }
1735    }
1736
1737    /// Decode model outputs with tracking, returning raw proto data for
1738    /// segmentation models.
1739    ///
1740    /// Accepts `TensorDyn` outputs directly from model inference.
1741    /// Returns `Ok(None)` for detection-only and ModelPack models.
1742    /// Returns `Ok(Some(ProtoData))` for YOLO segmentation models.
1743    ///
1744    /// # Arguments
1745    ///
1746    /// * `tracker` - The tracker instance to update
1747    /// * `timestamp` - Current frame timestamp
1748    /// * `outputs` - Tensor outputs from model inference
1749    /// * `output_boxes` - Destination for decoded detection boxes (cleared first)
1750    /// * `output_tracks` - Destination for track info (cleared first)
1751    ///
1752    /// # Errors
1753    ///
1754    /// Returns `DecoderError` if tensor mapping fails, dtypes are unsupported,
1755    /// or the outputs don't match the decoder's model configuration.
1756    pub fn decode_proto_tracked<TR: edgefirst_tracker::Tracker<DetectBox>>(
1757        &self,
1758        tracker: &mut TR,
1759        timestamp: u64,
1760        outputs: &[&edgefirst_tensor::TensorDyn],
1761        output_boxes: &mut Vec<DetectBox>,
1762        output_tracks: &mut Vec<edgefirst_tracker::TrackInfo>,
1763    ) -> Result<Option<ProtoData>, DecoderError> {
1764        // Per-scale fast path: route via the basic decode_proto then
1765        // update the tracker on the resulting boxes.
1766        if self.per_scale.is_some() {
1767            output_tracks.clear();
1768            let proto = self.decode_proto(outputs, output_boxes)?;
1769            Self::update_tracker(tracker, timestamp, output_boxes, output_tracks);
1770            return Ok(proto);
1771        }
1772
1773        let mapped = tensor_bridge::map_tensors(outputs)?;
1774        match &mapped {
1775            tensor_bridge::MappedOutputs::Quantized(maps) => {
1776                let views = tensor_bridge::quantized_views(maps)?;
1777                self.decode_tracked_quantized_proto(
1778                    tracker,
1779                    timestamp,
1780                    &views,
1781                    output_boxes,
1782                    output_tracks,
1783                )
1784            }
1785            tensor_bridge::MappedOutputs::Float16(maps) => {
1786                let views = tensor_bridge::f16_views(maps)?;
1787                self.decode_tracked_float_proto(
1788                    tracker,
1789                    timestamp,
1790                    &views,
1791                    output_boxes,
1792                    output_tracks,
1793                )
1794            }
1795            tensor_bridge::MappedOutputs::Float32(maps) => {
1796                let views = tensor_bridge::f32_views(maps)?;
1797                self.decode_tracked_float_proto(
1798                    tracker,
1799                    timestamp,
1800                    &views,
1801                    output_boxes,
1802                    output_tracks,
1803                )
1804            }
1805            tensor_bridge::MappedOutputs::Float64(maps) => {
1806                let views = tensor_bridge::f64_views(maps)?;
1807                self.decode_tracked_float_proto(
1808                    tracker,
1809                    timestamp,
1810                    &views,
1811                    output_boxes,
1812                    output_tracks,
1813                )
1814            }
1815        }
1816    }
1817}