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}