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}