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