Skip to main content

edgefirst_decoder/decoder/
builder.rs

1// SPDX-FileCopyrightText: Copyright 2025 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::HashSet;
5
6use super::config::ConfigOutputRef;
7use super::configs::{self, DecoderType, DecoderVersion, DimName, ModelType};
8use super::merge::DecodeProgram;
9use super::{ConfigOutput, ConfigOutputs, Decoder};
10use crate::per_scale::{DecodeDtype, PerScalePlan};
11use crate::schema::SchemaV2;
12use crate::DecoderError;
13
14/// Extract `(width, height)` from a schema [`crate::schema::InputSpec`].
15///
16/// Prefers named dims (`DimName::Width` / `DimName::Height`) when the
17/// `dshape` is populated, falling back to the NHWC convention
18/// (`shape[1] = H, shape[2] = W`) for 4-element shapes whenever either
19/// named dim is missing. Returns `None` for any other shape arity — the
20/// decoder treats that as "input dims unknown" and skips the EDGEAI-1303
21/// normalization path.
22///
23/// The fallback fires when **either** `Height` or `Width` is missing from
24/// the dshape (not only when both are absent), so a partially-named
25/// dshape (e.g. only `Width`) still resolves both dims via positional
26/// inference instead of silently disabling normalization.
27fn input_dims_from_spec(input: &crate::schema::InputSpec) -> Option<(usize, usize)> {
28    use crate::configs::DimName;
29    let mut h = None;
30    let mut w = None;
31    // `SchemaV2::validate()` doesn't currently enforce
32    // `dshape.len() <= shape.len()`, so a malformed schema can trip an
33    // out-of-bounds index here. Use `shape.get(i)` to silently skip
34    // dshape entries that overflow the shape — the caller treats
35    // missing dims as "unknown" and disables EDGEAI-1303 normalization.
36    for (i, (name, _)) in input.dshape.iter().enumerate() {
37        match name {
38            DimName::Height => h = input.shape.get(i).copied(),
39            DimName::Width => w = input.shape.get(i).copied(),
40            _ => {}
41        }
42    }
43    if (h.is_none() || w.is_none()) && input.shape.len() == 4 {
44        // NHWC default: [N, H, W, C]. Mirrors the per-scale `extract_hw`
45        // fallback (`crates/decoder/src/per_scale/plan.rs::extract_hw`).
46        // Only fill the missing axis so a partial named dshape still
47        // resolves both dims.
48        h = h.or(Some(input.shape[1]));
49        w = w.or(Some(input.shape[2]));
50    }
51    match (w, h) {
52        (Some(w), Some(h)) => Some((w, h)),
53        _ => None,
54    }
55}
56
57#[cfg(test)]
58mod input_dims_from_spec_tests {
59    use super::input_dims_from_spec;
60    use crate::configs::DimName;
61    use crate::schema::InputSpec;
62
63    #[test]
64    fn named_dshape_resolves_dims() {
65        let spec = InputSpec {
66            shape: vec![1, 480, 640, 3],
67            dshape: vec![
68                (DimName::Batch, 1),
69                (DimName::Height, 480),
70                (DimName::Width, 640),
71                (DimName::NumFeatures, 3),
72            ],
73            cameraadaptor: None,
74        };
75        assert_eq!(input_dims_from_spec(&spec), Some((640, 480)));
76    }
77
78    #[test]
79    fn empty_dshape_falls_back_to_nhwc_for_4d() {
80        let spec = InputSpec {
81            shape: vec![1, 480, 640, 3],
82            dshape: vec![],
83            cameraadaptor: None,
84        };
85        assert_eq!(input_dims_from_spec(&spec), Some((640, 480)));
86    }
87
88    #[test]
89    fn malformed_dshape_longer_than_shape_does_not_panic() {
90        // Regression for Copilot review on PR #63: indexing
91        // `input.shape[i]` while iterating dshape can OOB-panic when
92        // `dshape.len() > shape.len()`. The fix uses `shape.get(i)`
93        // and silently treats overflow as "dim missing".
94        let spec = InputSpec {
95            shape: vec![640, 480], // 2-D shape
96            dshape: vec![
97                (DimName::Width, 640),
98                (DimName::Height, 480),
99                (DimName::NumFeatures, 3), // overflow — index 2 ≥ shape.len()
100            ],
101            cameraadaptor: None,
102        };
103        // First two dshape entries resolve via `shape.get()`; the third
104        // is a no-op. Width/Height both resolved, so we expect Some.
105        assert_eq!(input_dims_from_spec(&spec), Some((640, 480)));
106    }
107
108    #[test]
109    fn malformed_dshape_only_overflow_returns_none() {
110        // All dshape entries are past the shape boundary — width and
111        // height stay None and the 4-D NHWC fallback doesn't fire
112        // (shape.len() == 1), so we get None instead of a panic.
113        let spec = InputSpec {
114            shape: vec![1],
115            dshape: vec![
116                (DimName::NumFeatures, 3),
117                (DimName::Height, 480),
118                (DimName::Width, 640),
119            ],
120            cameraadaptor: None,
121        };
122        assert_eq!(input_dims_from_spec(&spec), None);
123    }
124}
125
126/// Default [`DecoderBuilder::with_pre_nms_top_k`] cap for argmax decode.
127pub const DEFAULT_PRE_NMS_TOP_K: usize = 300;
128
129/// Default [`DecoderBuilder::with_pre_nms_top_k`] cap when multi-label decode
130/// is on: Ultralytics' `max_nms`, since multi-label emits up to
131/// anchors × classes candidates.
132pub const MULTI_LABEL_PRE_NMS_TOP_K: usize = 30_000;
133
134#[derive(Debug, Clone, PartialEq)]
135pub struct DecoderBuilder {
136    config_src: Option<ConfigSource>,
137    iou_threshold: f32,
138    score_threshold: f32,
139    /// NMS mode.
140    ///
141    /// - `Some(Nms::Auto)` — resolve from config or fallback to
142    ///   `ClassAware` (builder default)
143    /// - `Some(Nms::ClassAgnostic)` — explicit class-agnostic override
144    /// - `Some(Nms::ClassAware)` — explicit class-aware override
145    /// - `None` — bypass NMS entirely
146    nms: Option<configs::Nms>,
147    /// Output dtype for the per-scale fast path. Has no effect on
148    /// schemas without per-scale children (which use the legacy decode
149    /// path).
150    decode_dtype: DecodeDtype,
151    /// Explicit pre-NMS cap; `None` takes [`DEFAULT_PRE_NMS_TOP_K`], or
152    /// [`MULTI_LABEL_PRE_NMS_TOP_K`] when multi-label decode is on.
153    pre_nms_top_k: Option<usize>,
154    max_det: usize,
155    /// Explicit override for the model input dimensions `(width, height)`,
156    /// consumed by EDGEAI-1303 normalization. When set, takes precedence
157    /// over schema-derived dims; when `None`, the value is read from the
158    /// schema's `input.shape` / `input.dshape` at build time.
159    input_dims: Option<(usize, usize)>,
160    /// Explicit multi-label override; `None` defers to the config's
161    /// `nms_multi_label`, then to off.
162    multi_label: Option<bool>,
163}
164
165#[derive(Debug, Clone, PartialEq)]
166enum ConfigSource {
167    Yaml(String),
168    Json(String),
169    Config(ConfigOutputs),
170    /// Schema v2 metadata. During build the schema is either converted
171    /// to a legacy [`ConfigOutputs`] (flat case) or used to construct a
172    /// [`DecodeProgram`] that performs per-child dequant + merge at
173    /// decode time.
174    Schema(SchemaV2),
175}
176
177impl Default for DecoderBuilder {
178    /// Creates a default DecoderBuilder with no configuration and 0.5 score
179    /// threshold and 0.5 IoU threshold.
180    ///
181    /// A valid configuration must be provided before building the Decoder.
182    ///
183    /// # Examples
184    /// ```rust
185    /// # use edgefirst_decoder::{DecoderBuilder, DecoderResult};
186    /// # fn main() -> DecoderResult<()> {
187    /// #  let config_yaml = edgefirst_bench::testdata::read_to_string("modelpack_split.yaml").to_string();
188    /// let decoder = DecoderBuilder::default()
189    ///     .with_config_yaml_str(config_yaml)
190    ///     .build()?;
191    /// assert_eq!(decoder.score_threshold, 0.5);
192    /// assert_eq!(decoder.iou_threshold, 0.5);
193    ///
194    /// # Ok(())
195    /// # }
196    /// ```
197    fn default() -> Self {
198        Self {
199            config_src: None,
200            iou_threshold: 0.5,
201            score_threshold: 0.5,
202            nms: Some(configs::Nms::Auto),
203            decode_dtype: DecodeDtype::F32,
204            pre_nms_top_k: None,
205            max_det: 300,
206            input_dims: None,
207            multi_label: None,
208        }
209    }
210}
211
212impl DecoderBuilder {
213    /// Creates a default DecoderBuilder with no configuration and 0.5 score
214    /// threshold and 0.5 IoU threshold.
215    ///
216    /// A valid configuration must be provided before building the Decoder.
217    ///
218    /// # Examples
219    /// ```rust
220    /// # use edgefirst_decoder::{DecoderBuilder, DecoderResult};
221    /// # fn main() -> DecoderResult<()> {
222    /// #  let config_yaml = edgefirst_bench::testdata::read_to_string("modelpack_split.yaml").to_string();
223    /// let decoder = DecoderBuilder::new()
224    ///     .with_config_yaml_str(config_yaml)
225    ///     .build()?;
226    /// assert_eq!(decoder.score_threshold, 0.5);
227    /// assert_eq!(decoder.iou_threshold, 0.5);
228    ///
229    /// # Ok(())
230    /// # }
231    /// ```
232    pub fn new() -> Self {
233        Self::default()
234    }
235
236    /// Loads a model configuration in YAML format. Does not check if the string
237    /// is a correct configuration file. Use `DecoderBuilder.build()` to
238    /// deserialize the YAML and parse the model configuration.
239    ///
240    /// # Examples
241    /// ```rust,no_run
242    /// # use edgefirst_decoder::DecoderBuilder;
243    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
244    /// let config_yaml = std::fs::read_to_string("modelpack_split.yaml")?;
245    /// let decoder = DecoderBuilder::new()
246    ///     .with_config_yaml_str(config_yaml)
247    ///     .build()?;
248    ///
249    /// # Ok(())
250    /// # }
251    /// ```
252    pub fn with_config_yaml_str(mut self, yaml_str: String) -> Self {
253        self.config_src.replace(ConfigSource::Yaml(yaml_str));
254        self
255    }
256
257    /// Loads a model configuration in JSON format. Does not check if the string
258    /// is a correct configuration file. Use `DecoderBuilder.build()` to
259    /// deserialize the JSON and parse the model configuration.
260    ///
261    /// # Examples
262    /// ```rust,no_run
263    /// # use edgefirst_decoder::DecoderBuilder;
264    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
265    /// let config_json = std::fs::read_to_string("modelpack_split.json")?;
266    /// let decoder = DecoderBuilder::new()
267    ///     .with_config_json_str(config_json)
268    ///     .build()?;
269    ///
270    /// # Ok(())
271    /// # }
272    /// ```
273    pub fn with_config_json_str(mut self, json_str: String) -> Self {
274        self.config_src.replace(ConfigSource::Json(json_str));
275        self
276    }
277
278    /// Loads a model configuration. Does not check if the configuration is
279    /// correct. Intended to be used when the user needs control over the
280    /// deserialize of the configuration information. Use
281    /// `DecoderBuilder.build()` to parse the model configuration.
282    ///
283    /// # Examples
284    /// ```rust,no_run
285    /// # use edgefirst_decoder::{DecoderBuilder, ConfigOutputs};
286    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
287    /// let config_json = std::fs::read_to_string("modelpack_split.json")?;
288    /// let config: ConfigOutputs = serde_json::from_str(&config_json)?;
289    /// let decoder = DecoderBuilder::new().with_config(config).build()?;
290    ///
291    /// # Ok(())
292    /// # }
293    /// ```
294    pub fn with_config(mut self, config: ConfigOutputs) -> Self {
295        self.config_src.replace(ConfigSource::Config(config));
296        self
297    }
298
299    /// Configure the decoder from a schema v2 metadata document.
300    ///
301    /// Accepts a [`SchemaV2`] as produced by [`SchemaV2::parse_json`],
302    /// [`SchemaV2::parse_yaml`], [`SchemaV2::parse_file`], or
303    /// constructed programmatically. The builder validates the schema,
304    /// compiles a decode program for any split logical outputs
305    /// (per-scale or channel sub-splits), and downconverts the
306    /// logical-level semantics to the legacy [`ConfigOutputs`]
307    /// representation consumed by the existing decoder dispatch.
308    ///
309    /// # Examples
310    ///
311    /// ```rust,no_run
312    /// use edgefirst_decoder::{DecoderBuilder, DecoderResult};
313    /// use edgefirst_decoder::schema::SchemaV2;
314    ///
315    /// # fn main() -> DecoderResult<()> {
316    /// let schema = SchemaV2::parse_file("model/edgefirst.json")?;
317    /// let decoder = DecoderBuilder::new()
318    ///     .with_schema(schema)
319    ///     .with_score_threshold(0.25)
320    ///     .build()?;
321    /// # Ok(())
322    /// # }
323    /// ```
324    pub fn with_schema(mut self, schema: SchemaV2) -> Self {
325        self.config_src.replace(ConfigSource::Schema(schema));
326        self
327    }
328
329    /// Choose the output dtype for the per-scale decoder pipeline.
330    ///
331    /// Defaults to [`DecodeDtype::F32`]. Has no effect on schemas
332    /// without per-scale children (which use the legacy decode path).
333    /// `F16` saves ~2× memory bandwidth at the cost of 10-bit mantissa
334    /// precision; empirically safe for YOLO-family models.
335    ///
336    /// # Examples
337    ///
338    /// ```rust,no_run
339    /// use edgefirst_decoder::{DecodeDtype, DecoderBuilder, DecoderResult};
340    /// use edgefirst_decoder::schema::SchemaV2;
341    ///
342    /// # fn main() -> DecoderResult<()> {
343    /// let schema = SchemaV2::parse_file("model/edgefirst.json")?;
344    /// let decoder = DecoderBuilder::new()
345    ///     .with_schema(schema)
346    ///     .with_decode_dtype(DecodeDtype::F32)
347    ///     .build()?;
348    /// # Ok(())
349    /// # }
350    /// ```
351    pub fn with_decode_dtype(mut self, dtype: DecodeDtype) -> Self {
352        self.decode_dtype = dtype;
353        self
354    }
355
356    /// Loads a YOLO detection model configuration.  Use
357    /// `DecoderBuilder.build()` to parse the model configuration.
358    ///
359    /// # Examples
360    /// ```rust
361    /// # use edgefirst_decoder::{ DecoderBuilder, DecoderResult, configs };
362    /// # fn main() -> DecoderResult<()> {
363    /// let decoder = DecoderBuilder::new()
364    ///     .with_config_yolo_det(
365    ///         configs::Detection {
366    ///             anchors: None,
367    ///             decoder: configs::DecoderType::Ultralytics,
368    ///             quantization: Some(configs::QuantTuple(0.012345, 26)),
369    ///             shape: vec![1, 84, 8400],
370    ///             dshape: Vec::new(),
371    ///             normalized: Some(true),
372    ///         },
373    ///         None,
374    ///     )
375    ///     .build()?;
376    ///
377    /// # Ok(())
378    /// # }
379    /// ```
380    pub fn with_config_yolo_det(
381        mut self,
382        boxes: configs::Detection,
383        version: Option<DecoderVersion>,
384    ) -> Self {
385        let config = ConfigOutputs {
386            outputs: vec![ConfigOutput::Detection(boxes)],
387            decoder_version: version,
388            ..Default::default()
389        };
390        self.config_src.replace(ConfigSource::Config(config));
391        self
392    }
393
394    /// Loads a YOLO split detection model configuration.  Use
395    /// `DecoderBuilder.build()` to parse the model configuration.
396    ///
397    /// # Examples
398    /// ```rust
399    /// # use edgefirst_decoder::{ DecoderBuilder, DecoderResult, configs };
400    /// # fn main() -> DecoderResult<()> {
401    /// let boxes_config = configs::Boxes {
402    ///     decoder: configs::DecoderType::Ultralytics,
403    ///     quantization: Some(configs::QuantTuple(0.012345, 26)),
404    ///     shape: vec![1, 4, 8400],
405    ///     dshape: Vec::new(),
406    ///     normalized: Some(true),
407    /// };
408    /// let scores_config = configs::Scores {
409    ///     decoder: configs::DecoderType::Ultralytics,
410    ///     quantization: Some(configs::QuantTuple(0.0064123, -31)),
411    ///     shape: vec![1, 80, 8400],
412    ///     dshape: Vec::new(),
413    /// };
414    /// let decoder = DecoderBuilder::new()
415    ///     .with_config_yolo_split_det(boxes_config, scores_config)
416    ///     .build()?;
417    /// # Ok(())
418    /// # }
419    /// ```
420    pub fn with_config_yolo_split_det(
421        mut self,
422        boxes: configs::Boxes,
423        scores: configs::Scores,
424    ) -> Self {
425        let config = ConfigOutputs {
426            outputs: vec![ConfigOutput::Boxes(boxes), ConfigOutput::Scores(scores)],
427            ..Default::default()
428        };
429        self.config_src.replace(ConfigSource::Config(config));
430        self
431    }
432
433    /// Loads a YOLO segmentation model configuration.  Use
434    /// `DecoderBuilder.build()` to parse the model configuration.
435    ///
436    /// # Examples
437    /// ```rust
438    /// # use edgefirst_decoder::{ DecoderBuilder, DecoderResult, configs };
439    /// # fn main() -> DecoderResult<()> {
440    /// let seg_config = configs::Detection {
441    ///     decoder: configs::DecoderType::Ultralytics,
442    ///     quantization: Some(configs::QuantTuple(0.012345, 26)),
443    ///     shape: vec![1, 116, 8400],
444    ///     anchors: None,
445    ///     dshape: Vec::new(),
446    ///     normalized: Some(true),
447    /// };
448    /// let protos_config = configs::Protos {
449    ///     decoder: configs::DecoderType::Ultralytics,
450    ///     quantization: Some(configs::QuantTuple(0.0064123, -31)),
451    ///     shape: vec![1, 160, 160, 32],
452    ///     dshape: Vec::new(),
453    /// };
454    /// let decoder = DecoderBuilder::new()
455    ///     .with_config_yolo_segdet(
456    ///         seg_config,
457    ///         protos_config,
458    ///         Some(configs::DecoderVersion::Yolov8),
459    ///     )
460    ///     .build()?;
461    /// # Ok(())
462    /// # }
463    /// ```
464    pub fn with_config_yolo_segdet(
465        mut self,
466        boxes: configs::Detection,
467        protos: configs::Protos,
468        version: Option<DecoderVersion>,
469    ) -> Self {
470        let config = ConfigOutputs {
471            outputs: vec![ConfigOutput::Detection(boxes), ConfigOutput::Protos(protos)],
472            decoder_version: version,
473            ..Default::default()
474        };
475        self.config_src.replace(ConfigSource::Config(config));
476        self
477    }
478
479    /// Loads a YOLO split segmentation model configuration.  Use
480    /// `DecoderBuilder.build()` to parse the model configuration.
481    ///
482    /// # Examples
483    /// ```rust
484    /// # use edgefirst_decoder::{ DecoderBuilder, DecoderResult, configs };
485    /// # fn main() -> DecoderResult<()> {
486    /// let boxes_config = configs::Boxes {
487    ///     decoder: configs::DecoderType::Ultralytics,
488    ///     quantization: Some(configs::QuantTuple(0.012345, 26)),
489    ///     shape: vec![1, 4, 8400],
490    ///     dshape: Vec::new(),
491    ///     normalized: Some(true),
492    /// };
493    /// let scores_config = configs::Scores {
494    ///     decoder: configs::DecoderType::Ultralytics,
495    ///     quantization: Some(configs::QuantTuple(0.012345, 14)),
496    ///     shape: vec![1, 80, 8400],
497    ///     dshape: Vec::new(),
498    /// };
499    /// let mask_config = configs::MaskCoefficients {
500    ///     decoder: configs::DecoderType::Ultralytics,
501    ///     quantization: Some(configs::QuantTuple(0.0064123, 125)),
502    ///     shape: vec![1, 32, 8400],
503    ///     dshape: Vec::new(),
504    /// };
505    /// let protos_config = configs::Protos {
506    ///     decoder: configs::DecoderType::Ultralytics,
507    ///     quantization: Some(configs::QuantTuple(0.0064123, -31)),
508    ///     shape: vec![1, 160, 160, 32],
509    ///     dshape: Vec::new(),
510    /// };
511    /// let decoder = DecoderBuilder::new()
512    ///     .with_config_yolo_split_segdet(boxes_config, scores_config, mask_config, protos_config)
513    ///     .build()?;
514    /// # Ok(())
515    /// # }
516    /// ```
517    pub fn with_config_yolo_split_segdet(
518        mut self,
519        boxes: configs::Boxes,
520        scores: configs::Scores,
521        mask_coefficients: configs::MaskCoefficients,
522        protos: configs::Protos,
523    ) -> Self {
524        let config = ConfigOutputs {
525            outputs: vec![
526                ConfigOutput::Boxes(boxes),
527                ConfigOutput::Scores(scores),
528                ConfigOutput::MaskCoefficients(mask_coefficients),
529                ConfigOutput::Protos(protos),
530            ],
531            ..Default::default()
532        };
533        self.config_src.replace(ConfigSource::Config(config));
534        self
535    }
536
537    /// Loads a ModelPack detection model configuration.  Use
538    /// `DecoderBuilder.build()` to parse the model configuration.
539    ///
540    /// # Examples
541    /// ```rust
542    /// # use edgefirst_decoder::{ DecoderBuilder, DecoderResult, configs };
543    /// # fn main() -> DecoderResult<()> {
544    /// let boxes_config = configs::Boxes {
545    ///     decoder: configs::DecoderType::ModelPack,
546    ///     quantization: Some(configs::QuantTuple(0.012345, 26)),
547    ///     shape: vec![1, 8400, 1, 4],
548    ///     dshape: Vec::new(),
549    ///     normalized: Some(true),
550    /// };
551    /// let scores_config = configs::Scores {
552    ///     decoder: configs::DecoderType::ModelPack,
553    ///     quantization: Some(configs::QuantTuple(0.0064123, -31)),
554    ///     shape: vec![1, 8400, 3],
555    ///     dshape: Vec::new(),
556    /// };
557    /// let decoder = DecoderBuilder::new()
558    ///     .with_config_modelpack_det(boxes_config, scores_config)
559    ///     .build()?;
560    /// # Ok(())
561    /// # }
562    /// ```
563    pub fn with_config_modelpack_det(
564        mut self,
565        boxes: configs::Boxes,
566        scores: configs::Scores,
567    ) -> Self {
568        let config = ConfigOutputs {
569            outputs: vec![ConfigOutput::Boxes(boxes), ConfigOutput::Scores(scores)],
570            ..Default::default()
571        };
572        self.config_src.replace(ConfigSource::Config(config));
573        self
574    }
575
576    /// Loads a ModelPack split detection model configuration. Use
577    /// `DecoderBuilder.build()` to parse the model configuration.
578    ///
579    /// # Examples
580    /// ```rust
581    /// # use edgefirst_decoder::{ DecoderBuilder, DecoderResult, configs };
582    /// # fn main() -> DecoderResult<()> {
583    /// let config0 = configs::Detection {
584    ///     anchors: Some(vec![
585    ///         [0.13750000298023224, 0.2074074000120163],
586    ///         [0.2541666626930237, 0.21481481194496155],
587    ///         [0.23125000298023224, 0.35185185074806213],
588    ///     ]),
589    ///     decoder: configs::DecoderType::ModelPack,
590    ///     quantization: Some(configs::QuantTuple(0.012345, 26)),
591    ///     shape: vec![1, 17, 30, 18],
592    ///     dshape: Vec::new(),
593    ///     normalized: Some(true),
594    /// };
595    /// let config1 = configs::Detection {
596    ///     anchors: Some(vec![
597    ///         [0.36666667461395264, 0.31481480598449707],
598    ///         [0.38749998807907104, 0.4740740656852722],
599    ///         [0.5333333611488342, 0.644444465637207],
600    ///     ]),
601    ///     decoder: configs::DecoderType::ModelPack,
602    ///     quantization: Some(configs::QuantTuple(0.0064123, -31)),
603    ///     shape: vec![1, 9, 15, 18],
604    ///     dshape: Vec::new(),
605    ///     normalized: Some(true),
606    /// };
607    ///
608    /// let decoder = DecoderBuilder::new()
609    ///     .with_config_modelpack_det_split(vec![config0, config1])
610    ///     .build()?;
611    /// # Ok(())
612    /// # }
613    /// ```
614    pub fn with_config_modelpack_det_split(mut self, boxes: Vec<configs::Detection>) -> Self {
615        let outputs = boxes.into_iter().map(ConfigOutput::Detection).collect();
616        let config = ConfigOutputs {
617            outputs,
618            ..Default::default()
619        };
620        self.config_src.replace(ConfigSource::Config(config));
621        self
622    }
623
624    /// Loads a ModelPack segmentation detection model configuration. Use
625    /// `DecoderBuilder.build()` to parse the model configuration.
626    ///
627    /// # Examples
628    /// ```rust
629    /// # use edgefirst_decoder::{ DecoderBuilder, DecoderResult, configs };
630    /// # fn main() -> DecoderResult<()> {
631    /// let boxes_config = configs::Boxes {
632    ///     decoder: configs::DecoderType::ModelPack,
633    ///     quantization: Some(configs::QuantTuple(0.012345, 26)),
634    ///     shape: vec![1, 8400, 1, 4],
635    ///     dshape: Vec::new(),
636    ///     normalized: Some(true),
637    /// };
638    /// let scores_config = configs::Scores {
639    ///     decoder: configs::DecoderType::ModelPack,
640    ///     quantization: Some(configs::QuantTuple(0.0064123, -31)),
641    ///     shape: vec![1, 8400, 2],
642    ///     dshape: Vec::new(),
643    /// };
644    /// let seg_config = configs::Segmentation {
645    ///     decoder: configs::DecoderType::ModelPack,
646    ///     quantization: Some(configs::QuantTuple(0.0064123, -31)),
647    ///     shape: vec![1, 640, 640, 3],
648    ///     dshape: Vec::new(),
649    /// };
650    /// let decoder = DecoderBuilder::new()
651    ///     .with_config_modelpack_segdet(boxes_config, scores_config, seg_config)
652    ///     .build()?;
653    /// # Ok(())
654    /// # }
655    /// ```
656    pub fn with_config_modelpack_segdet(
657        mut self,
658        boxes: configs::Boxes,
659        scores: configs::Scores,
660        segmentation: configs::Segmentation,
661    ) -> Self {
662        let config = ConfigOutputs {
663            outputs: vec![
664                ConfigOutput::Boxes(boxes),
665                ConfigOutput::Scores(scores),
666                ConfigOutput::Segmentation(segmentation),
667            ],
668            ..Default::default()
669        };
670        self.config_src.replace(ConfigSource::Config(config));
671        self
672    }
673
674    /// Loads a ModelPack segmentation split detection model configuration. Use
675    /// `DecoderBuilder.build()` to parse the model configuration.
676    ///
677    /// # Examples
678    /// ```rust
679    /// # use edgefirst_decoder::{ DecoderBuilder, DecoderResult, configs };
680    /// # fn main() -> DecoderResult<()> {
681    /// let config0 = configs::Detection {
682    ///     anchors: Some(vec![
683    ///         [0.36666667461395264, 0.31481480598449707],
684    ///         [0.38749998807907104, 0.4740740656852722],
685    ///         [0.5333333611488342, 0.644444465637207],
686    ///     ]),
687    ///     decoder: configs::DecoderType::ModelPack,
688    ///     quantization: Some(configs::QuantTuple(0.08547406643629074, 174)),
689    ///     shape: vec![1, 9, 15, 18],
690    ///     dshape: Vec::new(),
691    ///     normalized: Some(true),
692    /// };
693    /// let config1 = configs::Detection {
694    ///     anchors: Some(vec![
695    ///         [0.13750000298023224, 0.2074074000120163],
696    ///         [0.2541666626930237, 0.21481481194496155],
697    ///         [0.23125000298023224, 0.35185185074806213],
698    ///     ]),
699    ///     decoder: configs::DecoderType::ModelPack,
700    ///     quantization: Some(configs::QuantTuple(0.09929127991199493, 183)),
701    ///     shape: vec![1, 17, 30, 18],
702    ///     dshape: Vec::new(),
703    ///     normalized: Some(true),
704    /// };
705    /// let seg_config = configs::Segmentation {
706    ///     decoder: configs::DecoderType::ModelPack,
707    ///     quantization: Some(configs::QuantTuple(0.0064123, -31)),
708    ///     shape: vec![1, 640, 640, 2],
709    ///     dshape: Vec::new(),
710    /// };
711    /// let decoder = DecoderBuilder::new()
712    ///     .with_config_modelpack_segdet_split(vec![config0, config1], seg_config)
713    ///     .build()?;
714    /// # Ok(())
715    /// # }
716    /// ```
717    pub fn with_config_modelpack_segdet_split(
718        mut self,
719        boxes: Vec<configs::Detection>,
720        segmentation: configs::Segmentation,
721    ) -> Self {
722        let mut outputs = boxes
723            .into_iter()
724            .map(ConfigOutput::Detection)
725            .collect::<Vec<_>>();
726        outputs.push(ConfigOutput::Segmentation(segmentation));
727        let config = ConfigOutputs {
728            outputs,
729            ..Default::default()
730        };
731        self.config_src.replace(ConfigSource::Config(config));
732        self
733    }
734
735    /// Loads a ModelPack segmentation model configuration. Use
736    /// `DecoderBuilder.build()` to parse the model configuration.
737    ///
738    /// # Examples
739    /// ```rust
740    /// # use edgefirst_decoder::{ DecoderBuilder, DecoderResult, configs };
741    /// # fn main() -> DecoderResult<()> {
742    /// let seg_config = configs::Segmentation {
743    ///     decoder: configs::DecoderType::ModelPack,
744    ///     quantization: Some(configs::QuantTuple(0.0064123, -31)),
745    ///     shape: vec![1, 640, 640, 3],
746    ///     dshape: Vec::new(),
747    /// };
748    /// let decoder = DecoderBuilder::new()
749    ///     .with_config_modelpack_seg(seg_config)
750    ///     .build()?;
751    /// # Ok(())
752    /// # }
753    /// ```
754    pub fn with_config_modelpack_seg(mut self, segmentation: configs::Segmentation) -> Self {
755        let config = ConfigOutputs {
756            outputs: vec![ConfigOutput::Segmentation(segmentation)],
757            ..Default::default()
758        };
759        self.config_src.replace(ConfigSource::Config(config));
760        self
761    }
762
763    /// Add an output to the decoder configuration.
764    ///
765    /// Incrementally builds the model configuration by adding outputs one at
766    /// a time. The decoder resolves the model type from the combination of
767    /// outputs during `build()`.
768    ///
769    /// If `dshape` is non-empty on the output, `shape` is automatically
770    /// derived from it (the size component of each named dimension). This
771    /// prevents conflicts between `shape` and `dshape`.
772    ///
773    /// This uses the programmatic config path. Calling this after
774    /// `with_config_json_str()` or `with_config_yaml_str()` replaces the
775    /// string-based config source.
776    ///
777    /// # Examples
778    /// ```rust
779    /// # use edgefirst_decoder::{DecoderBuilder, DecoderResult, ConfigOutput, configs};
780    /// # fn main() -> DecoderResult<()> {
781    /// let decoder = DecoderBuilder::new()
782    ///     .add_output(ConfigOutput::Scores(configs::Scores {
783    ///         decoder: configs::DecoderType::Ultralytics,
784    ///         dshape: vec![
785    ///             (configs::DimName::Batch, 1),
786    ///             (configs::DimName::NumClasses, 80),
787    ///             (configs::DimName::NumBoxes, 8400),
788    ///         ],
789    ///         ..Default::default()
790    ///     }))
791    ///     .add_output(ConfigOutput::Boxes(configs::Boxes {
792    ///         decoder: configs::DecoderType::Ultralytics,
793    ///         dshape: vec![
794    ///             (configs::DimName::Batch, 1),
795    ///             (configs::DimName::BoxCoords, 4),
796    ///             (configs::DimName::NumBoxes, 8400),
797    ///         ],
798    ///         ..Default::default()
799    ///     }))
800    ///     .build()?;
801    /// # Ok(())
802    /// # }
803    /// ```
804    pub fn add_output(mut self, output: ConfigOutput) -> Self {
805        if !matches!(self.config_src, Some(ConfigSource::Config(_))) {
806            self.config_src = Some(ConfigSource::Config(ConfigOutputs::default()));
807        }
808        if let Some(ConfigSource::Config(ref mut config)) = self.config_src {
809            config.outputs.push(Self::normalize_output(output));
810        }
811        self
812    }
813
814    /// Sets the decoder version for Ultralytics models.
815    ///
816    /// This is used with `add_output()` to specify the YOLO architecture
817    /// version when it cannot be inferred from the output shapes alone.
818    ///
819    /// - `Yolov5`, `Yolov8`, `Yolo11`: Traditional models requiring external
820    ///   NMS
821    /// - `Yolo26`: End-to-end models with NMS embedded in the model graph
822    ///
823    /// # Examples
824    /// ```rust
825    /// # use edgefirst_decoder::{DecoderBuilder, DecoderResult, ConfigOutput, configs};
826    /// # fn main() -> DecoderResult<()> {
827    /// let decoder = DecoderBuilder::new()
828    ///     .add_output(ConfigOutput::Detection(configs::Detection {
829    ///         decoder: configs::DecoderType::Ultralytics,
830    ///         dshape: vec![
831    ///             (configs::DimName::Batch, 1),
832    ///             (configs::DimName::NumBoxes, 100),
833    ///             (configs::DimName::NumFeatures, 6),
834    ///         ],
835    ///         ..Default::default()
836    ///     }))
837    ///     .with_decoder_version(configs::DecoderVersion::Yolo26)
838    ///     .build()?;
839    /// # Ok(())
840    /// # }
841    /// ```
842    pub fn with_decoder_version(mut self, version: configs::DecoderVersion) -> Self {
843        if !matches!(self.config_src, Some(ConfigSource::Config(_))) {
844            self.config_src = Some(ConfigSource::Config(ConfigOutputs::default()));
845        }
846        if let Some(ConfigSource::Config(ref mut config)) = self.config_src {
847            config.decoder_version = Some(version);
848        }
849        self
850    }
851
852    /// Normalize an output: if dshape is non-empty, derive shape from it.
853    fn normalize_output(mut output: ConfigOutput) -> ConfigOutput {
854        fn normalize_shape(shape: &mut Vec<usize>, dshape: &[(configs::DimName, usize)]) {
855            if !dshape.is_empty() {
856                *shape = dshape.iter().map(|(_, size)| *size).collect();
857            }
858        }
859        match &mut output {
860            ConfigOutput::Detection(c) => normalize_shape(&mut c.shape, &c.dshape),
861            ConfigOutput::Boxes(c) => normalize_shape(&mut c.shape, &c.dshape),
862            ConfigOutput::Scores(c) => normalize_shape(&mut c.shape, &c.dshape),
863            ConfigOutput::Protos(c) => normalize_shape(&mut c.shape, &c.dshape),
864            ConfigOutput::Segmentation(c) => normalize_shape(&mut c.shape, &c.dshape),
865            ConfigOutput::MaskCoefficients(c) => normalize_shape(&mut c.shape, &c.dshape),
866            ConfigOutput::Mask(c) => normalize_shape(&mut c.shape, &c.dshape),
867            ConfigOutput::Classes(c) => normalize_shape(&mut c.shape, &c.dshape),
868        }
869        output
870    }
871
872    /// Sets the scores threshold of the decoder
873    ///
874    /// # Examples
875    /// ```rust
876    /// # use edgefirst_decoder::{DecoderBuilder, DecoderResult};
877    /// # fn main() -> DecoderResult<()> {
878    /// # let config_json = edgefirst_bench::testdata::read_to_string("modelpack_split.json").to_string();
879    /// let decoder = DecoderBuilder::new()
880    ///     .with_config_json_str(config_json)
881    ///     .with_score_threshold(0.654)
882    ///     .build()?;
883    /// assert_eq!(decoder.score_threshold, 0.654);
884    /// # Ok(())
885    /// # }
886    /// ```
887    pub fn with_score_threshold(mut self, score_threshold: f32) -> Self {
888        self.score_threshold = score_threshold;
889        self
890    }
891
892    /// Sets the IOU threshold of the decoder. Has no effect when NMS is set to
893    /// `None`
894    ///
895    /// # Examples
896    /// ```rust
897    /// # use edgefirst_decoder::{DecoderBuilder, DecoderResult};
898    /// # fn main() -> DecoderResult<()> {
899    /// # let config_json = edgefirst_bench::testdata::read_to_string("modelpack_split.json").to_string();
900    /// let decoder = DecoderBuilder::new()
901    ///     .with_config_json_str(config_json)
902    ///     .with_iou_threshold(0.654)
903    ///     .build()?;
904    /// assert_eq!(decoder.iou_threshold, 0.654);
905    /// # Ok(())
906    /// # }
907    /// ```
908    pub fn with_iou_threshold(mut self, iou_threshold: f32) -> Self {
909        self.iou_threshold = iou_threshold;
910        self
911    }
912
913    /// Sets the NMS mode for the decoder.
914    ///
915    /// - `Some(Nms::Auto)` — resolve from model config (e.g. `edgefirst.json`)
916    ///   or fall back to `ClassAware` (this is the builder default)
917    /// - `Some(Nms::ClassAgnostic)` — class-agnostic NMS: suppress overlapping
918    ///   boxes regardless of class label
919    /// - `Some(Nms::ClassAware)` — class-aware NMS: only suppress boxes that
920    ///   share the same class label AND overlap above the IoU threshold
921    /// - `None` — bypass NMS entirely (for end-to-end models with embedded NMS)
922    ///
923    /// # Examples
924    /// ```rust
925    /// # use edgefirst_decoder::{DecoderBuilder, DecoderResult, configs::Nms};
926    /// # fn main() -> DecoderResult<()> {
927    /// # let config_json = edgefirst_bench::testdata::read_to_string("modelpack_split.json").to_string();
928    /// let decoder = DecoderBuilder::new()
929    ///     .with_config_json_str(config_json)
930    ///     .with_nms(Some(Nms::ClassAware))
931    ///     .build()?;
932    /// assert_eq!(decoder.nms, Some(Nms::ClassAware));
933    /// # Ok(())
934    /// # }
935    /// ```
936    pub fn with_nms(mut self, nms: Option<configs::Nms>) -> Self {
937        self.nms = nms;
938        self
939    }
940
941    /// Enable multi-label box decode for validation runs.
942    ///
943    /// When `true`, the decoder emits one candidate per `(anchor, class)` pair
944    /// for every class whose score meets the threshold — matching the
945    /// Ultralytics `val` multi-label decode that drives COCO mAP evaluation.
946    ///
947    /// **Default:** the model config's `nms_multi_label`, else `false`
948    /// (argmax: one class per anchor). An explicit call overrides the config
949    /// either way; `build()` logs a warning when the config turns it on.
950    /// [`Decoder::decode_for_tracking`](crate::Decoder::decode_for_tracking),
951    /// and `Decoder::decode_tracked` / `Decoder::decode_proto_tracked` (behind
952    /// the `tracker` feature), ignore it and always decode one label per box.
953    ///
954    /// The candidates go through the decoder's NMS mode unchanged. Class-aware
955    /// NMS (the default when neither the caller nor the config sets a mode)
956    /// keeps every class of an anchor; an explicit `Nms::ClassAgnostic`
957    /// suppresses across classes, as Ultralytics
958    /// `non_max_suppression(multi_label=True, agnostic=True)` does.
959    ///
960    /// Unless [`with_pre_nms_top_k`](Self::with_pre_nms_top_k) is called,
961    /// multi-label decode raises the pre-NMS cap to
962    /// [`MULTI_LABEL_PRE_NMS_TOP_K`].
963    ///
964    /// # Examples
965    /// ```rust
966    /// # use edgefirst_decoder::{DecoderBuilder, DecoderResult};
967    /// # fn main() -> DecoderResult<()> {
968    /// # let config_json = edgefirst_bench::testdata::read_to_string("modelpack_split.json").to_string();
969    /// let decoder = DecoderBuilder::new()
970    ///     .with_config_json_str(config_json)
971    ///     .with_multi_label(true)
972    ///     .build()?;
973    /// # Ok(())
974    /// # }
975    /// ```
976    pub fn with_multi_label(mut self, v: bool) -> Self {
977        self.multi_label = Some(v);
978        self
979    }
980
981    /// Sets the maximum number of candidate boxes fed into NMS after score
982    /// filtering; `0` means unbounded. Uses partial sort (O(N)) to select the
983    /// top-K candidates, dramatically reducing the O(N²) NMS cost when many
984    /// low-confidence proposals pass the threshold (common with mAP eval at
985    /// 0.001). Equal scores at the cut keep the lower anchor index. Applies
986    /// to every NMS decode path; ignored when NMS is bypassed.
987    ///
988    /// Default: [`DEFAULT_PRE_NMS_TOP_K`] (300), or
989    /// [`MULTI_LABEL_PRE_NMS_TOP_K`] (30 000, Ultralytics' `max_nms`) when
990    /// multi-label decode is on, whether from
991    /// [`with_multi_label`](Self::with_multi_label) or the model's
992    /// `nms_multi_label`; argmax decodes of a multi-label decoder (the
993    /// tracking entry points) keep 300. Unset, the built decoder's
994    /// [`Decoder::pre_nms_top_k`](crate::Decoder::pre_nms_top_k) is `None` and
995    /// is resolved per decode; an explicit call stores `Some(n)`, which every
996    /// decode uses as given.
997    ///
998    /// # ⚠️ Validation vs Deployment
999    ///
1000    /// The default is appropriate for **deployment** where
1001    /// `score_threshold ≥ 0.25` means few anchors survive filtering and
1002    /// top-K is effectively a no-op.
1003    ///
1004    /// For **COCO mAP evaluation** (`score_threshold ≈ 0.001`), set this to
1005    /// `0` (unbounded) so that all score-passing candidates reach NMS.
1006    /// Under argmax decode there is one candidate per anchor, so the anchor
1007    /// count (8 400 for a 640 × 640 YOLO model) is also enough; under
1008    /// multi-label decode there are up to anchors × classes candidates, so use
1009    /// `0` or leave the multi-label default. Failing to raise the cap causes
1010    /// **~9 pp box mAP loss** — the decoder math is correct but the
1011    /// evaluation protocol requires full recall across the confidence range.
1012    ///
1013    /// Post-processing latency scales with candidate count. At deployment
1014    /// thresholds the cost difference is negligible; at validation thresholds
1015    /// it is measurable but necessary for correct results.
1016    ///
1017    /// # Examples
1018    ///
1019    /// Deployment (default top-K, high score threshold):
1020    /// ```rust
1021    /// # use edgefirst_decoder::{DecoderBuilder, DecoderResult};
1022    /// # fn main() -> DecoderResult<()> {
1023    /// # let config_json = edgefirst_bench::testdata::read_to_string("modelpack_split.json").to_string();
1024    /// let decoder = DecoderBuilder::new()
1025    ///     .with_config_json_str(config_json)
1026    ///     .with_score_threshold(0.25)
1027    ///     // pre_nms_top_k defaults to 300 — appropriate here
1028    ///     .build()?;
1029    /// assert_eq!(decoder.pre_nms_top_k, None);
1030    /// assert_eq!(decoder.pre_nms_top_k_for(false), 300);
1031    /// # Ok(())
1032    /// # }
1033    /// ```
1034    ///
1035    /// COCO mAP evaluation (pass every candidate to NMS):
1036    /// ```rust
1037    /// # use edgefirst_decoder::{DecoderBuilder, DecoderResult};
1038    /// # fn main() -> DecoderResult<()> {
1039    /// # let config_json = edgefirst_bench::testdata::read_to_string("modelpack_split.json").to_string();
1040    /// let decoder = DecoderBuilder::new()
1041    ///     .with_config_json_str(config_json)
1042    ///     .with_score_threshold(0.001)
1043    ///     .with_pre_nms_top_k(0) // unbounded
1044    ///     .with_max_det(300)
1045    ///     .build()?;
1046    /// assert_eq!(decoder.pre_nms_top_k, Some(0));
1047    /// # Ok(())
1048    /// # }
1049    /// ```
1050    pub fn with_pre_nms_top_k(mut self, pre_nms_top_k: usize) -> Self {
1051        self.pre_nms_top_k = Some(pre_nms_top_k);
1052        self
1053    }
1054
1055    /// Sets the maximum number of detections returned after NMS.
1056    /// Matches the Ultralytics `max_det` parameter.
1057    ///
1058    /// Default: 300.
1059    ///
1060    /// # Examples
1061    /// ```rust
1062    /// # use edgefirst_decoder::{DecoderBuilder, DecoderResult};
1063    /// # fn main() -> DecoderResult<()> {
1064    /// # let config_json = edgefirst_bench::testdata::read_to_string("modelpack_split.json").to_string();
1065    /// let decoder = DecoderBuilder::new()
1066    ///     .with_config_json_str(config_json)
1067    ///     .with_max_det(100)
1068    ///     .build()?;
1069    /// assert_eq!(decoder.max_det, 100);
1070    /// # Ok(())
1071    /// # }
1072    /// ```
1073    pub fn with_max_det(mut self, max_det: usize) -> Self {
1074        self.max_det = max_det;
1075        self
1076    }
1077
1078    /// Sets the model input dimensions `(width, height)` consumed by the
1079    /// EDGEAI-1303 normalization path. Use this when building via
1080    /// [`with_config`](Self::with_config) / [`add_output`](Self::add_output)
1081    /// (no schema) and the model emits pixel-space boxes that need to be
1082    /// divided by `(W, H)` before NMS.
1083    ///
1084    /// When the builder is also configured with [`with_schema`](Self::with_schema)
1085    /// (or `with_config_json_str` / `with_config_yaml_str`) and the schema's
1086    /// `input` block carries usable dims, this explicit override **takes
1087    /// precedence** so callers can correct schemas with missing or wrong
1088    /// input specs without rewriting the schema.
1089    ///
1090    /// # Examples
1091    /// ```rust
1092    /// # use edgefirst_decoder::{DecoderBuilder, DecoderResult};
1093    /// # fn main() -> DecoderResult<()> {
1094    /// # let config_yaml = edgefirst_bench::testdata::read_to_string("modelpack_split.yaml").to_string();
1095    /// let decoder = DecoderBuilder::new()
1096    ///     .with_config_yaml_str(config_yaml)
1097    ///     .with_input_dims(640, 640)
1098    ///     .build()?;
1099    /// assert_eq!(decoder.input_dims(), Some((640, 640)));
1100    /// # Ok(())
1101    /// # }
1102    /// ```
1103    pub fn with_input_dims(mut self, width: usize, height: usize) -> Self {
1104        self.input_dims = Some((width, height));
1105        self
1106    }
1107
1108    /// Builds the decoder with the given settings. If the config is a JSON or
1109    /// YAML string, this will deserialize the JSON or YAML and then parse the
1110    /// configuration information.
1111    ///
1112    /// # Examples
1113    /// ```rust
1114    /// # use edgefirst_decoder::{DecoderBuilder, DecoderResult};
1115    /// # fn main() -> DecoderResult<()> {
1116    /// # let config_json = edgefirst_bench::testdata::read_to_string("modelpack_split.json").to_string();
1117    /// let decoder = DecoderBuilder::new()
1118    ///     .with_config_json_str(config_json)
1119    ///     .with_score_threshold(0.654)
1120    ///     .build()?;
1121    /// # Ok(())
1122    /// # }
1123    /// ```
1124    pub fn build(self) -> Result<Decoder, DecoderError> {
1125        let decode_dtype = self.decode_dtype;
1126        let explicit_input_dims = self.input_dims;
1127        let (config, decode_program, per_scale_plan, schema_input_dims) = match self.config_src {
1128            Some(ConfigSource::Json(s)) => {
1129                Self::build_from_schema(SchemaV2::parse_json(&s)?, decode_dtype)?
1130            }
1131            Some(ConfigSource::Yaml(s)) => {
1132                Self::build_from_schema(SchemaV2::parse_yaml(&s)?, decode_dtype)?
1133            }
1134            Some(ConfigSource::Config(c)) => (c, None, None, None),
1135            Some(ConfigSource::Schema(schema)) => Self::build_from_schema(schema, decode_dtype)?,
1136            None => return Err(DecoderError::NoConfig),
1137        };
1138        // Explicit `with_input_dims(W, H)` overrides any schema-derived
1139        // value so callers can fix schemas with missing or wrong input
1140        // specs without rewriting the schema (EDGEAI-1303).
1141        let input_dims = explicit_input_dims.or(schema_input_dims);
1142
1143        // Enforce the physical-order contract: when dshape is present
1144        // it must describe the same axes as shape in the same order,
1145        // listed from outermost to innermost. Ambiguous-layout roles
1146        // (Protos, Boxes, Scores, MaskCoefficients, Classes, Detection)
1147        // may still omit dshape when shape is already in the decoder's
1148        // canonical order.
1149        for output in &config.outputs {
1150            Decoder::validate_output_layout(output.into())?;
1151        }
1152
1153        // Extract normalized flag from config outputs.
1154        //
1155        // The per-scale subsystem (DFL/LTRB → dist2bbox → sigmoid) emits
1156        // boxes in pixel coordinates by design — `(grid + dist) * stride`
1157        // — independently of any `normalized: true` annotation in the
1158        // schema. Override to `Some(false)` so the per-scale bridge's
1159        // call to `yolo::maybe_normalize_boxes_in_place` fires and
1160        // divides by `input_dims`, yielding `[0, 1]` output. The
1161        // accessor `Decoder::normalized_boxes()` applies the
1162        // pixel→normalized upgrade for the per-scale path and for any
1163        // legacy `ModelType` whose every entry point normalizes
1164        // uniformly (currently `YoloSegDet`, `YoloSplitSegDet`, and
1165        // `YoloSegDet2Way`); other paths surface the raw flag.
1166        let normalized = if per_scale_plan.is_some() {
1167            Some(false)
1168        } else {
1169            Self::get_normalized(&config.outputs)
1170        };
1171
1172        let (multi_label, multi_label_source) = match (self.multi_label, config.nms_multi_label) {
1173            (Some(v), _) => (v, super::MultiLabelSource::Explicit),
1174            (None, Some(v)) => (v, super::MultiLabelSource::Metadata),
1175            (None, None) => (false, super::MultiLabelSource::Default),
1176        };
1177        if multi_label && multi_label_source == super::MultiLabelSource::Metadata {
1178            tracing::warn!(
1179                "model metadata sets nms_multi_label: decode emits one box per class above the \
1180                 score threshold (validation decode); pass multi_label=false to override"
1181            );
1182        }
1183        let pre_nms_top_k = self.pre_nms_top_k;
1184
1185        // NMS precedence:
1186        //   Some(ClassAgnostic|ClassAware) → explicit user override
1187        //   Some(Auto) → resolve from config, fallback to ClassAware
1188        //   None → NMS disabled (explicit)
1189        //
1190        // `Auto` is always resolved to a concrete mode here — it never
1191        // persists into the built `Decoder`, even if the config itself
1192        // contains `Auto`.
1193        let resolve_auto = |nms: Option<configs::Nms>| match nms {
1194            Some(configs::Nms::Auto) | None => Some(configs::Nms::ClassAware),
1195            concrete => concrete,
1196        };
1197        let nms = match self.nms {
1198            Some(configs::Nms::Auto) => resolve_auto(config.nms),
1199            other => other,
1200        };
1201        // When the per-scale path is active, the per_scale subsystem owns
1202        // model decoding entirely — `decode` / `decode_proto` short-circuit
1203        // on `per_scale.is_some()` before reading `model_type`. Skip the
1204        // legacy ModelType validation, which otherwise rejects per-scale
1205        // schemas that carry `decoder_version: yolo26` (an
1206        // "end-to-end" hint) but use the per-scale split outputs rather
1207        // than the end-to-end split-output shape the legacy validator
1208        // expects. We keep a placeholder `ModelType` so the field remains
1209        // valid; it is dead state for per-scale Decoders.
1210        let model_type = if per_scale_plan.is_some() {
1211            // Drop the un-needed config; the per-scale subsystem owns it.
1212            drop(config);
1213            ModelType::PerScale
1214        } else {
1215            Self::get_model_type(config)?
1216        };
1217
1218        let per_scale = per_scale_plan
1219            .map(|plan| std::sync::Mutex::new(crate::per_scale::PerScaleDecoder::new(plan)));
1220
1221        debug_assert!(
1222            !matches!(nms, Some(configs::Nms::Auto)),
1223            "Nms::Auto must be resolved to a concrete mode before building Decoder"
1224        );
1225
1226        Ok(Decoder {
1227            model_type,
1228            iou_threshold: self.iou_threshold,
1229            score_threshold: self.score_threshold,
1230            nms,
1231            pre_nms_top_k,
1232            max_det: self.max_det,
1233            normalized,
1234            input_dims,
1235            multi_label,
1236            multi_label_source,
1237            tracked_multi_label_warned: std::sync::atomic::AtomicBool::new(false),
1238            decode_program,
1239            per_scale,
1240        })
1241    }
1242
1243    /// Validate a [`SchemaV2`] and lower it to the (legacy `ConfigOutputs`,
1244    /// optional `DecodeProgram`, optional `PerScalePlan`) tuple the rest
1245    /// of `build()` consumes.
1246    ///
1247    /// Centralises the v2 lowering so JSON, YAML, and direct
1248    /// `with_schema` callers all go through the same validation,
1249    /// merge-program, and per-scale plan construction. `SchemaV2::parse_json`
1250    /// / `parse_yaml` already auto-detect v1 vs v2 input and return a v2
1251    /// schema either way (v1 inputs are upgraded in memory via
1252    /// `from_v1`), so this helper is the sole place that turns a
1253    /// schema into builder-ready state.
1254    #[allow(clippy::type_complexity)]
1255    fn build_from_schema(
1256        schema: SchemaV2,
1257        decode_dtype: DecodeDtype,
1258    ) -> Result<
1259        (
1260            ConfigOutputs,
1261            Option<DecodeProgram>,
1262            Option<PerScalePlan>,
1263            Option<(usize, usize)>,
1264        ),
1265        DecoderError,
1266    > {
1267        schema.validate()?;
1268        // The per-scale subsystem claims per-scale schemas in full and owns
1269        // their decode end-to-end (`decode` / `decode_proto` short-circuit on
1270        // `per_scale.is_some()`). Build it first and use its claim as the
1271        // single source of truth: only fall back to the schema-v2 merge
1272        // program for split schemas it does NOT claim (e.g. ARA-2 channel
1273        // sub-splits). This keeps `decode_program` `None` for per-scale
1274        // schemas so the merge path never sees per-scale logicals.
1275        let per_scale = PerScalePlan::try_from_schema(&schema, decode_dtype)?;
1276        let program = if per_scale.is_some() {
1277            None
1278        } else {
1279            DecodeProgram::try_from_schema(&schema)?
1280        };
1281        // Extract model input (W, H) from `input.shape`/`dshape`. Used by
1282        // the legacy decode path to honour `normalized: false` (see
1283        // EDGEAI-1303). `None` is fine when the schema omits the input
1284        // spec — the decoder falls back to the protobox `>2.0` reject.
1285        let input_dims = schema.input.as_ref().and_then(input_dims_from_spec);
1286        let legacy = schema.to_legacy_config_outputs()?;
1287        Ok((legacy, program, per_scale, input_dims))
1288    }
1289
1290    /// Extracts the normalized flag from config outputs.
1291    /// - `Some(true)`: Boxes are in normalized [0,1] coordinates
1292    /// - `Some(false)`: Boxes are in pixel coordinates
1293    /// - `None`: Unknown (not specified in config), caller must infer
1294    fn get_normalized(outputs: &[ConfigOutput]) -> Option<bool> {
1295        for output in outputs {
1296            match output {
1297                ConfigOutput::Detection(det) => return det.normalized,
1298                ConfigOutput::Boxes(boxes) => return boxes.normalized,
1299                _ => {}
1300            }
1301        }
1302        None // not specified
1303    }
1304
1305    fn get_model_type(configs: ConfigOutputs) -> Result<ModelType, DecoderError> {
1306        // yolo or modelpack
1307        let mut yolo = false;
1308        let mut modelpack = false;
1309        for c in &configs.outputs {
1310            match c.decoder() {
1311                DecoderType::ModelPack => modelpack = true,
1312                DecoderType::Ultralytics => yolo = true,
1313            }
1314        }
1315        match (modelpack, yolo) {
1316            (true, true) => Err(DecoderError::InvalidConfig(
1317                "Both ModelPack and Yolo outputs found in config".to_string(),
1318            )),
1319            (true, false) => Self::get_model_type_modelpack(configs),
1320            (false, true) => Self::get_model_type_yolo(configs),
1321            (false, false) => Err(DecoderError::InvalidConfig(
1322                "No outputs found in config".to_string(),
1323            )),
1324        }
1325    }
1326
1327    fn get_model_type_yolo(configs: ConfigOutputs) -> Result<ModelType, DecoderError> {
1328        let mut boxes = None;
1329        let mut protos = None;
1330        let mut split_boxes = None;
1331        let mut split_scores = None;
1332        let mut split_mask_coeff = None;
1333        let mut split_classes = None;
1334        for c in configs.outputs {
1335            match c {
1336                ConfigOutput::Detection(detection) => boxes = Some(detection),
1337                ConfigOutput::Segmentation(_) => {
1338                    return Err(DecoderError::InvalidConfig(
1339                        "Invalid Segmentation output with Yolo decoder".to_string(),
1340                    ));
1341                }
1342                ConfigOutput::Protos(protos_) => protos = Some(protos_),
1343                ConfigOutput::Mask(_) => {
1344                    return Err(DecoderError::InvalidConfig(
1345                        "Invalid Mask output with Yolo decoder".to_string(),
1346                    ));
1347                }
1348                ConfigOutput::Scores(scores) => split_scores = Some(scores),
1349                ConfigOutput::Boxes(boxes) => split_boxes = Some(boxes),
1350                ConfigOutput::MaskCoefficients(mask_coeff) => split_mask_coeff = Some(mask_coeff),
1351                ConfigOutput::Classes(classes) => split_classes = Some(classes),
1352            }
1353        }
1354
1355        // Use end-to-end model types when:
1356        // 1. decoder_version is explicitly set to Yolo26 (definitive), OR
1357        //    decoder_version is not set but the dshapes are (batch, num_boxes,
1358        //    num_features)
1359        let is_end_to_end_dshape = boxes.as_ref().is_some_and(|b| {
1360            let dims = b.dshape.iter().map(|(d, _)| *d).collect::<Vec<_>>();
1361            dims == vec![DimName::Batch, DimName::NumBoxes, DimName::NumFeatures]
1362        });
1363
1364        let is_end_to_end = configs
1365            .decoder_version
1366            .map(|v| v.is_end_to_end())
1367            .unwrap_or(is_end_to_end_dshape);
1368
1369        if is_end_to_end {
1370            if let Some(boxes) = boxes {
1371                if let Some(protos) = protos {
1372                    Self::verify_yolo_seg_det_26(&boxes, &protos)?;
1373                    return Ok(ModelType::YoloEndToEndSegDet { boxes, protos });
1374                } else {
1375                    Self::verify_yolo_det_26(&boxes)?;
1376                    return Ok(ModelType::YoloEndToEndDet { boxes });
1377                }
1378            } else if let (Some(split_boxes), Some(split_scores), Some(split_classes)) =
1379                (split_boxes, split_scores, split_classes)
1380            {
1381                if let (Some(split_mask_coeff), Some(protos)) = (split_mask_coeff, protos) {
1382                    Self::verify_yolo_split_end_to_end_segdet(
1383                        &split_boxes,
1384                        &split_scores,
1385                        &split_classes,
1386                        &split_mask_coeff,
1387                        &protos,
1388                    )?;
1389                    return Ok(ModelType::YoloSplitEndToEndSegDet {
1390                        boxes: split_boxes,
1391                        scores: split_scores,
1392                        classes: split_classes,
1393                        mask_coeff: split_mask_coeff,
1394                        protos,
1395                    });
1396                }
1397                Self::verify_yolo_split_end_to_end_det(
1398                    &split_boxes,
1399                    &split_scores,
1400                    &split_classes,
1401                )?;
1402                return Ok(ModelType::YoloSplitEndToEndDet {
1403                    boxes: split_boxes,
1404                    scores: split_scores,
1405                    classes: split_classes,
1406                });
1407            } else {
1408                return Err(DecoderError::InvalidConfig(
1409                    "Invalid Yolo end-to-end model outputs".to_string(),
1410                ));
1411            }
1412        }
1413
1414        if let Some(boxes) = boxes {
1415            match (split_mask_coeff, protos) {
1416                (Some(mask_coeff), Some(protos)) => {
1417                    // 2-way split: combined detection + separate mask_coeff + protos
1418                    Self::verify_yolo_seg_det_2way(&boxes, &mask_coeff, &protos)?;
1419                    Ok(ModelType::YoloSegDet2Way {
1420                        boxes,
1421                        mask_coeff,
1422                        protos,
1423                    })
1424                }
1425                (_, Some(protos)) => {
1426                    // Unsplit: mask_coefs embedded in detection tensor
1427                    Self::verify_yolo_seg_det(&boxes, &protos)?;
1428                    Ok(ModelType::YoloSegDet { boxes, protos })
1429                }
1430                _ => {
1431                    Self::verify_yolo_det(&boxes)?;
1432                    Ok(ModelType::YoloDet { boxes })
1433                }
1434            }
1435        } else if let (Some(boxes), Some(scores)) = (split_boxes, split_scores) {
1436            if let (Some(mask_coeff), Some(protos)) = (split_mask_coeff, protos) {
1437                Self::verify_yolo_split_segdet(&boxes, &scores, &mask_coeff, &protos)?;
1438                Ok(ModelType::YoloSplitSegDet {
1439                    boxes,
1440                    scores,
1441                    mask_coeff,
1442                    protos,
1443                })
1444            } else {
1445                Self::verify_yolo_split_det(&boxes, &scores)?;
1446                Ok(ModelType::YoloSplitDet { boxes, scores })
1447            }
1448        } else {
1449            Err(DecoderError::InvalidConfig(
1450                "Invalid Yolo model outputs".to_string(),
1451            ))
1452        }
1453    }
1454
1455    fn verify_yolo_det(detect: &configs::Detection) -> Result<(), DecoderError> {
1456        if detect.shape.len() != 3 {
1457            return Err(DecoderError::InvalidConfig(format!(
1458                "Invalid Yolo Detection shape {:?}",
1459                detect.shape
1460            )));
1461        }
1462
1463        Self::verify_dshapes(
1464            &detect.dshape,
1465            &detect.shape,
1466            "Detection",
1467            &[DimName::Batch, DimName::NumFeatures, DimName::NumBoxes],
1468        )?;
1469        if !detect.dshape.is_empty() {
1470            Self::get_class_count(&detect.dshape, None, None)?;
1471        } else {
1472            Self::get_class_count_no_dshape(detect.into(), None)?;
1473        }
1474
1475        Ok(())
1476    }
1477
1478    fn verify_yolo_det_26(detect: &configs::Detection) -> Result<(), DecoderError> {
1479        if detect.shape.len() != 3 {
1480            return Err(DecoderError::InvalidConfig(format!(
1481                "Invalid Yolo Detection shape {:?}",
1482                detect.shape
1483            )));
1484        }
1485
1486        Self::verify_dshapes(
1487            &detect.dshape,
1488            &detect.shape,
1489            "Detection",
1490            &[DimName::Batch, DimName::NumFeatures, DimName::NumBoxes],
1491        )?;
1492
1493        if !detect.shape.contains(&6) {
1494            return Err(DecoderError::InvalidConfig(
1495                "Yolo26 Detection must have 6 features".to_string(),
1496            ));
1497        }
1498
1499        Ok(())
1500    }
1501
1502    fn verify_yolo_seg_det(
1503        detection: &configs::Detection,
1504        protos: &configs::Protos,
1505    ) -> Result<(), DecoderError> {
1506        if detection.shape.len() != 3 {
1507            return Err(DecoderError::InvalidConfig(format!(
1508                "Invalid Yolo Detection shape {:?}",
1509                detection.shape
1510            )));
1511        }
1512        if protos.shape.len() != 4 {
1513            return Err(DecoderError::InvalidConfig(format!(
1514                "Invalid Yolo Protos shape {:?}",
1515                protos.shape
1516            )));
1517        }
1518
1519        Self::verify_dshapes(
1520            &detection.dshape,
1521            &detection.shape,
1522            "Detection",
1523            &[DimName::Batch, DimName::NumFeatures, DimName::NumBoxes],
1524        )?;
1525        Self::verify_dshapes(
1526            &protos.dshape,
1527            &protos.shape,
1528            "Protos",
1529            &[
1530                DimName::Batch,
1531                DimName::Height,
1532                DimName::Width,
1533                DimName::NumProtos,
1534            ],
1535        )?;
1536
1537        let protos_count = Self::get_protos_count(&protos.dshape)
1538            .unwrap_or_else(|| protos.shape[1].min(protos.shape[3]));
1539        log::debug!("Protos count: {}", protos_count);
1540        log::debug!("Detection dshape: {:?}", detection.dshape);
1541        let classes = if !detection.dshape.is_empty() {
1542            Self::get_class_count(&detection.dshape, Some(protos_count), None)?
1543        } else {
1544            Self::get_class_count_no_dshape(detection.into(), Some(protos_count))?
1545        };
1546
1547        if classes == 0 {
1548            return Err(DecoderError::InvalidConfig(
1549                "Yolo Segmentation Detection has zero classes".to_string(),
1550            ));
1551        }
1552
1553        Ok(())
1554    }
1555
1556    fn verify_yolo_seg_det_2way(
1557        detection: &configs::Detection,
1558        mask_coeff: &configs::MaskCoefficients,
1559        protos: &configs::Protos,
1560    ) -> Result<(), DecoderError> {
1561        if detection.shape.len() != 3 {
1562            return Err(DecoderError::InvalidConfig(format!(
1563                "Invalid Yolo 2-Way Detection shape {:?}",
1564                detection.shape
1565            )));
1566        }
1567        if mask_coeff.shape.len() != 3 {
1568            return Err(DecoderError::InvalidConfig(format!(
1569                "Invalid Yolo 2-Way Mask Coefficients shape {:?}",
1570                mask_coeff.shape
1571            )));
1572        }
1573        if protos.shape.len() != 4 {
1574            return Err(DecoderError::InvalidConfig(format!(
1575                "Invalid Yolo 2-Way Protos shape {:?}",
1576                protos.shape
1577            )));
1578        }
1579
1580        Self::verify_dshapes(
1581            &detection.dshape,
1582            &detection.shape,
1583            "Detection",
1584            &[DimName::Batch, DimName::NumFeatures, DimName::NumBoxes],
1585        )?;
1586        Self::verify_dshapes(
1587            &mask_coeff.dshape,
1588            &mask_coeff.shape,
1589            "Mask Coefficients",
1590            &[DimName::Batch, DimName::NumProtos, DimName::NumBoxes],
1591        )?;
1592        Self::verify_dshapes(
1593            &protos.dshape,
1594            &protos.shape,
1595            "Protos",
1596            &[
1597                DimName::Batch,
1598                DimName::Height,
1599                DimName::Width,
1600                DimName::NumProtos,
1601            ],
1602        )?;
1603
1604        // Validate num_boxes match between detection and mask_coeff
1605        let det_num = Self::get_box_count(&detection.dshape).unwrap_or(detection.shape[2]);
1606        let mask_num = Self::get_box_count(&mask_coeff.dshape).unwrap_or(mask_coeff.shape[2]);
1607        if det_num != mask_num {
1608            return Err(DecoderError::InvalidConfig(format!(
1609                "Yolo 2-Way Detection num_boxes {} incompatible with Mask Coefficients num_boxes {}",
1610                det_num, mask_num
1611            )));
1612        }
1613
1614        // Validate mask_coeff channels match protos channels
1615        let mask_channels = if !mask_coeff.dshape.is_empty() {
1616            Self::get_protos_count(&mask_coeff.dshape).ok_or_else(|| {
1617                DecoderError::InvalidConfig(
1618                    "Could not find num_protos in mask_coeff config".to_string(),
1619                )
1620            })?
1621        } else {
1622            mask_coeff.shape[1]
1623        };
1624        let proto_channels = if !protos.dshape.is_empty() {
1625            Self::get_protos_count(&protos.dshape).ok_or_else(|| {
1626                DecoderError::InvalidConfig(
1627                    "Could not find num_protos in protos config".to_string(),
1628                )
1629            })?
1630        } else {
1631            protos.shape[1].min(protos.shape[3])
1632        };
1633        if mask_channels != proto_channels {
1634            return Err(DecoderError::InvalidConfig(format!(
1635                "Yolo 2-Way Protos channels {} incompatible with Mask Coefficients channels {}",
1636                proto_channels, mask_channels
1637            )));
1638        }
1639
1640        // Validate detection has classes (nc+4 features, no mask_coefs embedded)
1641        if !detection.dshape.is_empty() {
1642            Self::get_class_count(&detection.dshape, None, None)?;
1643        } else {
1644            Self::get_class_count_no_dshape(detection.into(), None)?;
1645        }
1646
1647        Ok(())
1648    }
1649
1650    fn verify_yolo_seg_det_26(
1651        detection: &configs::Detection,
1652        protos: &configs::Protos,
1653    ) -> Result<(), DecoderError> {
1654        if detection.shape.len() != 3 {
1655            return Err(DecoderError::InvalidConfig(format!(
1656                "Invalid Yolo Detection shape {:?}",
1657                detection.shape
1658            )));
1659        }
1660        if protos.shape.len() != 4 {
1661            return Err(DecoderError::InvalidConfig(format!(
1662                "Invalid Yolo Protos shape {:?}",
1663                protos.shape
1664            )));
1665        }
1666
1667        Self::verify_dshapes(
1668            &detection.dshape,
1669            &detection.shape,
1670            "Detection",
1671            &[DimName::Batch, DimName::NumFeatures, DimName::NumBoxes],
1672        )?;
1673        Self::verify_dshapes(
1674            &protos.dshape,
1675            &protos.shape,
1676            "Protos",
1677            &[
1678                DimName::Batch,
1679                DimName::Height,
1680                DimName::Width,
1681                DimName::NumProtos,
1682            ],
1683        )?;
1684
1685        let protos_count = Self::get_protos_count(&protos.dshape)
1686            .unwrap_or_else(|| protos.shape[1].min(protos.shape[3]));
1687        log::debug!("Protos count: {}", protos_count);
1688        log::debug!("Detection dshape: {:?}", detection.dshape);
1689
1690        if !detection.shape.contains(&(6 + protos_count)) {
1691            return Err(DecoderError::InvalidConfig(format!(
1692                "Yolo26 Segmentation Detection must have num_features be 6 + num_protos = {}",
1693                6 + protos_count
1694            )));
1695        }
1696
1697        Ok(())
1698    }
1699
1700    fn verify_yolo_split_det(
1701        boxes: &configs::Boxes,
1702        scores: &configs::Scores,
1703    ) -> Result<(), DecoderError> {
1704        if boxes.shape.len() != 3 {
1705            return Err(DecoderError::InvalidConfig(format!(
1706                "Invalid Yolo Split Boxes shape {:?}",
1707                boxes.shape
1708            )));
1709        }
1710        if scores.shape.len() != 3 {
1711            return Err(DecoderError::InvalidConfig(format!(
1712                "Invalid Yolo Split Scores shape {:?}",
1713                scores.shape
1714            )));
1715        }
1716
1717        Self::verify_dshapes(
1718            &boxes.dshape,
1719            &boxes.shape,
1720            "Boxes",
1721            &[DimName::Batch, DimName::BoxCoords, DimName::NumBoxes],
1722        )?;
1723        Self::verify_dshapes(
1724            &scores.dshape,
1725            &scores.shape,
1726            "Scores",
1727            &[DimName::Batch, DimName::NumClasses, DimName::NumBoxes],
1728        )?;
1729
1730        let boxes_num = Self::get_box_count(&boxes.dshape).unwrap_or(boxes.shape[2]);
1731        let scores_num = Self::get_box_count(&scores.dshape).unwrap_or(scores.shape[2]);
1732
1733        if boxes_num != scores_num {
1734            return Err(DecoderError::InvalidConfig(format!(
1735                "Yolo Split Detection Boxes num {} incompatible with Scores num {}",
1736                boxes_num, scores_num
1737            )));
1738        }
1739
1740        Ok(())
1741    }
1742
1743    fn verify_yolo_split_segdet(
1744        boxes: &configs::Boxes,
1745        scores: &configs::Scores,
1746        mask_coeff: &configs::MaskCoefficients,
1747        protos: &configs::Protos,
1748    ) -> Result<(), DecoderError> {
1749        if boxes.shape.len() != 3 {
1750            return Err(DecoderError::InvalidConfig(format!(
1751                "Invalid Yolo Split Boxes shape {:?}",
1752                boxes.shape
1753            )));
1754        }
1755        if scores.shape.len() != 3 {
1756            return Err(DecoderError::InvalidConfig(format!(
1757                "Invalid Yolo Split Scores shape {:?}",
1758                scores.shape
1759            )));
1760        }
1761
1762        if mask_coeff.shape.len() != 3 {
1763            return Err(DecoderError::InvalidConfig(format!(
1764                "Invalid Yolo Split Mask Coefficients shape {:?}",
1765                mask_coeff.shape
1766            )));
1767        }
1768
1769        if protos.shape.len() != 4 {
1770            return Err(DecoderError::InvalidConfig(format!(
1771                "Invalid Yolo Protos shape {:?}",
1772                mask_coeff.shape
1773            )));
1774        }
1775
1776        Self::verify_dshapes(
1777            &boxes.dshape,
1778            &boxes.shape,
1779            "Boxes",
1780            &[DimName::Batch, DimName::BoxCoords, DimName::NumBoxes],
1781        )?;
1782        Self::verify_dshapes(
1783            &scores.dshape,
1784            &scores.shape,
1785            "Scores",
1786            &[DimName::Batch, DimName::NumClasses, DimName::NumBoxes],
1787        )?;
1788        Self::verify_dshapes(
1789            &mask_coeff.dshape,
1790            &mask_coeff.shape,
1791            "Mask Coefficients",
1792            &[DimName::Batch, DimName::NumProtos, DimName::NumBoxes],
1793        )?;
1794        Self::verify_dshapes(
1795            &protos.dshape,
1796            &protos.shape,
1797            "Protos",
1798            &[
1799                DimName::Batch,
1800                DimName::Height,
1801                DimName::Width,
1802                DimName::NumProtos,
1803            ],
1804        )?;
1805
1806        let boxes_num = Self::get_box_count(&boxes.dshape).unwrap_or(boxes.shape[2]);
1807        let scores_num = Self::get_box_count(&scores.dshape).unwrap_or(scores.shape[2]);
1808        let mask_num = Self::get_box_count(&mask_coeff.dshape).unwrap_or(mask_coeff.shape[2]);
1809
1810        let mask_channels = if !mask_coeff.dshape.is_empty() {
1811            Self::get_protos_count(&mask_coeff.dshape).ok_or_else(|| {
1812                DecoderError::InvalidConfig("Could not find num_protos in config".to_string())
1813            })?
1814        } else {
1815            mask_coeff.shape[1]
1816        };
1817        let proto_channels = if !protos.dshape.is_empty() {
1818            Self::get_protos_count(&protos.dshape).ok_or_else(|| {
1819                DecoderError::InvalidConfig("Could not find num_protos in config".to_string())
1820            })?
1821        } else {
1822            protos.shape[1].min(protos.shape[3])
1823        };
1824
1825        if boxes_num != scores_num {
1826            return Err(DecoderError::InvalidConfig(format!(
1827                "Yolo Split Detection Boxes num {} incompatible with Scores num {}",
1828                boxes_num, scores_num
1829            )));
1830        }
1831
1832        if boxes_num != mask_num {
1833            return Err(DecoderError::InvalidConfig(format!(
1834                "Yolo Split Detection Boxes num {} incompatible with Mask Coefficients num {}",
1835                boxes_num, mask_num
1836            )));
1837        }
1838
1839        if proto_channels != mask_channels {
1840            return Err(DecoderError::InvalidConfig(format!(
1841                "Yolo Protos channels {} incompatible with Mask Coefficients channels {}",
1842                proto_channels, mask_channels
1843            )));
1844        }
1845
1846        Ok(())
1847    }
1848
1849    fn verify_yolo_split_end_to_end_det(
1850        boxes: &configs::Boxes,
1851        scores: &configs::Scores,
1852        classes: &configs::Classes,
1853    ) -> Result<(), DecoderError> {
1854        if boxes.shape.len() != 3 || !boxes.shape.contains(&4) {
1855            return Err(DecoderError::InvalidConfig(format!(
1856                "Split end-to-end boxes must be [batch, N, 4], got {:?}",
1857                boxes.shape
1858            )));
1859        }
1860        if scores.shape.len() != 3 || !scores.shape.contains(&1) {
1861            return Err(DecoderError::InvalidConfig(format!(
1862                "Split end-to-end scores must be [batch, N, 1], got {:?}",
1863                scores.shape
1864            )));
1865        }
1866        if classes.shape.len() != 3 || !classes.shape.contains(&1) {
1867            return Err(DecoderError::InvalidConfig(format!(
1868                "Split end-to-end classes must be [batch, N, 1], got {:?}",
1869                classes.shape
1870            )));
1871        }
1872        Ok(())
1873    }
1874
1875    fn verify_yolo_split_end_to_end_segdet(
1876        boxes: &configs::Boxes,
1877        scores: &configs::Scores,
1878        classes: &configs::Classes,
1879        mask_coeff: &configs::MaskCoefficients,
1880        protos: &configs::Protos,
1881    ) -> Result<(), DecoderError> {
1882        Self::verify_yolo_split_end_to_end_det(boxes, scores, classes)?;
1883        if mask_coeff.shape.len() != 3 {
1884            return Err(DecoderError::InvalidConfig(format!(
1885                "Invalid split end-to-end mask coefficients shape {:?}",
1886                mask_coeff.shape
1887            )));
1888        }
1889        if protos.shape.len() != 4 {
1890            return Err(DecoderError::InvalidConfig(format!(
1891                "Invalid protos shape {:?}",
1892                protos.shape
1893            )));
1894        }
1895        Ok(())
1896    }
1897
1898    fn get_model_type_modelpack(configs: ConfigOutputs) -> Result<ModelType, DecoderError> {
1899        let mut split_decoders = Vec::new();
1900        let mut segment_ = None;
1901        let mut scores_ = None;
1902        let mut boxes_ = None;
1903        for c in configs.outputs {
1904            match c {
1905                ConfigOutput::Detection(detection) => split_decoders.push(detection),
1906                ConfigOutput::Segmentation(segmentation) => segment_ = Some(segmentation),
1907                ConfigOutput::Mask(_) => {}
1908                ConfigOutput::Protos(_) => {
1909                    return Err(DecoderError::InvalidConfig(
1910                        "ModelPack should not have protos".to_string(),
1911                    ));
1912                }
1913                ConfigOutput::Scores(scores) => scores_ = Some(scores),
1914                ConfigOutput::Boxes(boxes) => boxes_ = Some(boxes),
1915                ConfigOutput::MaskCoefficients(_) => {
1916                    return Err(DecoderError::InvalidConfig(
1917                        "ModelPack should not have mask coefficients".to_string(),
1918                    ));
1919                }
1920                ConfigOutput::Classes(_) => {
1921                    return Err(DecoderError::InvalidConfig(
1922                        "ModelPack should not have classes output".to_string(),
1923                    ));
1924                }
1925            }
1926        }
1927
1928        if let Some(segmentation) = segment_ {
1929            if !split_decoders.is_empty() {
1930                let classes = Self::verify_modelpack_split_det(&split_decoders)?;
1931                Self::verify_modelpack_seg(&segmentation, Some(classes))?;
1932                Ok(ModelType::ModelPackSegDetSplit {
1933                    detection: split_decoders,
1934                    segmentation,
1935                })
1936            } else if let (Some(scores), Some(boxes)) = (scores_, boxes_) {
1937                let classes = Self::verify_modelpack_det(&boxes, &scores)?;
1938                Self::verify_modelpack_seg(&segmentation, Some(classes))?;
1939                Ok(ModelType::ModelPackSegDet {
1940                    boxes,
1941                    scores,
1942                    segmentation,
1943                })
1944            } else {
1945                Self::verify_modelpack_seg(&segmentation, None)?;
1946                Ok(ModelType::ModelPackSeg { segmentation })
1947            }
1948        } else if !split_decoders.is_empty() {
1949            Self::verify_modelpack_split_det(&split_decoders)?;
1950            Ok(ModelType::ModelPackDetSplit {
1951                detection: split_decoders,
1952            })
1953        } else if let (Some(scores), Some(boxes)) = (scores_, boxes_) {
1954            Self::verify_modelpack_det(&boxes, &scores)?;
1955            Ok(ModelType::ModelPackDet { boxes, scores })
1956        } else {
1957            Err(DecoderError::InvalidConfig(
1958                "Invalid ModelPack model outputs".to_string(),
1959            ))
1960        }
1961    }
1962
1963    fn verify_modelpack_det(
1964        boxes: &configs::Boxes,
1965        scores: &configs::Scores,
1966    ) -> Result<usize, DecoderError> {
1967        if boxes.shape.len() != 4 {
1968            return Err(DecoderError::InvalidConfig(format!(
1969                "Invalid ModelPack Boxes shape {:?}",
1970                boxes.shape
1971            )));
1972        }
1973        if scores.shape.len() != 3 {
1974            return Err(DecoderError::InvalidConfig(format!(
1975                "Invalid ModelPack Scores shape {:?}",
1976                scores.shape
1977            )));
1978        }
1979
1980        Self::verify_dshapes(
1981            &boxes.dshape,
1982            &boxes.shape,
1983            "Boxes",
1984            &[
1985                DimName::Batch,
1986                DimName::NumBoxes,
1987                DimName::Padding,
1988                DimName::BoxCoords,
1989            ],
1990        )?;
1991        Self::verify_dshapes(
1992            &scores.dshape,
1993            &scores.shape,
1994            "Scores",
1995            &[DimName::Batch, DimName::NumBoxes, DimName::NumClasses],
1996        )?;
1997
1998        let boxes_num = Self::get_box_count(&boxes.dshape).unwrap_or(boxes.shape[1]);
1999        let scores_num = Self::get_box_count(&scores.dshape).unwrap_or(scores.shape[1]);
2000
2001        if boxes_num != scores_num {
2002            return Err(DecoderError::InvalidConfig(format!(
2003                "ModelPack Detection Boxes num {} incompatible with Scores num {}",
2004                boxes_num, scores_num
2005            )));
2006        }
2007
2008        let num_classes = if !scores.dshape.is_empty() {
2009            Self::get_class_count(&scores.dshape, None, None)?
2010        } else {
2011            Self::get_class_count_no_dshape(scores.into(), None)?
2012        };
2013
2014        Ok(num_classes)
2015    }
2016
2017    fn verify_modelpack_split_det(boxes: &[configs::Detection]) -> Result<usize, DecoderError> {
2018        let mut num_classes = None;
2019        for b in boxes {
2020            let Some(num_anchors) = b.anchors.as_ref().map(|a| a.len()) else {
2021                return Err(DecoderError::InvalidConfig(
2022                    "ModelPack Split Detection missing anchors".to_string(),
2023                ));
2024            };
2025
2026            if num_anchors == 0 {
2027                return Err(DecoderError::InvalidConfig(
2028                    "ModelPack Split Detection has zero anchors".to_string(),
2029                ));
2030            }
2031
2032            if b.shape.len() != 4 {
2033                return Err(DecoderError::InvalidConfig(format!(
2034                    "Invalid ModelPack Split Detection shape {:?}",
2035                    b.shape
2036                )));
2037            }
2038
2039            Self::verify_dshapes(
2040                &b.dshape,
2041                &b.shape,
2042                "Split Detection",
2043                &[
2044                    DimName::Batch,
2045                    DimName::Height,
2046                    DimName::Width,
2047                    DimName::NumAnchorsXFeatures,
2048                ],
2049            )?;
2050            let classes = if !b.dshape.is_empty() {
2051                Self::get_class_count(&b.dshape, None, Some(num_anchors))?
2052            } else {
2053                Self::get_class_count_no_dshape(b.into(), None)?
2054            };
2055
2056            match num_classes {
2057                Some(n) => {
2058                    if n != classes {
2059                        return Err(DecoderError::InvalidConfig(format!(
2060                            "ModelPack Split Detection inconsistent number of classes: previous {}, current {}",
2061                            n, classes
2062                        )));
2063                    }
2064                }
2065                None => {
2066                    num_classes = Some(classes);
2067                }
2068            }
2069        }
2070
2071        Ok(num_classes.unwrap_or(0))
2072    }
2073
2074    fn verify_modelpack_seg(
2075        segmentation: &configs::Segmentation,
2076        classes: Option<usize>,
2077    ) -> Result<(), DecoderError> {
2078        if segmentation.shape.len() != 4 {
2079            return Err(DecoderError::InvalidConfig(format!(
2080                "Invalid ModelPack Segmentation shape {:?}",
2081                segmentation.shape
2082            )));
2083        }
2084        // ModelPack segmentation is canonical NHWC: [Batch, Height, Width,
2085        // NumClasses]. Structural integrity of the dshape (rank, per-axis size
2086        // agreement with `shape`, no duplicate axes) is already enforced for
2087        // every output by `Decoder::validate_output_layout` at build entry, so
2088        // here we only require the spatial axes to be named and recover the
2089        // class count from the channel axis below.
2090        //
2091        // Metadata producers that lack embedded metadata (the validator's
2092        // shape-inference fallback for ModelPack `.keras` models) tag the
2093        // channel axis as `num_protos` or omit its name, which previously
2094        // hard-failed here with "Segmentation dshape missing required
2095        // dimension NumClasses" (DE-2651 / DE-2628). That is unnecessarily
2096        // strict: an unnamed or mis-named channel axis at position 3 is sorted
2097        // to the canonical tail by `swap_axes_if_needed`, so the physical NHWC
2098        // order — and therefore decode — is unaffected. Infer the class count
2099        // from the NHWC shape instead of rejecting the model.
2100        if !segmentation.dshape.is_empty() {
2101            let names =
2102                HashSet::<DimName>::from_iter(segmentation.dshape.iter().map(|(name, _)| *name));
2103            for dim in [DimName::Batch, DimName::Height, DimName::Width] {
2104                if !names.contains(&dim) {
2105                    return Err(DecoderError::InvalidConfig(format!(
2106                        "Segmentation dshape missing required dimension {dim:?}"
2107                    )));
2108                }
2109            }
2110        }
2111
2112        if let Some(classes) = classes {
2113            let seg_classes = Self::modelpack_seg_num_classes(segmentation);
2114
2115            if seg_classes != classes + 1 {
2116                return Err(DecoderError::InvalidConfig(format!(
2117                    "ModelPack Segmentation channels {} incompatible with number of classes {}",
2118                    seg_classes, classes
2119                )));
2120            }
2121        }
2122        Ok(())
2123    }
2124
2125    /// Number of segmentation classes for a ModelPack segmentation output.
2126    ///
2127    /// Prefers an explicit `NumClasses` dshape dimension when present;
2128    /// otherwise falls back to the canonical NHWC channel axis (`shape[3]`).
2129    /// This tolerates metadata that omits or mis-tags the class axis while
2130    /// still honouring a correctly-tagged dshape. Only called after
2131    /// `verify_modelpack_seg` has validated that `shape` is rank-4 (it returns
2132    /// `InvalidConfig` otherwise), so `shape[3]` cannot panic.
2133    fn modelpack_seg_num_classes(segmentation: &configs::Segmentation) -> usize {
2134        for (dim_name, dim_size) in &segmentation.dshape {
2135            if *dim_name == DimName::NumClasses {
2136                return *dim_size;
2137            }
2138        }
2139        segmentation.shape[3]
2140    }
2141
2142    // verifies that dshapes match the given shape
2143    fn verify_dshapes(
2144        dshape: &[(DimName, usize)],
2145        shape: &[usize],
2146        name: &str,
2147        dims: &[DimName],
2148    ) -> Result<(), DecoderError> {
2149        for s in shape {
2150            if *s == 0 {
2151                return Err(DecoderError::InvalidConfig(format!(
2152                    "{} shape has zero dimension",
2153                    name
2154                )));
2155            }
2156        }
2157
2158        if shape.len() != dims.len() {
2159            return Err(DecoderError::InvalidConfig(format!(
2160                "{} shape length {} does not match expected dims length {}",
2161                name,
2162                shape.len(),
2163                dims.len()
2164            )));
2165        }
2166
2167        if dshape.is_empty() {
2168            return Ok(());
2169        }
2170        // check the dshape lengths match the shape lengths
2171        if dshape.len() != shape.len() {
2172            return Err(DecoderError::InvalidConfig(format!(
2173                "{} dshape length does not match shape length",
2174                name
2175            )));
2176        }
2177
2178        // check the dshape values match the shape values
2179        for ((dim_name, dim_size), shape_size) in dshape.iter().zip(shape) {
2180            if dim_size != shape_size {
2181                return Err(DecoderError::InvalidConfig(format!(
2182                    "{} dshape dimension {} size {} does not match shape size {}",
2183                    name, dim_name, dim_size, shape_size
2184                )));
2185            }
2186            if *dim_name == DimName::Padding && *dim_size != 1 {
2187                return Err(DecoderError::InvalidConfig(
2188                    "Padding dimension size must be 1".to_string(),
2189                ));
2190            }
2191
2192            if *dim_name == DimName::BoxCoords && *dim_size != 4 {
2193                return Err(DecoderError::InvalidConfig(
2194                    "BoxCoords dimension size must be 4".to_string(),
2195                ));
2196            }
2197        }
2198
2199        let dims_present = HashSet::<DimName>::from_iter(dshape.iter().map(|(name, _)| *name));
2200        for dim in dims {
2201            if !dims_present.contains(dim) {
2202                return Err(DecoderError::InvalidConfig(format!(
2203                    "{} dshape missing required dimension {:?}",
2204                    name, dim
2205                )));
2206            }
2207        }
2208
2209        Ok(())
2210    }
2211
2212    fn get_box_count(dshape: &[(DimName, usize)]) -> Option<usize> {
2213        for (dim_name, dim_size) in dshape {
2214            if *dim_name == DimName::NumBoxes {
2215                return Some(*dim_size);
2216            }
2217        }
2218        None
2219    }
2220
2221    fn get_class_count_no_dshape(
2222        config: ConfigOutputRef,
2223        protos: Option<usize>,
2224    ) -> Result<usize, DecoderError> {
2225        match config {
2226            ConfigOutputRef::Detection(detection) => match detection.decoder {
2227                DecoderType::Ultralytics => {
2228                    if detection.shape[1] <= 4 + protos.unwrap_or(0) {
2229                        return Err(DecoderError::InvalidConfig(format!(
2230                            "Invalid shape: Yolo num_features {} must be greater than {}",
2231                            detection.shape[1],
2232                            4 + protos.unwrap_or(0),
2233                        )));
2234                    }
2235                    Ok(detection.shape[1] - 4 - protos.unwrap_or(0))
2236                }
2237                DecoderType::ModelPack => {
2238                    let Some(num_anchors) = detection.anchors.as_ref().map(|a| a.len()) else {
2239                        return Err(DecoderError::Internal(
2240                            "ModelPack Detection missing anchors".to_string(),
2241                        ));
2242                    };
2243                    let anchors_x_features = detection.shape[3];
2244                    if anchors_x_features <= num_anchors * 5 {
2245                        return Err(DecoderError::InvalidConfig(format!(
2246                            "Invalid ModelPack Split Detection shape: anchors_x_features {} not greater than number of anchors * 5 = {}",
2247                            anchors_x_features,
2248                            num_anchors * 5,
2249                        )));
2250                    }
2251
2252                    if !anchors_x_features.is_multiple_of(num_anchors) {
2253                        return Err(DecoderError::InvalidConfig(format!(
2254                            "Invalid ModelPack Split Detection shape: anchors_x_features {} not a multiple of number of anchors {}",
2255                            anchors_x_features, num_anchors
2256                        )));
2257                    }
2258                    Ok(anchors_x_features / num_anchors - 5)
2259                }
2260            },
2261
2262            ConfigOutputRef::Scores(scores) => match scores.decoder {
2263                DecoderType::Ultralytics => Ok(scores.shape[1]),
2264                DecoderType::ModelPack => Ok(scores.shape[2]),
2265            },
2266            _ => Err(DecoderError::Internal(
2267                "Attempted to get class count from unsupported config output".to_owned(),
2268            )),
2269        }
2270    }
2271
2272    // get the class count from dshape or calculate from num_features
2273    fn get_class_count(
2274        dshape: &[(DimName, usize)],
2275        protos: Option<usize>,
2276        anchors: Option<usize>,
2277    ) -> Result<usize, DecoderError> {
2278        if dshape.is_empty() {
2279            return Ok(0);
2280        }
2281        // if it has num_classes in dshape, return it
2282        for (dim_name, dim_size) in dshape {
2283            if *dim_name == DimName::NumClasses {
2284                return Ok(*dim_size);
2285            }
2286        }
2287
2288        // number of classes can be calculated from num_features - 4 for yolo.  If the
2289        // model has protos, we also subtract the number of protos.
2290        for (dim_name, dim_size) in dshape {
2291            if *dim_name == DimName::NumFeatures {
2292                let protos = protos.unwrap_or(0);
2293                if protos + 4 >= *dim_size {
2294                    return Err(DecoderError::InvalidConfig(format!(
2295                        "Invalid shape: Yolo num_features {} must be greater than {}",
2296                        *dim_size,
2297                        protos + 4,
2298                    )));
2299                }
2300                return Ok(*dim_size - 4 - protos);
2301            }
2302        }
2303
2304        // number of classes can be calculated from number of anchors for modelpack
2305        // split detection
2306        if let Some(num_anchors) = anchors {
2307            for (dim_name, dim_size) in dshape {
2308                if *dim_name == DimName::NumAnchorsXFeatures {
2309                    let anchors_x_features = *dim_size;
2310                    if anchors_x_features <= num_anchors * 5 {
2311                        return Err(DecoderError::InvalidConfig(format!(
2312                            "Invalid ModelPack Split Detection shape: anchors_x_features {} not greater than number of anchors * 5 = {}",
2313                            anchors_x_features,
2314                            num_anchors * 5,
2315                        )));
2316                    }
2317
2318                    if !anchors_x_features.is_multiple_of(num_anchors) {
2319                        return Err(DecoderError::InvalidConfig(format!(
2320                            "Invalid ModelPack Split Detection shape: anchors_x_features {} not a multiple of number of anchors {}",
2321                            anchors_x_features, num_anchors
2322                        )));
2323                    }
2324                    return Ok((anchors_x_features / num_anchors) - 5);
2325                }
2326            }
2327        }
2328        Err(DecoderError::InvalidConfig(
2329            "Cannot determine number of classes from dshape".to_owned(),
2330        ))
2331    }
2332
2333    fn get_protos_count(dshape: &[(DimName, usize)]) -> Option<usize> {
2334        for (dim_name, dim_size) in dshape {
2335            if *dim_name == DimName::NumProtos {
2336                return Some(*dim_size);
2337            }
2338        }
2339        None
2340    }
2341}