edgefirst_decoder/decoder/
configs.rs

1// SPDX-FileCopyrightText: Copyright 2025 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4//! The vocabulary types a model config is built from.
5//!
6//! Per-role output descriptors ([`Detection`], [`Boxes`], [`Scores`],
7//! [`Classes`], [`Segmentation`], [`Protos`], [`MaskCoefficients`],
8//! [`Mask`]), the axis names used to declare physical layout ([`DimName`]),
9//! quantization pairs ([`QuantTuple`]), and the enums that steer decoding:
10//! [`DecoderType`], [`DecoderVersion`], [`Nms`], and the resolved
11//! [`ModelType`].
12//!
13//! [`Nms`] is worth a note: it has no `None` variant. Bypassing suppression is
14//! expressed as `Option<Nms>::None` on the decoder configuration, and
15//! [`Nms::Auto`] means "take the mode from the config document", falling back
16//! to [`Nms::ClassAgnostic`] when the document is silent.
17
18use std::collections::HashMap;
19use std::fmt::Display;
20
21use serde::{Deserialize, Serialize};
22
23/// Deserialize dshape from either array-of-tuples or array-of-single-key-dicts.
24///
25/// The metadata spec produces `[{"batch": 1}, {"num_features": 84}]` (dict format),
26/// while serde's default `Vec<(A, B)>` expects `[["batch", 1]]` (tuple format).
27/// This deserializer accepts both.
28pub fn deserialize_dshape<'de, D>(deserializer: D) -> Result<Vec<(DimName, usize)>, D::Error>
29where
30    D: serde::Deserializer<'de>,
31{
32    #[derive(Deserialize)]
33    #[serde(untagged)]
34    enum DShapeItem {
35        Tuple(DimName, usize),
36        Map(HashMap<DimName, usize>),
37    }
38
39    let items: Vec<DShapeItem> = Vec::deserialize(deserializer)?;
40    items
41        .into_iter()
42        .map(|item| match item {
43            DShapeItem::Tuple(name, size) => Ok((name, size)),
44            DShapeItem::Map(map) => {
45                if map.len() != 1 {
46                    return Err(serde::de::Error::custom(
47                        "dshape map entry must have exactly one key",
48                    ));
49                }
50                let (name, size) = map.into_iter().next().unwrap();
51                Ok((name, size))
52            }
53        })
54        .collect()
55}
56
57#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Copy)]
58pub struct QuantTuple(pub f32, pub i32);
59impl From<QuantTuple> for (f32, i32) {
60    fn from(value: QuantTuple) -> Self {
61        (value.0, value.1)
62    }
63}
64
65impl From<(f32, i32)> for QuantTuple {
66    fn from(value: (f32, i32)) -> Self {
67        QuantTuple(value.0, value.1)
68    }
69}
70
71#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Default)]
72pub struct Segmentation {
73    #[serde(default)]
74    pub decoder: DecoderType,
75    #[serde(default)]
76    pub quantization: Option<QuantTuple>,
77    #[serde(default)]
78    pub shape: Vec<usize>,
79    #[serde(default, deserialize_with = "deserialize_dshape")]
80    pub dshape: Vec<(DimName, usize)>,
81}
82
83#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Default)]
84pub struct Protos {
85    #[serde(default)]
86    pub decoder: DecoderType,
87    #[serde(default)]
88    pub quantization: Option<QuantTuple>,
89    #[serde(default)]
90    pub shape: Vec<usize>,
91    #[serde(default, deserialize_with = "deserialize_dshape")]
92    pub dshape: Vec<(DimName, usize)>,
93}
94
95#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Default)]
96pub struct MaskCoefficients {
97    #[serde(default)]
98    pub decoder: DecoderType,
99    #[serde(default)]
100    pub quantization: Option<QuantTuple>,
101    #[serde(default)]
102    pub shape: Vec<usize>,
103    #[serde(default, deserialize_with = "deserialize_dshape")]
104    pub dshape: Vec<(DimName, usize)>,
105}
106
107#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Default)]
108pub struct Mask {
109    #[serde(default)]
110    pub decoder: DecoderType,
111    #[serde(default)]
112    pub quantization: Option<QuantTuple>,
113    #[serde(default)]
114    pub shape: Vec<usize>,
115    #[serde(default, deserialize_with = "deserialize_dshape")]
116    pub dshape: Vec<(DimName, usize)>,
117}
118
119#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Default)]
120pub struct Detection {
121    #[serde(default)]
122    pub anchors: Option<Vec<[f32; 2]>>,
123    #[serde(default)]
124    pub decoder: DecoderType,
125    #[serde(default)]
126    pub quantization: Option<QuantTuple>,
127    #[serde(default)]
128    pub shape: Vec<usize>,
129    #[serde(default, deserialize_with = "deserialize_dshape")]
130    pub dshape: Vec<(DimName, usize)>,
131    /// Whether box coordinates are normalized to `[0,1]` range.
132    /// - `Some(true)`: Coordinates in `[0,1]` range relative to model input
133    /// - `Some(false)`: Pixel coordinates relative to model input
134    ///   (letterboxed)
135    /// - `None`: Unknown, caller must infer (e.g., check if any coordinate
136    ///   > 1.0)
137    #[serde(default)]
138    pub normalized: Option<bool>,
139}
140
141#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Default)]
142pub struct Scores {
143    #[serde(default)]
144    pub decoder: DecoderType,
145    #[serde(default)]
146    pub quantization: Option<QuantTuple>,
147    #[serde(default)]
148    pub shape: Vec<usize>,
149    #[serde(default, deserialize_with = "deserialize_dshape")]
150    pub dshape: Vec<(DimName, usize)>,
151}
152
153#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Default)]
154pub struct Boxes {
155    #[serde(default)]
156    pub decoder: DecoderType,
157    #[serde(default)]
158    pub quantization: Option<QuantTuple>,
159    #[serde(default)]
160    pub shape: Vec<usize>,
161    #[serde(default, deserialize_with = "deserialize_dshape")]
162    pub dshape: Vec<(DimName, usize)>,
163    /// Whether box coordinates are normalized to `[0,1]` range.
164    /// - `Some(true)`: Coordinates in `[0,1]` range relative to model input
165    /// - `Some(false)`: Pixel coordinates relative to model input
166    ///   (letterboxed)
167    /// - `None`: Unknown, caller must infer (e.g., check if any coordinate
168    ///   > 1.0)
169    #[serde(default)]
170    pub normalized: Option<bool>,
171}
172
173#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Default)]
174pub struct Classes {
175    #[serde(default)]
176    pub decoder: DecoderType,
177    #[serde(default)]
178    pub quantization: Option<QuantTuple>,
179    #[serde(default)]
180    pub shape: Vec<usize>,
181    #[serde(default, deserialize_with = "deserialize_dshape")]
182    pub dshape: Vec<(DimName, usize)>,
183}
184
185#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Copy, Hash, Eq)]
186pub enum DimName {
187    #[serde(rename = "batch")]
188    Batch,
189    #[serde(rename = "height")]
190    Height,
191    #[serde(rename = "width")]
192    Width,
193    #[serde(rename = "num_classes")]
194    NumClasses,
195    #[serde(rename = "num_features")]
196    NumFeatures,
197    #[serde(rename = "num_boxes")]
198    NumBoxes,
199    #[serde(rename = "num_protos")]
200    NumProtos,
201    #[serde(rename = "num_anchors_x_features")]
202    NumAnchorsXFeatures,
203    #[serde(rename = "padding")]
204    Padding,
205    #[serde(rename = "box_coords")]
206    BoxCoords,
207    /// Any axis name the HAL does not recognise (e.g. a producer's
208    /// `channels` on the input dshape). Preserved so the dshape length still
209    /// matches the shape and the axis sorts to the canonical tail in
210    /// `swap_axes_if_needed`, but it never satisfies a required-dimension
211    /// check. Keeps metadata parsing tolerant of unknown axis names instead
212    /// of failing the whole decoder build (DE-2651).
213    #[serde(other)]
214    Unknown,
215}
216
217impl Display for DimName {
218    /// Formats the DimName for display
219    /// # Examples
220    /// ```rust
221    /// # use edgefirst_decoder::configs::DimName;
222    /// let dim = DimName::Height;
223    /// assert_eq!(format!("{}", dim), "height");
224    /// # let s = format!("{} {} {} {} {} {} {} {} {} {}", DimName::Batch, DimName::Height, DimName::Width, DimName::NumClasses, DimName::NumFeatures, DimName::NumBoxes, DimName::NumProtos, DimName::NumAnchorsXFeatures, DimName::Padding, DimName::BoxCoords);
225    /// # assert_eq!(s, "batch height width num_classes num_features num_boxes num_protos num_anchors_x_features padding box_coords");
226    /// ```
227    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
228        match self {
229            DimName::Batch => write!(f, "batch"),
230            DimName::Height => write!(f, "height"),
231            DimName::Width => write!(f, "width"),
232            DimName::NumClasses => write!(f, "num_classes"),
233            DimName::NumFeatures => write!(f, "num_features"),
234            DimName::NumBoxes => write!(f, "num_boxes"),
235            DimName::NumProtos => write!(f, "num_protos"),
236            DimName::NumAnchorsXFeatures => write!(f, "num_anchors_x_features"),
237            DimName::Padding => write!(f, "padding"),
238            DimName::BoxCoords => write!(f, "box_coords"),
239            DimName::Unknown => write!(f, "unknown"),
240        }
241    }
242}
243
244#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Copy, Hash, Eq, Default)]
245pub enum DecoderType {
246    #[serde(rename = "modelpack")]
247    ModelPack,
248    #[default]
249    #[serde(rename = "ultralytics", alias = "yolov8")]
250    Ultralytics,
251}
252
253/// Decoder version for Ultralytics models.
254///
255/// Specifies the YOLO architecture version, which determines the decoding
256/// strategy:
257/// - `Yolov5`, `Yolov8`, `Yolo11`: Traditional models requiring external
258///   NMS
259/// - `Yolo26`: End-to-end models with NMS embedded in the model
260///   architecture
261///
262/// When `decoder_version` is set to `Yolo26`, the decoder uses end-to-end
263/// model types regardless of the `nms` setting.
264#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Copy, Hash, Eq)]
265#[serde(rename_all = "lowercase")]
266pub enum DecoderVersion {
267    /// YOLOv5 - anchor-free DFL decoder, requires external NMS
268    #[serde(rename = "yolov5")]
269    Yolov5,
270    /// YOLOv8 - anchor-free DFL decoder, requires external NMS
271    #[serde(rename = "yolov8")]
272    Yolov8,
273    /// YOLO11 - anchor-free DFL decoder, requires external NMS
274    #[serde(rename = "yolo11")]
275    Yolo11,
276    /// YOLO26 - end-to-end model with embedded NMS (one-to-one matching
277    /// heads)
278    #[serde(rename = "yolo26")]
279    Yolo26,
280}
281
282impl DecoderVersion {
283    /// Returns true if this version uses end-to-end inference (embedded
284    /// NMS).
285    pub fn is_end_to_end(&self) -> bool {
286        matches!(self, DecoderVersion::Yolo26)
287    }
288}
289
290/// NMS (Non-Maximum Suppression) mode for filtering overlapping detections.
291///
292/// This enum is used with `Option<Nms>`:
293/// - `Some(Nms::Auto)` — resolve from config or fall back to `ClassAgnostic`
294/// - `Some(Nms::ClassAgnostic)` — class-agnostic NMS: suppress overlapping
295///   boxes regardless of class label
296/// - `Some(Nms::ClassAware)` — class-aware NMS: only suppress boxes that
297///   share the same class label AND overlap above the IoU threshold
298/// - `None` — bypass NMS entirely (for end-to-end models with embedded NMS)
299#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Copy, Hash, Eq, Default)]
300#[serde(rename_all = "snake_case")]
301pub enum Nms {
302    /// Let the builder resolve NMS mode from the model config (e.g.
303    /// `edgefirst.json`).  Falls back to [`Nms::ClassAgnostic`] when no
304    /// config specifies a mode.  This is the builder default — callers
305    /// should only use an explicit variant when they need to override
306    /// the config.
307    Auto,
308    /// Suppress overlapping boxes regardless of class label (default
309    /// concrete behavior).
310    #[default]
311    ClassAgnostic,
312    /// Only suppress boxes with the same class label that overlap.
313    ClassAware,
314}
315
316#[derive(Debug, Clone, PartialEq)]
317pub enum ModelType {
318    ModelPackSegDet {
319        boxes: Boxes,
320        scores: Scores,
321        segmentation: Segmentation,
322    },
323    ModelPackSegDetSplit {
324        detection: Vec<Detection>,
325        segmentation: Segmentation,
326    },
327    ModelPackDet {
328        boxes: Boxes,
329        scores: Scores,
330    },
331    ModelPackDetSplit {
332        detection: Vec<Detection>,
333    },
334    ModelPackSeg {
335        segmentation: Segmentation,
336    },
337    YoloDet {
338        boxes: Detection,
339    },
340    YoloSegDet {
341        boxes: Detection,
342        protos: Protos,
343    },
344    YoloSplitDet {
345        boxes: Boxes,
346        scores: Scores,
347    },
348    YoloSplitSegDet {
349        boxes: Boxes,
350        scores: Scores,
351        mask_coeff: MaskCoefficients,
352        protos: Protos,
353    },
354    /// 2-way split YOLO segmentation detection.
355    /// Combined detection tensor (boxes + scores) with separate mask
356    /// coefficients and prototype masks.
357    /// - detection: [1, nc+4, N] — boxes and scores combined
358    /// - mask_coeff: [1, 32, N] — mask coefficients (separate tensor)
359    /// - protos: [1, H/4, W/4, 32] — prototype masks
360    YoloSegDet2Way {
361        boxes: Detection,
362        mask_coeff: MaskCoefficients,
363        protos: Protos,
364    },
365    /// End-to-end YOLO detection (post-NMS output from model)
366    /// Input shape: (1, N, 6+) where columns are [x1, y1, x2, y2, conf,
367    /// class, ...]
368    YoloEndToEndDet {
369        boxes: Detection,
370    },
371    /// End-to-end YOLO detection + segmentation (post-NMS output from
372    /// model) Input shape: (1, N, 6 + num_protos) where columns are
373    /// [x1, y1, x2, y2, conf, class, mask_coeff_0, ..., mask_coeff_31]
374    YoloEndToEndSegDet {
375        boxes: Detection,
376        protos: Protos,
377    },
378    /// Split end-to-end YOLO detection (onnx2tf splits `[1,N,6]` into 3
379    /// tensors) boxes: [batch, N, 4] xyxy, scores: [batch, N, 1],
380    /// classes: [batch, N, 1]
381    YoloSplitEndToEndDet {
382        boxes: Boxes,
383        scores: Scores,
384        classes: Classes,
385    },
386    /// Split end-to-end YOLO seg detection (onnx2tf splits into 5
387    /// tensors)
388    YoloSplitEndToEndSegDet {
389        boxes: Boxes,
390        scores: Scores,
391        classes: Classes,
392        mask_coeff: MaskCoefficients,
393        protos: Protos,
394    },
395    /// Per-scale (physical-output-decomposition) YOLO model. The
396    /// per-scale subsystem (`crates/decoder/src/per_scale/`) owns
397    /// model decoding entirely; this variant exists as a marker so the
398    /// `Decoder::model_type` field has a sensible value for per-scale
399    /// Decoders that bypass the legacy `ModelType`-driven dispatch.
400    PerScale,
401}
edgefirst_decoder/decoder/configs.rs

edgefirst_decoder/decoder/
configs.rs
