Skip to main content

edgefirst_decoder/
infer.rs

1// SPDX-FileCopyrightText: Copyright 2025-2026 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4//! Ultralytics YOLO schema inference from raw model I/O signals.
5//!
6//! Model export pipelines (ONNX, TFLite, CoreML) carry Ultralytics-authored
7//! metadata (class names, task, input size) alongside the tensor shapes
8//! and dtypes the runtime reports. This module turns that raw signal into
9//! a [`crate::schema::SchemaV2`] the decoder can act on, without requiring
10//! a hand-written `edgefirst.json` for every Ultralytics export.
11
12use std::collections::{BTreeMap, BTreeSet};
13use std::fmt;
14
15use crate::configs::DimName;
16use crate::schema::{
17    BoxEncoding, DType, DecoderKind, DecoderVersion, InputSpec, LogicalOutput, LogicalType,
18    NmsMode, Quantization, SchemaV2, ScoreFormat,
19};
20
21/// A single input or output tensor as reported by the inference runtime.
22#[derive(Debug, Clone, PartialEq)]
23pub struct TensorInfo {
24    /// Tensor name as reported by the runtime.
25    pub name: String,
26    /// Tensor shape in the model's native layout.
27    pub shape: Vec<usize>,
28    /// Tensor element dtype.
29    pub dtype: DType,
30    /// Quantization parameters, if the tensor is quantized.
31    pub quantization: Option<Quantization>,
32}
33
34edgefirst_tensor::ef_vocabulary! {
35    /// Container format of the model the signals were read from. Drives
36    /// Ultralytics per-format conventions (box normalization, tensor order).
37    ///
38    /// The code is the one the C `EF_MODEL_SOURCE_*` enumerators use; the
39    /// string form is the one the Python bindings accept.
40    pub enum ModelSource {
41        /// An ONNX export. Ultralytics emits pixel-space box coordinates here,
42        /// so the inferred schema sets `normalized: false`.
43        Onnx = 0, "onnx", ONNX,
44        /// A TFLite/LiteRT export. Ultralytics emits boxes normalized to
45        /// `[0, 1]` here, so the inferred schema sets `normalized: true`.
46        TfLite = 1, "tflite", TFLITE,
47        /// A CoreML export (`.mlpackage` / `.mlmodelc`). Ultralytics' CoreML
48        /// converter traces the same graph the ONNX exporter traces — the
49        /// `1/w` normalize lives only in `IOSDetectModel`, which is used
50        /// solely on the `nms=True` pipeline path — so a no-NMS export emits
51        /// pixel-space box coordinates and the inferred schema sets
52        /// `normalized: false`, matching [`ModelSource::Onnx`].
53        ///
54        /// Established by source-tracing the Ultralytics exporter and
55        /// confirmed by runtime magnitude measurement against a matched fp16
56        /// ONNX pair; `testdata/infer/NOTES.md` records the reasoning, the
57        /// measurement and the captured fixture.
58        ///
59        /// This resolution covers the no-NMS export only. A `nms=True` export
60        /// runs Apple's NMS pipeline (`image`/`iouThreshold`/
61        /// `confidenceThreshold` inputs, dynamically-shaped `confidence`/
62        /// `coordinates` outputs) — a different artifact entirely, not the
63        /// anchor-grid tensor this resolution reasons about. Inference does
64        /// not classify it as a schema: neither output's shape can match the
65        /// `[1, 4+nc, A]` anchor-grid layout this module looks for.
66        CoreMl = 3, "coreml", COREML,
67        /// Any other container. Inference **refuses** this with
68        /// [`InferError::UnknownBoxConvention`] rather than assuming a box
69        /// convention: whether coordinates are pixel-space or `[0, 1]` follows
70        /// the exporter, is not derivable from tensor shapes, and getting it
71        /// wrong scales every box by the input size. The variant exists so a
72        /// signal collector can report honestly what it read; add a measured
73        /// container here rather than routing it through `Other`.
74        Other = 2, "other", OTHER,
75    }
76    #[doc(hidden)]
77    pub mod model_source_code;
78}
79
80/// Raw model I/O signals: tensor shapes/dtypes plus unparsed metadata, as
81/// reported by the inference runtime that loaded the model.
82#[derive(Debug, Clone, PartialEq)]
83pub struct ModelSignals {
84    /// Container format the signals were read from.
85    pub source: ModelSource,
86    /// Input tensors.
87    pub inputs: Vec<TensorInfo>,
88    /// Output tensors.
89    pub outputs: Vec<TensorInfo>,
90    /// Raw model metadata key/values (ONNX metadata_props, TFLite
91    /// metadata entries, CoreML's own metadata alongside Ultralytics'
92    /// props). Values are passed verbatim; parsing happens here.
93    pub metadata: BTreeMap<String, String>,
94}
95
96/// Result of inferring a schema from [`ModelSignals`].
97#[derive(Debug, Clone, PartialEq)]
98pub struct InferredSchema {
99    /// The inferred schema.
100    pub schema: SchemaV2,
101    /// Class names, in index order.
102    pub labels: Vec<String>,
103    /// Human-readable summary, e.g. "Ultralytics YOLO26 segment, 80 classes".
104    pub description: String,
105}
106
107/// Errors that can occur while inferring an Ultralytics schema.
108#[derive(Debug)]
109pub enum InferError {
110    /// Model metadata carries no Ultralytics signature.
111    NotUltralytics(String),
112    /// The Ultralytics `names` metadata field could not be parsed.
113    BadNames(String),
114    /// A metadata field is present but carries a value that cannot be
115    /// interpreted. Distinct from an absent field, which is allowed to
116    /// select a fallback.
117    BadMetadata(String),
118    /// The Ultralytics `task` metadata field names an unsupported task.
119    UnsupportedTask(String),
120    /// The number of classes in `names` does not match what the output
121    /// tensor shapes imply.
122    ClassCountMismatch { expected: usize, found: usize },
123    /// The output tensor layout does not match any supported Ultralytics
124    /// convention.
125    UnsupportedLayout(String),
126    /// More than one output matches the same layout, so which tensor
127    /// carries the detections depends on the order the runtime reported
128    /// them. Reported rather than resolved by position.
129    AmbiguousLayout(String),
130    /// A boundary tensor carries quantization the decoder cannot consume.
131    /// Currently this means per-channel quantization: the decoder supports
132    /// per-tensor only, so such a schema would build a decoder that fails.
133    UnsupportedQuantization(String),
134    /// The signals came from [`ModelSource::Other`], for which the box
135    /// coordinate convention is unmeasured. Only the ONNX (pixel-space),
136    /// TFLite (`[0, 1]`), and CoreML (pixel-space) conventions are
137    /// characterized, and picking the wrong one silently scales every box
138    /// by the input size.
139    UnknownBoxConvention,
140}
141
142impl fmt::Display for InferError {
143    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
144        match self {
145            InferError::NotUltralytics(s) => {
146                write!(f, "model metadata carries no Ultralytics signature: {s}")
147            }
148            InferError::BadNames(s) => {
149                write!(f, "cannot parse Ultralytics `names` metadata: {s}")
150            }
151            InferError::BadMetadata(s) => {
152                write!(f, "malformed Ultralytics metadata: {s}")
153            }
154            InferError::UnsupportedTask(s) => write!(
155                f,
156                "unsupported Ultralytics task `{s}` (supported: detect, segment)"
157            ),
158            InferError::ClassCountMismatch { expected, found } => write!(
159                f,
160                "class-count mismatch: names give {expected} classes but output features imply {found}"
161            ),
162            InferError::UnsupportedLayout(s) => write!(f, "unsupported output layout: {s}"),
163            InferError::AmbiguousLayout(s) => write!(f, "ambiguous output layout: {s}"),
164            InferError::UnsupportedQuantization(s) => {
165                write!(f, "unsupported quantization: {s}")
166            }
167            InferError::UnknownBoxConvention => write!(
168                f,
169                "box coordinate convention is unknown for this container; only \
170                 onnx (pixel-space), tflite ([0,1]), and coreml (pixel-space) \
171                 are characterized"
172            ),
173        }
174    }
175}
176
177impl std::error::Error for InferError {}
178
179/// Supported Ultralytics export tasks.
180#[derive(Debug, Clone, Copy, PartialEq, Eq)]
181enum Task {
182    Detect,
183    Segment,
184}
185
186impl Task {
187    fn parse(s: &str) -> Result<Self, InferError> {
188        match s.trim().to_ascii_lowercase().as_str() {
189            "detect" => Ok(Task::Detect),
190            "segment" => Ok(Task::Segment),
191            other => Err(InferError::UnsupportedTask(other.to_string())),
192        }
193    }
194}
195
196/// Parsed Ultralytics export metadata, in whichever raw form (ONNX flat
197/// string props, or TFLite `metadata.json` envelope) it was captured.
198#[derive(Debug, Clone, PartialEq)]
199struct UltralyticsMeta {
200    names: Vec<String>,
201    task: Option<Task>,
202    /// YOLO26's built-in NMS-free end-to-end head signal. `Some(true)` for
203    /// YOLO26 exports, `Some(false)` for NMS-required heads (v8/v11),
204    /// `None` when the metadata omits the field entirely.
205    end2end: Option<bool>,
206}
207
208/// Refuses metadata that positively identifies a *different* vendor.
209///
210/// A `names` map alone is not an Ultralytics signature — any exporter can
211/// write class names — and this function then applies Ultralytics box,
212/// score and normalization conventions to whatever it was given. Real
213/// exports carry `author: "Ultralytics"` and a `docs` URL, so when either
214/// is present and names someone else, that is a positive contradiction and
215/// an error.
216///
217/// Absence is not treated as a contradiction: intermediate tools
218/// (quantizers, format converters) routinely preserve `names` while
219/// dropping provenance fields, and those exports are still the ones this
220/// module exists to read. The layout cross-checks that follow remain the
221/// substantive guard — a model whose shapes don't fit is rejected whatever
222/// its metadata claims.
223fn reject_foreign_vendor(author: Option<&String>, docs: Option<&String>) -> Result<(), InferError> {
224    let names_ultralytics = |v: &String| v.to_ascii_lowercase().contains("ultralytics");
225    for (field, value) in [("author", author), ("docs", docs)] {
226        if let Some(v) = value {
227            if !v.trim().is_empty() && !names_ultralytics(v) {
228                return Err(InferError::NotUltralytics(format!(
229                    "metadata `{field}` is `{v}`, not Ultralytics"
230                )));
231            }
232        }
233    }
234    Ok(())
235}
236
237impl UltralyticsMeta {
238    /// Extracts Ultralytics metadata from a model's raw metadata map.
239    ///
240    /// Handles two forms: flat string props (ONNX `metadata_props`, one
241    /// key per field) and the TFLite `metadata.json` envelope (a single
242    /// entry whose *value* is a JSON document carrying the same fields).
243    fn from_metadata(meta: &BTreeMap<String, String>) -> Result<Self, InferError> {
244        if let Some(names_raw) = meta.get("names") {
245            reject_foreign_vendor(meta.get("author"), meta.get("docs"))?;
246            return Self::from_flat_props(meta, names_raw);
247        }
248
249        // TFLite form: scan every metadata value for a JSON object that
250        // carries a `names` field.
251        for value in meta.values() {
252            if let Ok(serde_json::Value::Object(obj)) = serde_json::from_str(value) {
253                if let Some(names) = obj.get("names") {
254                    let author = obj
255                        .get("author")
256                        .and_then(|v| v.as_str())
257                        .map(str::to_string);
258                    let docs = obj.get("docs").and_then(|v| v.as_str()).map(str::to_string);
259                    reject_foreign_vendor(author.as_ref(), docs.as_ref())?;
260                    return Self::from_json_envelope(&obj, names);
261                }
262            }
263        }
264
265        Err(InferError::NotUltralytics(
266            "no `names` metadata field or metadata.json envelope found".into(),
267        ))
268    }
269
270    fn from_flat_props(
271        meta: &BTreeMap<String, String>,
272        names_raw: &str,
273    ) -> Result<Self, InferError> {
274        let names = parse_names(names_raw)?;
275        let task = meta.get("task").map(|t| Task::parse(t)).transpose()?;
276        let end2end = meta.get("end2end").map(|s| parse_bool_str(s)).transpose()?;
277        Ok(UltralyticsMeta {
278            names,
279            task,
280            end2end,
281        })
282    }
283
284    /// Takes the `names` value directly rather than re-looking it up:
285    /// [`Self::from_metadata`] only reaches here after finding the key, so
286    /// a second lookup would add an error arm that can never fire.
287    fn from_json_envelope(
288        obj: &serde_json::Map<String, serde_json::Value>,
289        names_value: &serde_json::Value,
290    ) -> Result<Self, InferError> {
291        let names = names_from_json_value(names_value)?;
292        let task = obj
293            .get("task")
294            .and_then(|v| v.as_str())
295            .map(Task::parse)
296            .transpose()?;
297        // The envelope types it properly, so anything non-boolean is
298        // corrupt rather than merely absent.
299        let end2end = match obj.get("end2end") {
300            None | Some(serde_json::Value::Null) => None,
301            Some(serde_json::Value::Bool(b)) => Some(*b),
302            Some(other) => {
303                return Err(InferError::BadMetadata(format!(
304                    "`end2end` is {other}, expected a boolean"
305                )))
306            }
307        };
308        Ok(UltralyticsMeta {
309            names,
310            task,
311            end2end,
312        })
313    }
314}
315
316/// Parses an Ultralytics `end2end` metadata string value (`"True"` /
317/// `"False"`, case-insensitively).
318///
319/// An unparseable value is an error rather than `None`: `None` means "the
320/// key is absent", which selects the shape-only fallback. Letting a
321/// corrupt value collapse into that would classify by shape while the
322/// model was actively claiming something else.
323fn parse_bool_str(s: &str) -> Result<bool, InferError> {
324    match s.trim().to_ascii_lowercase().as_str() {
325        "true" => Ok(true),
326        "false" => Ok(false),
327        other => Err(InferError::BadMetadata(format!(
328            "`end2end` is `{other}`, expected true or false"
329        ))),
330    }
331}
332
333/// Parses the Ultralytics `names` metadata field, which appears in two
334/// forms: a stringified Python dict (`"{0: 'person', 1: 'bicycle'}"`, as
335/// captured verbatim from ONNX `metadata_props`) or a JSON object with
336/// string keys (`{"0": "person", "1": "bicycle"}`).
337fn parse_names(s: &str) -> Result<Vec<String>, InferError> {
338    let trimmed = s.trim();
339
340    if let Ok(map) = serde_json::from_str::<BTreeMap<String, String>>(trimmed) {
341        return names_from_str_keyed_map(map);
342    }
343
344    let body = trimmed
345        .strip_prefix('{')
346        .and_then(|s| s.strip_suffix('}'))
347        .ok_or_else(|| InferError::BadNames(format!("not a dict: {trimmed}")))?;
348
349    let mut indexed: BTreeMap<usize, String> = BTreeMap::new();
350    for item in split_top_level_commas(body) {
351        let item = item.trim();
352        if item.is_empty() {
353            continue;
354        }
355        let (key, value) = item
356            .split_once(':')
357            .ok_or_else(|| InferError::BadNames(format!("malformed entry `{item}`")))?;
358        let idx: usize = key
359            .trim()
360            .parse()
361            .map_err(|_| InferError::BadNames(format!("non-numeric key `{key}`")))?;
362        let value = unquote(value.trim())
363            .ok_or_else(|| InferError::BadNames(format!("unquoted value `{value}`")))?;
364        indexed.insert(idx, value);
365    }
366    contiguous_values(indexed)
367}
368
369/// Converts a `names` JSON value (a JSON object keyed by stringified
370/// index) into ordered class names.
371fn names_from_json_value(v: &serde_json::Value) -> Result<Vec<String>, InferError> {
372    let obj = v
373        .as_object()
374        .ok_or_else(|| InferError::BadNames("`names` is not a JSON object".into()))?;
375    let mut map = BTreeMap::new();
376    for (k, val) in obj {
377        let s = val
378            .as_str()
379            .ok_or_else(|| InferError::BadNames(format!("non-string name for key `{k}`")))?;
380        map.insert(k.clone(), s.to_string());
381    }
382    names_from_str_keyed_map(map)
383}
384
385/// Converts a string-keyed names map (JSON form) into ordered class names
386/// by parsing keys as indices, rather than relying on lexicographic string
387/// order (which would misorder e.g. "10" before "2").
388fn names_from_str_keyed_map(map: BTreeMap<String, String>) -> Result<Vec<String>, InferError> {
389    let mut indexed: BTreeMap<usize, String> = BTreeMap::new();
390    for (k, v) in map {
391        let idx: usize = k
392            .parse()
393            .map_err(|_| InferError::BadNames(format!("non-numeric key `{k}`")))?;
394        indexed.insert(idx, v);
395    }
396    contiguous_values(indexed)
397}
398
399/// Requires a `0..n` contiguous index range and returns the values in
400/// index order.
401fn contiguous_values(indexed: BTreeMap<usize, String>) -> Result<Vec<String>, InferError> {
402    // Zero classes is not a degenerate-but-valid model: `feat` would be
403    // `4 + 0`, and the decoder requires at least one class, so a schema
404    // built from it parses and then fails at `DecoderBuilder::build`.
405    // Refuse here, where the message can say what is actually missing.
406    if indexed.is_empty() {
407        return Err(InferError::BadNames("class map is empty".into()));
408    }
409    let n = indexed.len();
410    for i in 0..n {
411        if !indexed.contains_key(&i) {
412            return Err(InferError::BadNames("non-contiguous class indices".into()));
413        }
414    }
415    Ok(indexed.into_values().collect())
416}
417
418/// Splits a dict body on top-level commas, treating `'...'` and `"..."`
419/// spans as opaque so commas inside class names never split an entry.
420fn split_top_level_commas(s: &str) -> Vec<String> {
421    let mut parts = Vec::new();
422    let mut current = String::new();
423    let mut quote: Option<char> = None;
424    for c in s.chars() {
425        match quote {
426            Some(q) => {
427                current.push(c);
428                if c == q {
429                    quote = None;
430                }
431            }
432            None => match c {
433                '\'' | '"' => {
434                    quote = Some(c);
435                    current.push(c);
436                }
437                ',' => parts.push(std::mem::take(&mut current)),
438                _ => current.push(c),
439            },
440        }
441    }
442    if !current.is_empty() {
443        parts.push(current);
444    }
445    parts
446}
447
448/// Strips matching outer `'` or `"` quotes from a value, returning `None`
449/// if the value isn't quoted.
450fn unquote(s: &str) -> Option<String> {
451    let bytes = s.as_bytes();
452    if s.len() >= 2 {
453        let first = bytes[0];
454        let last = bytes[s.len() - 1];
455        if (first == b'\'' && last == b'\'') || (first == b'"' && last == b'"') {
456            return Some(s[1..s.len() - 1].to_string());
457        }
458    }
459    None
460}
461
462/// Classifies a rank-4 model input's spatial layout.
463///
464/// Channel dimension is wherever the shape holds a `3`: index 1 (NCHW) or
465/// index 3 (NHWC). Layout is never assumed from [`ModelSource`] — some
466/// TFLite exports are still NCHW at the tensor level.
467///
468/// Batch must be 1. Every output layout this module classifies is
469/// batch-1 (`classify_pre_nms`, `find_e2e_candidate` and `classify_proto`
470/// all require it), so accepting a batched input would emit a `dshape`
471/// whose `Batch` disagrees with the detection tensor it is paired with.
472#[allow(clippy::type_complexity)]
473fn classify_input_layout(
474    input: &TensorInfo,
475) -> Result<(usize, usize, Vec<(DimName, usize)>), InferError> {
476    debug_assert_eq!(
477        input.shape.len(),
478        4,
479        "caller selects the input by rank-4; see infer_ultralytics_schema"
480    );
481    if input.shape[0] != 1 {
482        return Err(InferError::UnsupportedLayout(format!(
483            "model input `{}` has batch {} but every supported output layout \
484             is batch-1; shape {:?}",
485            input.name, input.shape[0], input.shape
486        )));
487    }
488    if input.shape[1] == 3 {
489        let (h, w) = (input.shape[2], input.shape[3]);
490        Ok((
491            h,
492            w,
493            vec![
494                (DimName::Batch, input.shape[0]),
495                (DimName::NumFeatures, 3),
496                (DimName::Height, h),
497                (DimName::Width, w),
498            ],
499        ))
500    } else if input.shape[3] == 3 {
501        let (h, w) = (input.shape[1], input.shape[2]);
502        Ok((
503            h,
504            w,
505            vec![
506                (DimName::Batch, input.shape[0]),
507                (DimName::Height, h),
508                (DimName::Width, w),
509                (DimName::NumFeatures, 3),
510            ],
511        ))
512    } else {
513        Err(InferError::UnsupportedLayout(format!(
514            "cannot locate the channel dimension (=3) in input shape {:?}",
515            input.shape
516        )))
517    }
518}
519
520/// If `shape` is rank-3 with a unit leading (batch) dim and one of its
521/// remaining two dims equals `expected_anchors`, returns the other dim's
522/// value and its index (1 or 2) within `shape`. `None` if no dim matches,
523/// including when the leading dim isn't 1 — same requirement as
524/// [`find_e2e_candidate`], so a batch > 1 tensor is never silently
525/// mislabeled batch-1.
526fn find_anchors_dim(shape: &[usize], expected_anchors: usize) -> Option<(usize, usize)> {
527    if shape.len() != 3 || shape[0] != 1 {
528        return None;
529    }
530    if shape[2] == expected_anchors {
531        Some((shape[1], 1)) // features candidate at index 1: features-first
532    } else if shape[1] == expected_anchors {
533        Some((shape[2], 2)) // features candidate at index 2: anchors-first
534    } else {
535        None
536    }
537}
538
539/// Classifies the combined pre-NMS detection layout (rule 4): a rank-3
540/// output whose dims contain both `expected_anchors` and `feat` (`4 + nc`).
541///
542/// Every rank-3 output is scanned to completion. A tensor that carries the
543/// anchor count but the wrong feature width is remembered as a fallback
544/// [`InferError::ClassCountMismatch`] rather than returned immediately,
545/// because an export can legitimately publish more than one anchor-shaped
546/// output — a `[1, 4, 8400]` box tensor beside the `[1, 4+nc, 8400]`
547/// detection tensor, or a `[1, k, 8400]` mask-coefficient tensor — and
548/// returning on the first near-miss made the result depend on the order the
549/// runtime happened to report its outputs in. The mismatch is only
550/// surfaced when no output matched, where it is the more useful diagnostic:
551/// a dim equal to the anchor count unambiguously signals pre-NMS intent, so
552/// a feature-width mismatch there is a class-count problem rather than a
553/// layout question.
554#[allow(clippy::type_complexity)]
555fn classify_pre_nms<'a>(
556    rank3_outputs: &[&'a TensorInfo],
557    expected_anchors: usize,
558    feat: usize,
559    nc: usize,
560) -> Result<(&'a TensorInfo, Vec<(DimName, usize)>), InferError> {
561    let mut mismatch: Option<InferError> = None;
562    let mut matched: Option<(&'a TensorInfo, Vec<(DimName, usize)>)> = None;
563    for out in rank3_outputs {
564        let Some((other_dim, feat_idx)) = find_anchors_dim(&out.shape, expected_anchors) else {
565            continue;
566        };
567        if other_dim == feat {
568            let dshape = if feat_idx == 1 {
569                vec![
570                    (DimName::Batch, 1),
571                    (DimName::NumFeatures, feat),
572                    (DimName::NumBoxes, expected_anchors),
573                ]
574            } else {
575                vec![
576                    (DimName::Batch, 1),
577                    (DimName::NumBoxes, expected_anchors),
578                    (DimName::NumFeatures, feat),
579                ]
580            };
581            // Two tensors of the identical detection shape leave nothing to
582            // choose between them but the order the runtime listed them in.
583            if let Some((first, _)) = &matched {
584                return Err(InferError::AmbiguousLayout(format!(
585                    "outputs `{}` and `{}` both match the pre-NMS detection layout \
586                     ({:?}); cannot tell which carries the detections",
587                    first.name, out.name, out.shape
588                )));
589            }
590            matched = Some((out, dshape));
591            continue;
592        }
593        mismatch.get_or_insert_with(|| InferError::ClassCountMismatch {
594            expected: nc,
595            // `feat - nc` is the non-class base width: `4` for detection,
596            // `4 + k` for segmentation (proto channels).
597            found: other_dim.saturating_sub(feat.saturating_sub(nc)),
598        });
599    }
600    if let Some(found) = matched {
601        return Ok(found);
602    }
603    Err(mismatch.unwrap_or_else(|| {
604        InferError::UnsupportedLayout(format!(
605            "no output matches the Ultralytics pre-NMS layout (expected a dim == \
606             {expected_anchors} anchors and a dim == {feat} features [4 + {nc} classes]); \
607             outputs seen: {:?}",
608            rank3_outputs.iter().map(|o| &o.shape).collect::<Vec<_>>()
609        ))
610    }))
611}
612
613/// Finds an end-to-end (YOLO26) candidate: a rank-3 output shaped
614/// `[1, N, feat]`, where `feat` is `6` for detection or `6 + k` for
615/// segmentation (`k` being the proto channel count).
616///
617/// `N` is the export's `max_det`, which Ultralytics bakes into the graph
618/// rather than applying afterwards: the head's `postprocess` runs a TopK
619/// with `k = min(max_det, anchors)` (`nn/modules/head.py`), so the tensor
620/// carries exactly that many rows and no padding. The bound here is
621/// upstream's own clamp, `m.max_det = min(args.max_det, available)` where
622/// `available` is the anchor count summed over strides
623/// (`engine/exporter.py`) — the identical quantity this module computes as
624/// `expected_anchors`. So `N == expected_anchors` is a legitimate
625/// end-to-end shape (any `max_det >= anchors` produces it, verified by
626/// export), not a pre-NMS tensor to be excluded.
627fn find_e2e_candidate<'a>(
628    rank3_outputs: &[&'a TensorInfo],
629    expected_anchors: usize,
630    feat: usize,
631) -> Result<Option<(&'a TensorInfo, usize)>, InferError> {
632    let matches: Vec<(&'a TensorInfo, usize)> = rank3_outputs
633        .iter()
634        .filter_map(|o| {
635            let s = &o.shape;
636            debug_assert_eq!(s.len(), 3, "caller filters to rank-3 outputs");
637            // `N > 0`: a zero-row tensor satisfies every other condition but
638            // produces a schema the decoder rejects for a zero dimension.
639            (s[0] == 1 && s[2] == feat && s[1] > 0 && s[1] <= expected_anchors)
640                .then_some((*o, s[1]))
641        })
642        .collect();
643    // Two tensors of the same end-to-end shape leave nothing to choose
644    // between them but the order the runtime listed them in. This module
645    // reports ambiguity rather than picking.
646    if matches.len() > 1 {
647        return Err(InferError::AmbiguousLayout(format!(
648            "{} outputs match the end-to-end layout [1, N, {feat}]: {:?}; \
649             cannot tell which carries the detections",
650            matches.len(),
651            matches.iter().map(|(o, _)| &o.name).collect::<Vec<_>>()
652        )));
653    }
654    Ok(matches.into_iter().next())
655}
656
657/// Classifies a rank-4 proto tensor (rule 1): the mask-prototype output
658/// paired with a segmentation detection tensor. Expected spatial dims are
659/// `(H/4, W/4)` (Ultralytics proto stride 4); the remaining non-batch dim
660/// is the proto channel count `k`, never hardcoded. Supports both
661/// `[1, k, H/4, W/4]` (NCHW) and `[1, H/4, W/4, k]` (NHWC).
662///
663/// The NCHW check runs first, so a shape where `k` happens to equal one of
664/// the spatial dims (e.g. `[1, 160, 160, 160]` on a 640x640 input, where
665/// `k == H/4 == W/4`) is shape-undecidable and deliberately resolves NCHW
666/// rather than erroring — no real Ultralytics export produces a proto with
667/// that many channels, so this tie-break is a pragmatic default, not a
668/// verified convention.
669fn classify_proto(
670    proto: &TensorInfo,
671    h: usize,
672    w: usize,
673) -> Result<(usize, Vec<(DimName, usize)>), InferError> {
674    let (eh, ew) = (h / 4, w / 4);
675    let s = &proto.shape;
676    if s.len() != 4 || s[0] != 1 {
677        return Err(InferError::UnsupportedLayout(format!(
678            "expected a rank-4 proto tensor with batch 1, found shape {s:?}"
679        )));
680    }
681    if s[2] == eh && s[3] == ew {
682        let k = s[1];
683        Ok((
684            k,
685            vec![
686                (DimName::Batch, 1),
687                (DimName::NumProtos, k),
688                (DimName::Height, eh),
689                (DimName::Width, ew),
690            ],
691        ))
692    } else if s[1] == eh && s[2] == ew {
693        let k = s[3];
694        Ok((
695            k,
696            vec![
697                (DimName::Batch, 1),
698                (DimName::Height, eh),
699                (DimName::Width, ew),
700                (DimName::NumProtos, k),
701            ],
702        ))
703    } else {
704        Err(InferError::UnsupportedLayout(format!(
705            "proto tensor shape {s:?} doesn't match the expected ({eh}, {ew}) spatial dims"
706        )))
707    }
708}
709
710/// Default field values for a [`LogicalOutput`] whose semantic fields are
711/// filled in by the caller; every field not overridden is `None` or empty.
712fn default_logical() -> LogicalOutput {
713    LogicalOutput {
714        name: None,
715        type_: None,
716        shape: Vec::new(),
717        dshape: Vec::new(),
718        decoder: None,
719        encoding: None,
720        score_format: None,
721        normalized: None,
722        anchors: None,
723        stride: None,
724        dtype: None,
725        quantization: None,
726        outputs: Vec::new(),
727        activation_applied: None,
728        activation_required: None,
729        semantic: Default::default(),
730    }
731}
732
733/// Rejects an integer-dtype boundary tensor that carries no quantization
734/// parameters — there would be no way to dequantize it, so this is always
735/// an error rather than a value silently passed through as raw integers.
736///
737/// Real Ultralytics int8 TFLite exports never hit this: the int8 graph is
738/// wrapped in dequant/quant ops, so the boundary tensors the runtime
739/// reports are float32 (see `infer_int8_tflite_classifies_like_float32`).
740/// This guard only fires for a genuinely quantized boundary, which no
741/// captured fixture exhibits but a future export convention might.
742/// Per-channel quantization is refused here for the same reason: the
743/// decoder supports per-tensor quantization only, so a per-channel boundary
744/// would produce a schema that parses and then fails at
745/// `DecoderBuilder::build` with an error naming the decoder rather than the
746/// signal that caused it. Refusing at the boundary names the real problem
747/// while the caller still has the tensor in hand.
748fn check_quantized_boundary(tensor: &TensorInfo) -> Result<(), InferError> {
749    if tensor.dtype.is_integer() && tensor.quantization.is_none() {
750        return Err(InferError::UnsupportedLayout(format!(
751            "output `{}` has integer dtype {:?} but no quantization params",
752            tensor.name, tensor.dtype
753        )));
754    }
755    if let Some(q) = &tensor.quantization {
756        if q.scale.len() > 1 {
757            return Err(InferError::UnsupportedQuantization(format!(
758                "output `{}` carries per-channel quantization ({} scales); \
759                 the decoder supports per-tensor quantization only",
760                tensor.name,
761                q.scale.len()
762            )));
763        }
764        // Exactly one scale, not merely "not per-channel". An empty scale
765        // list passed the check above and produced `{"scale": []}`, which
766        // the decoder reads as a zero scale rather than rejecting.
767        if q.scale.is_empty() {
768            return Err(InferError::UnsupportedQuantization(format!(
769                "output `{}` is quantized but carries no scale",
770                tensor.name
771            )));
772        }
773        // Zero points are optional (symmetric quantization), but when
774        // present must pair with the scale; a longer list was silently
775        // truncated to its first element.
776        if let Some(zp) = &q.zero_point {
777            if zp.len() != q.scale.len() {
778                return Err(InferError::UnsupportedQuantization(format!(
779                    "output `{}` has {} scale(s) but {} zero point(s)",
780                    tensor.name,
781                    q.scale.len(),
782                    zp.len()
783                )));
784            }
785        }
786    }
787    Ok(())
788}
789
790/// Infers an Ultralytics YOLO [`SchemaV2`] from raw model I/O signals and
791/// metadata.
792///
793/// Detection and segmentation are both supported. Segmentation is
794/// recognised either by metadata (`task == segment`) or by the presence of
795/// a rank-4 proto tensor; the two signals are cross-checked (rule 5) so a
796/// `segment` task with no proto, or a `detect` task with a proto present,
797/// is always an error rather than a guess.
798///
799/// # Examples
800///
801/// A two-class YOLOv8 detection export, as an inference runtime would
802/// report it — a rank-4 input, one `[1, 4 + nc, anchors]` output, and the
803/// metadata the exporter embedded:
804///
805/// ```
806/// use std::collections::BTreeMap;
807///
808/// use edgefirst_decoder::schema::DType;
809/// use edgefirst_decoder::{
810///     infer_ultralytics_schema, ModelSignals, ModelSource, TensorInfo,
811/// };
812///
813/// let tensor = |name: &str, shape: &[usize]| TensorInfo {
814///     name: name.to_string(),
815///     shape: shape.to_vec(),
816///     dtype: DType::Float32,
817///     quantization: None,
818/// };
819///
820/// let mut metadata = BTreeMap::new();
821/// metadata.insert("task".to_string(), "detect".to_string());
822/// metadata.insert("end2end".to_string(), "False".to_string());
823/// metadata.insert("names".to_string(), "{0: 'person', 1: 'bicycle'}".to_string());
824///
825/// let signals = ModelSignals {
826///     source: ModelSource::Onnx,
827///     inputs: vec![tensor("images", &[1, 3, 640, 640])],
828///     // 8400 anchors for a 640x640 input (80² + 40² + 20²);
829///     // 6 features = 4 box coordinates + 2 classes.
830///     outputs: vec![tensor("output0", &[1, 6, 8400])],
831///     metadata,
832/// };
833///
834/// let inferred = infer_ultralytics_schema(&signals)?;
835/// assert_eq!(inferred.labels, ["person", "bicycle"]);
836/// assert_eq!(inferred.description, "Ultralytics YOLOv8/11 detect, 2 classes");
837///
838/// // Hand the schema to `DecoderBuilder` as `edgefirst.json` v2 JSON.
839/// let config_json = serde_json::to_string(&inferred.schema).unwrap();
840/// assert!(config_json.contains("\"decoder_version\":\"yolov8\""));
841/// # Ok::<(), edgefirst_decoder::InferError>(())
842/// ```
843///
844/// Signals that carry no Ultralytics metadata are refused with a typed
845/// error rather than a best-effort guess:
846///
847/// ```
848/// use std::collections::BTreeMap;
849///
850/// use edgefirst_decoder::{
851///     infer_ultralytics_schema, InferError, ModelSignals, ModelSource,
852/// };
853///
854/// let signals = ModelSignals {
855///     source: ModelSource::Onnx,
856///     inputs: Vec::new(),
857///     outputs: Vec::new(),
858///     metadata: BTreeMap::new(),
859/// };
860///
861/// assert!(matches!(
862///     infer_ultralytics_schema(&signals),
863///     Err(InferError::NotUltralytics(_))
864/// ));
865/// ```
866pub fn infer_ultralytics_schema(signals: &ModelSignals) -> Result<InferredSchema, InferError> {
867    infer_from_meta(signals, UltralyticsMeta::from_metadata(&signals.metadata)?)
868}
869
870/// Infers the schema, with no labels, from a class count other than the one
871/// `names` gives, and returns that count. `classes`, the count the caller
872/// states, is tried first; otherwise each anchor-shaped output proposes a
873/// count and exactly one must lay out, else the layout is ambiguous.
874pub(crate) fn infer_ultralytics_schema_ignoring_names(
875    signals: &ModelSignals,
876    classes: Option<usize>,
877) -> Result<(InferredSchema, usize), InferError> {
878    let meta = UltralyticsMeta::from_metadata(&signals.metadata)?;
879    let with_count = |nc: usize| {
880        let names = (0..nc).map(|i| format!("class{i}")).collect();
881        let candidate = UltralyticsMeta {
882            names,
883            task: meta.task,
884            end2end: meta.end2end,
885        };
886        infer_from_meta(signals, candidate)
887    };
888    const REMEDY: &str = "supply a labels override naming every class, or describe the \
889                          outputs in an edgefirst.json beside the model";
890    let stated = classes
891        .filter(|&nc| nc > 0)
892        .and_then(|nc| with_count(nc).ok());
893    let mut found: Option<InferredSchema> = stated;
894    if found.is_none() {
895        for nc in candidate_class_counts(signals)? {
896            let inferred = match with_count(nc) {
897                Ok(inferred) => inferred,
898                Err(InferError::AmbiguousLayout(why)) => {
899                    return Err(InferError::AmbiguousLayout(format!("{why}; {REMEDY}")))
900                }
901                Err(_) => continue,
902            };
903            if let Some(first) = &found {
904                return Err(InferError::AmbiguousLayout(format!(
905                    "`names` miscounts the classes and the outputs allow both {} and {nc}; \
906                     {REMEDY}",
907                    first.labels.len()
908                )));
909            }
910            found = Some(inferred);
911        }
912    }
913    let mut inferred = found.ok_or_else(|| {
914        InferError::UnsupportedLayout("no output carries a usable class count".into())
915    })?;
916    let nc = std::mem::take(&mut inferred.labels).len();
917    Ok((inferred, nc))
918}
919
920/// The class counts the anchor-shaped outputs imply, ascending.
921fn candidate_class_counts(signals: &ModelSignals) -> Result<BTreeSet<usize>, InferError> {
922    let input = signals
923        .inputs
924        .iter()
925        .find(|t| t.shape.len() == 4)
926        .ok_or_else(|| InferError::UnsupportedLayout("no rank-4 model input found".into()))?;
927    let (h, w, _) = classify_input_layout(input)?;
928    let k = match signals.outputs.iter().find(|o| o.shape.len() == 4) {
929        Some(proto) => classify_proto(proto, h, w)?.0,
930        None => 0,
931    };
932    Ok(signals
933        .outputs
934        .iter()
935        .filter_map(|o| find_anchors_dim(&o.shape, anchor_count(h, w)))
936        .filter_map(|(features, _)| features.checked_sub(4 + k))
937        .filter(|&nc| nc > 0)
938        .collect())
939}
940
941/// Anchors over the stride 8, 16 and 32 grids.
942fn anchor_count(h: usize, w: usize) -> usize {
943    (h / 8) * (w / 8) + (h / 16) * (w / 16) + (h / 32) * (w / 32)
944}
945
946fn infer_from_meta(
947    signals: &ModelSignals,
948    meta: UltralyticsMeta,
949) -> Result<InferredSchema, InferError> {
950    let proto = signals.outputs.iter().find(|o| o.shape.len() == 4);
951    match (meta.task, proto.is_some()) {
952        (Some(Task::Segment), false) => {
953            return Err(InferError::UnsupportedLayout(
954                "segment task but no proto tensor".into(),
955            ));
956        }
957        (Some(Task::Detect), true) => {
958            return Err(InferError::UnsupportedLayout(
959                "detect task but proto tensor present".into(),
960            ));
961        }
962        _ => {}
963    }
964    let is_segment = meta.task == Some(Task::Segment) || proto.is_some();
965
966    let input = signals
967        .inputs
968        .iter()
969        .find(|t| t.shape.len() == 4)
970        .ok_or_else(|| {
971            InferError::UnsupportedLayout(format!(
972                "no rank-4 model input found; inputs seen: {:?}",
973                signals.inputs.iter().map(|t| &t.shape).collect::<Vec<_>>()
974            ))
975        })?;
976    let (h, w, input_dshape) = classify_input_layout(input)?;
977
978    let expected_anchors = anchor_count(h, w);
979    let nc = meta.names.len();
980
981    // Segmentation detection features carry `k` extra mask coefficients
982    // (`4 + nc + k`); the end-to-end feature width grows by the same `k`
983    // (`6 + k`). `k` comes from the proto tensor itself (rule 1), never
984    // hardcoded, so both widths fall back to the plain detection case when
985    // there's no proto.
986    // `is_segment` is already true whenever a proto is present (the two
987    // are cross-checked above), so the proto's presence alone decides here.
988    let (k, proto_dshape) = match proto {
989        Some(proto) => {
990            let (k, dshape) = classify_proto(proto, h, w)?;
991            (k, Some(dshape))
992        }
993        None => (0, None),
994    };
995    let feat = 4 + nc + k;
996    let e2e_feat = 6 + k;
997
998    let rank3_outputs: Vec<&TensorInfo> = signals
999        .outputs
1000        .iter()
1001        .filter(|o| o.shape.len() == 3)
1002        .collect();
1003
1004    // The `end2end` metadata flag is the primary end-to-end signal (present
1005    // on every real export). The shape-based e2e fallback applies only
1006    // when the flag is absent (stripped metadata).
1007    let (out, dshape, version) = match meta.end2end {
1008        Some(true) => {
1009            let (out, n) = find_e2e_candidate(&rank3_outputs, expected_anchors, e2e_feat)?
1010                .ok_or_else(|| {
1011                    InferError::UnsupportedLayout(format!(
1012                        "metadata declares end2end=true but no output matches the YOLO26 \
1013                         end-to-end layout [1, N<={expected_anchors}, {e2e_feat}]; \
1014                         outputs seen: {:?}",
1015                        rank3_outputs.iter().map(|o| &o.shape).collect::<Vec<_>>()
1016                    ))
1017                })?;
1018            (
1019                out,
1020                vec![
1021                    (DimName::Batch, 1),
1022                    (DimName::NumBoxes, n),
1023                    (DimName::NumFeatures, e2e_feat),
1024                ],
1025                DecoderVersion::Yolo26,
1026            )
1027        }
1028        Some(false) => {
1029            let (out, dshape) = classify_pre_nms(&rank3_outputs, expected_anchors, feat, nc)?;
1030            (out, dshape, DecoderVersion::Yolov8)
1031        }
1032        None => match classify_pre_nms(&rank3_outputs, expected_anchors, feat, nc) {
1033            Ok((out, dshape)) => (out, dshape, DecoderVersion::Yolov8),
1034            Err(err @ InferError::ClassCountMismatch { .. }) => return Err(err),
1035            Err(_) => {
1036                let (out, n) = find_e2e_candidate(&rank3_outputs, expected_anchors, e2e_feat)?
1037                    .ok_or_else(|| {
1038                        InferError::UnsupportedLayout(format!(
1039                            "no output matches a known Ultralytics detection layout \
1040                             (pre-NMS or end-to-end); outputs seen: {:?}",
1041                            rank3_outputs.iter().map(|o| &o.shape).collect::<Vec<_>>()
1042                        ))
1043                    })?;
1044                (
1045                    out,
1046                    vec![
1047                        (DimName::Batch, 1),
1048                        (DimName::NumBoxes, n),
1049                        (DimName::NumFeatures, e2e_feat),
1050                    ],
1051                    DecoderVersion::Yolo26,
1052                )
1053            }
1054        },
1055    };
1056
1057    check_quantized_boundary(out)?;
1058    if let Some(proto) = proto {
1059        check_quantized_boundary(proto)?;
1060    }
1061
1062    // Box normalization is the one field shape cannot reveal -- it follows
1063    // the exporter. ONNX and TFLite are pinned by runtime magnitude
1064    // measurement (see testdata/infer/NOTES.md answer 5: 637.25 px vs
1065    // 0.9957 on the same image); CoreML resolves alongside ONNX on both
1066    // bases. Source-tracing the Ultralytics exporter shows the `1/w`
1067    // normalize lives only in `IOSDetectModel`, used solely on the
1068    // `nms=True` path, so a plain no-NMS CoreML export passes the raw model
1069    // through pixel-space exactly as ONNX does; the same runtime check run
1070    // against a matched fp16 ONNX pair agrees, with x reaching 640.25 on a
1071    // 640-pixel input and every coordinate statistic matching within half a
1072    // pixel -- fp16 rounding, where a convention difference would be a
1073    // factor of 640 (see testdata/infer/NOTES.md for both, and the captured
1074    // fixture). Guessing it wrong scales every box by the input size, which
1075    // is why `Other` is refused rather than defaulted: everywhere else this
1076    // module errors on ambiguity, and this is the field whose corruption
1077    // `tests/infer_builder.rs` exists to pin.
1078    let normalized = match signals.source {
1079        ModelSource::TfLite => true,
1080        ModelSource::Onnx | ModelSource::CoreMl => false,
1081        ModelSource::Other => {
1082            return Err(InferError::UnknownBoxConvention);
1083        }
1084    };
1085    let det = LogicalOutput {
1086        name: Some(out.name.clone()),
1087        // `Detections` (plural) is the schema vocabulary's fully-decoded
1088        // post-NMS type; `Detection` is the anchor-grid output that still
1089        // needs decoding. Both collapse to the same legacy path today
1090        // (schema.rs `logical_to_legacy`), but this JSON now leaves the
1091        // process through the C and Python bindings, so it has to name the
1092        // tensor it actually describes.
1093        type_: Some(match version {
1094            DecoderVersion::Yolo26 => LogicalType::Detections,
1095            _ => LogicalType::Detection,
1096        }),
1097        shape: out.shape.clone(),
1098        dshape,
1099        decoder: Some(DecoderKind::Ultralytics),
1100        encoding: Some(BoxEncoding::Direct), // graph already decoded DFL
1101        score_format: Some(ScoreFormat::PerClass),
1102        normalized: Some(normalized),
1103        anchors: None,
1104        stride: None,
1105        dtype: Some(out.dtype),
1106        quantization: out.quantization.clone(),
1107        outputs: Vec::new(),
1108        activation_applied: None,
1109        activation_required: None,
1110        semantic: Default::default(),
1111    };
1112
1113    let mut outputs = vec![det];
1114    if let (Some(proto), Some(proto_dshape)) = (proto, proto_dshape) {
1115        outputs.push(LogicalOutput {
1116            name: Some(proto.name.clone()),
1117            type_: Some(LogicalType::Protos),
1118            shape: proto.shape.clone(),
1119            dshape: proto_dshape,
1120            decoder: Some(DecoderKind::Ultralytics),
1121            dtype: Some(proto.dtype),
1122            quantization: proto.quantization.clone(),
1123            ..default_logical()
1124        });
1125    }
1126
1127    let schema = SchemaV2 {
1128        schema_version: 2,
1129        input: Some(InputSpec {
1130            shape: input.shape.clone(),
1131            dshape: input_dshape,
1132            cameraadaptor: Some("rgb".into()),
1133            preprocessing: None,
1134        }),
1135        outputs,
1136        // Ultralytics runs NMS with `agnostic=False`, i.e. class-aware; stated
1137        // explicitly so the schema describes the model rather than relying on
1138        // the builder's fallback. An explicit `.with_nms(..)` still overrides.
1139        //
1140        // YOLO26 end-to-end heads do their NMS in-graph, so they carry none:
1141        // this field describes what the *decoder* must add.
1142        nms: match version {
1143            DecoderVersion::Yolo26 => None,
1144            _ => Some(NmsMode::ClassAware),
1145        },
1146        nms_multi_label: None,
1147        decoder_version: Some(version),
1148    };
1149
1150    let version_label = if version == DecoderVersion::Yolo26 {
1151        "YOLO26"
1152    } else {
1153        "YOLOv8/11"
1154    };
1155    let task_label = if is_segment { "segment" } else { "detect" };
1156    let description = format!("Ultralytics {version_label} {task_label}, {nc} classes");
1157
1158    Ok(InferredSchema {
1159        schema,
1160        labels: meta.names,
1161        description,
1162    })
1163}
1164
1165#[cfg(test)]
1166mod tests {
1167    use super::*;
1168
1169    /// Root of the `testdata/infer` fixture directory.
1170    ///
1171    /// `EDGEFIRST_TESTDATA_DIR` wins when set -- `scripts/on-target-test.sh`
1172    /// and CI both mirror `crates/decoder/testdata/infer` into
1173    /// `$EDGEFIRST_TESTDATA_DIR/infer` (see `on-target-test.sh`'s testdata
1174    /// sync and `.github/workflows/test.yml`'s "Upload testdata as
1175    /// artifact" step). The manifest-relative fallback is baked in at
1176    /// **compile** time, so a cross-compiled binary carries the build
1177    /// host's absolute path and can never find fixtures on a different
1178    /// target -- same reasoning as `decoder_from_edgefirst_json.rs`'s
1179    /// `testdata_root()`.
1180    fn infer_fixture_dir() -> std::path::PathBuf {
1181        std::env::var("EDGEFIRST_TESTDATA_DIR")
1182            .map(std::path::PathBuf::from)
1183            .unwrap_or_else(|_| {
1184                std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("testdata")
1185            })
1186            .join("infer")
1187    }
1188
1189    #[test]
1190    fn parse_names_python_dict_repr() {
1191        let s = "{0: 'person', 1: 'bicycle', 2: \"fire hydrant\"}";
1192        assert_eq!(
1193            parse_names(s).unwrap(),
1194            vec!["person", "bicycle", "fire hydrant"]
1195        );
1196    }
1197
1198    #[test]
1199    fn parse_names_json_object() {
1200        let s = r#"{"0": "person", "1": "bicycle"}"#;
1201        assert_eq!(parse_names(s).unwrap(), vec!["person", "bicycle"]);
1202    }
1203
1204    #[test]
1205    fn parse_names_orders_by_index_not_insertion() {
1206        let s = "{1: 'b', 0: 'a'}";
1207        assert_eq!(parse_names(s).unwrap(), vec!["a", "b"]);
1208    }
1209
1210    #[test]
1211    fn parse_names_rejects_garbage() {
1212        assert!(matches!(
1213            parse_names("not a dict"),
1214            Err(InferError::BadNames(_))
1215        ));
1216    }
1217
1218    #[test]
1219    fn parse_names_rejects_non_contiguous_indices() {
1220        // Python dict-repr form: index 1 is missing.
1221        match parse_names("{0: 'a', 2: 'b'}") {
1222            Err(InferError::BadNames(msg)) => {
1223                assert!(
1224                    msg.contains("non-contiguous"),
1225                    "unexpected BadNames message: {msg}"
1226                );
1227            }
1228            other => panic!("expected BadNames, got {other:?}"),
1229        }
1230
1231        // JSON-object form: same gap.
1232        match parse_names(r#"{"0": "a", "2": "b"}"#) {
1233            Err(InferError::BadNames(msg)) => {
1234                assert!(
1235                    msg.contains("non-contiguous"),
1236                    "unexpected BadNames message: {msg}"
1237                );
1238            }
1239            other => panic!("expected BadNames, got {other:?}"),
1240        }
1241    }
1242
1243    #[test]
1244    fn meta_from_onnx_props() {
1245        let mut m = BTreeMap::new();
1246        m.insert("names".into(), "{0: 'person'}".into());
1247        m.insert("task".into(), "detect".into());
1248        let u = UltralyticsMeta::from_metadata(&m).unwrap();
1249        assert_eq!(u.names, vec!["person"]);
1250        assert_eq!(u.task, Some(Task::Detect));
1251    }
1252
1253    #[test]
1254    fn meta_task_pose_rejected() {
1255        let mut m = BTreeMap::new();
1256        m.insert("names".into(), "{0: 'person'}".into());
1257        m.insert("task".into(), "pose".into());
1258        assert!(matches!(
1259            UltralyticsMeta::from_metadata(&m),
1260            Err(InferError::UnsupportedTask(_))
1261        ));
1262    }
1263
1264    #[test]
1265    fn meta_missing_names_is_not_ultralytics() {
1266        assert!(matches!(
1267            UltralyticsMeta::from_metadata(&BTreeMap::new()),
1268            Err(InferError::NotUltralytics(_))
1269        ));
1270    }
1271
1272    #[test]
1273    fn meta_end2end_true_or_false_from_onnx_flat_props() {
1274        let mut m = BTreeMap::new();
1275        m.insert("names".into(), "{0: 'person'}".into());
1276        m.insert("end2end".into(), "True".into());
1277        assert_eq!(
1278            UltralyticsMeta::from_metadata(&m).unwrap().end2end,
1279            Some(true)
1280        );
1281
1282        m.insert("end2end".into(), "False".into());
1283        assert_eq!(
1284            UltralyticsMeta::from_metadata(&m).unwrap().end2end,
1285            Some(false)
1286        );
1287    }
1288
1289    /// Loads a captured Task-0 fixture's `metadata` map, exactly as
1290    /// `ModelSignals::metadata` would carry it: raw string key/values,
1291    /// regardless of whether the source was ONNX (flat props) or TFLite
1292    /// (single `metadata.json` envelope entry).
1293    fn load_fixture_metadata(name: &str) -> BTreeMap<String, String> {
1294        let path = infer_fixture_dir().join(format!("{name}.signals.json"));
1295        let content = std::fs::read_to_string(&path)
1296            .unwrap_or_else(|e| panic!("failed to read fixture {}: {e}", path.display()));
1297        let json: serde_json::Value =
1298            serde_json::from_str(&content).expect("fixture is valid JSON");
1299        json["metadata"]
1300            .as_object()
1301            .expect("fixture has a metadata object")
1302            .iter()
1303            .map(|(k, v)| {
1304                (
1305                    k.clone(),
1306                    v.as_str().expect("metadata value is a string").to_string(),
1307                )
1308            })
1309            .collect()
1310    }
1311
1312    #[test]
1313    fn meta_from_real_tflite_metadata_json_envelope() {
1314        let meta = load_fixture_metadata("yolov8n_float32");
1315        let u = UltralyticsMeta::from_metadata(&meta).unwrap();
1316        assert_eq!(u.names.len(), 80);
1317        assert_eq!(u.task, Some(Task::Detect));
1318    }
1319
1320    #[test]
1321    fn meta_end2end_true_on_real_yolo26_tflite_fixture() {
1322        let meta = load_fixture_metadata("yolo26n_float32");
1323        let u = UltralyticsMeta::from_metadata(&meta).unwrap();
1324        assert_eq!(u.end2end, Some(true));
1325    }
1326
1327    #[test]
1328    fn meta_end2end_false_on_real_yolov8_onnx_fixture() {
1329        let meta = load_fixture_metadata("yolov8n");
1330        let u = UltralyticsMeta::from_metadata(&meta).unwrap();
1331        assert_eq!(u.end2end, Some(false));
1332    }
1333
1334    /// Loads a captured Task-0 fixture into `ModelSignals`, exactly as the
1335    /// inference runtime would report it: tensor names/shapes/dtypes plus
1336    /// the raw metadata map.
1337    fn signals_from_fixture(name: &str) -> ModelSignals {
1338        let path = infer_fixture_dir().join(format!("{name}.signals.json"));
1339        let content = std::fs::read_to_string(&path)
1340            .unwrap_or_else(|e| panic!("failed to read fixture {}: {e}", path.display()));
1341        let json: serde_json::Value =
1342            serde_json::from_str(&content).expect("fixture is valid JSON");
1343
1344        let source = match json["source"]
1345            .as_str()
1346            .expect("fixture `source` is a string")
1347        {
1348            "onnx" => ModelSource::Onnx,
1349            "tflite" => ModelSource::TfLite,
1350            "coreml" => ModelSource::CoreMl,
1351            other => panic!("fixture {name}: unknown source `{other}`"),
1352        };
1353
1354        fn parse_dtype(s: &str) -> DType {
1355            match s {
1356                "float32" => DType::Float32,
1357                "float16" => DType::Float16,
1358                "int8" => DType::Int8,
1359                "uint8" => DType::Uint8,
1360                "int16" => DType::Int16,
1361                "uint16" => DType::Uint16,
1362                "int32" => DType::Int32,
1363                "uint32" => DType::Uint32,
1364                other => panic!("unknown dtype `{other}`"),
1365            }
1366        }
1367
1368        fn parse_tensor(v: &serde_json::Value) -> TensorInfo {
1369            TensorInfo {
1370                name: v["name"]
1371                    .as_str()
1372                    .expect("tensor `name` is a string")
1373                    .to_string(),
1374                shape: v["shape"]
1375                    .as_array()
1376                    .expect("tensor `shape` is an array")
1377                    .iter()
1378                    .map(|d| d.as_u64().expect("shape dim is a number") as usize)
1379                    .collect(),
1380                dtype: parse_dtype(v["dtype"].as_str().expect("tensor `dtype` is a string")),
1381                // Every captured fixture reports `"quantization": null`
1382                // (see testdata/infer/NOTES.md); real quantized signals
1383                // are out of scope for this task.
1384                quantization: None,
1385            }
1386        }
1387
1388        let inputs = json["inputs"]
1389            .as_array()
1390            .expect("fixture `inputs` is an array")
1391            .iter()
1392            .map(parse_tensor)
1393            .collect();
1394        let outputs = json["outputs"]
1395            .as_array()
1396            .expect("fixture `outputs` is an array")
1397            .iter()
1398            .map(parse_tensor)
1399            .collect();
1400        let metadata = json["metadata"]
1401            .as_object()
1402            .expect("fixture `metadata` is an object")
1403            .iter()
1404            .map(|(k, v)| {
1405                (
1406                    k.clone(),
1407                    v.as_str().expect("metadata value is a string").to_string(),
1408                )
1409            })
1410            .collect();
1411
1412        ModelSignals {
1413            source,
1414            inputs,
1415            outputs,
1416            metadata,
1417        }
1418    }
1419
1420    /// Builds a synthetic Ultralytics-style metadata map (Python dict-repr
1421    /// `names`, `task`, `end2end`) for hand-built layout tests that don't
1422    /// need a full captured fixture.
1423    fn synthetic_metadata(nc: usize, task: &str, end2end: &str) -> BTreeMap<String, String> {
1424        let names = (0..nc)
1425            .map(|i| format!("{i}: 'c{i}'"))
1426            .collect::<Vec<_>>()
1427            .join(", ");
1428        let mut m = BTreeMap::new();
1429        m.insert("names".into(), format!("{{{names}}}"));
1430        m.insert("task".into(), task.into());
1431        m.insert("end2end".into(), end2end.into());
1432        m
1433    }
1434
1435    #[test]
1436    fn infer_transposed_det_layout_classified() {
1437        // Anchors-first layout ([1, A, 4+nc]), as produced by older
1438        // Ultralytics TFLite exporters. Rule 4 must still classify this
1439        // correctly from tensor shape alone, and the builder's dshape-order
1440        // e2e-inference pitfall must be guarded by an explicit
1441        // decoder_version.
1442        let s = ModelSignals {
1443            source: ModelSource::TfLite,
1444            inputs: vec![TensorInfo {
1445                name: "images".into(),
1446                shape: vec![1, 640, 640, 3],
1447                dtype: DType::Float32,
1448                quantization: None,
1449            }],
1450            outputs: vec![TensorInfo {
1451                name: "output0".into(),
1452                shape: vec![1, 8400, 84],
1453                dtype: DType::Float32,
1454                quantization: None,
1455            }],
1456            metadata: synthetic_metadata(80, "detect", "False"),
1457        };
1458        let r = infer_ultralytics_schema(&s).unwrap();
1459        let o = &r.schema.outputs[0];
1460        assert_eq!(
1461            o.dshape,
1462            vec![
1463                (DimName::Batch, 1),
1464                (DimName::NumBoxes, 8400),
1465                (DimName::NumFeatures, 84),
1466            ]
1467        );
1468        assert_eq!(r.schema.decoder_version, Some(DecoderVersion::Yolov8));
1469    }
1470
1471    #[test]
1472    fn infer_yolo26n_onnx_end_to_end() {
1473        let s = signals_from_fixture("yolo26n");
1474        let r = infer_ultralytics_schema(&s).unwrap();
1475        assert_eq!(r.schema.decoder_version, Some(DecoderVersion::Yolo26));
1476        assert_eq!(r.schema.outputs[0].shape[2], 6);
1477    }
1478
1479    #[test]
1480    fn infer_class_count_mismatch_rejected() {
1481        let mut s = signals_from_fixture("yolov8n");
1482        s.metadata.insert("names".into(), "{0: 'only-one'}".into()); // nc=1, f=84
1483        assert!(matches!(
1484            infer_ultralytics_schema(&s),
1485            Err(InferError::ClassCountMismatch {
1486                expected: 1,
1487                found: 80
1488            })
1489        ));
1490    }
1491
1492    #[test]
1493    fn infer_seg_class_count_mismatch_rejected() {
1494        // Covers the generalized ClassCountMismatch math for segmentation,
1495        // where the non-class base width is `4 + k` (proto channels), not
1496        // the plain-detection `4`. Real yolov8n-seg det tensor is
1497        // [1, 116, 8400] (features-first, feat = 4+80+32); truncating
1498        // `names` to 1 class makes the expected feat 4+1+32=37, so the
1499        // real 116-feature tensor mismatches with found = 116-(37-1) = 80
1500        // — the true class count the features imply, unaffected by the
1501        // metadata truncation.
1502        let mut s = signals_from_fixture("yolov8n-seg");
1503        s.metadata.insert("names".into(), "{0: 'only-one'}".into()); // nc=1
1504        assert!(matches!(
1505            infer_ultralytics_schema(&s),
1506            Err(InferError::ClassCountMismatch {
1507                expected: 1,
1508                found: 80
1509            })
1510        ));
1511    }
1512
1513    #[test]
1514    fn infer_end2end_flag_shape_conflict_rejected() {
1515        // The end2end flag is the primary signal: a real yolov8n ONNX
1516        // export with its pre-NMS [1, 84, 8400] output, but with `end2end`
1517        // forced to "True", must never be guessed into an end-to-end
1518        // classification — flag/shape disagreement is always an error.
1519        let mut s = signals_from_fixture("yolov8n");
1520        s.metadata.insert("end2end".into(), "True".into());
1521        assert!(matches!(
1522            infer_ultralytics_schema(&s),
1523            Err(InferError::UnsupportedLayout(_))
1524        ));
1525    }
1526
1527    #[test]
1528    fn infer_yolov8n_seg_onnx() {
1529        let r = infer_ultralytics_schema(&signals_from_fixture("yolov8n-seg")).unwrap();
1530        assert_eq!(r.schema.outputs.len(), 2);
1531        let proto = r
1532            .schema
1533            .outputs
1534            .iter()
1535            .find(|o| o.type_ == Some(LogicalType::Protos))
1536            .unwrap();
1537        assert_eq!(proto.dshape[1].0, DimName::NumProtos); // NCHW
1538        let det = r
1539            .schema
1540            .outputs
1541            .iter()
1542            .find(|o| o.type_ == Some(LogicalType::Detection))
1543            .unwrap();
1544        assert!(det.shape.contains(&116)); // 4+80+32
1545    }
1546
1547    #[test]
1548    fn infer_yolov8n_seg_tflite_proto_nchw() {
1549        // The real yolov8n-seg_float32 TFLite export's proto is
1550        // [1, 32, 160, 160], NCHW-ordered — same as its ONNX counterpart.
1551        // An earlier draft of this test assumed TFLite seg exports carry a
1552        // NHWC-transposed proto; no captured fixture confirms that (see
1553        // infer_nhwc_proto_layout_classified for the synthetic NHWC coverage
1554        // the shape-driven classifier still supports).
1555        let r = infer_ultralytics_schema(&signals_from_fixture("yolov8n-seg_float32")).unwrap();
1556        let proto = r
1557            .schema
1558            .outputs
1559            .iter()
1560            .find(|o| o.type_ == Some(LogicalType::Protos))
1561            .unwrap();
1562        assert_eq!(proto.dshape[1].0, DimName::NumProtos); // NCHW
1563        assert_eq!(proto.dshape[1].1, 32);
1564    }
1565
1566    #[test]
1567    fn infer_yolo26n_seg_end_to_end() {
1568        let r = infer_ultralytics_schema(&signals_from_fixture("yolo26n-seg")).unwrap();
1569        assert_eq!(r.schema.decoder_version, Some(DecoderVersion::Yolo26));
1570        assert_eq!(r.schema.outputs.len(), 2);
1571    }
1572
1573    #[test]
1574    fn infer_nhwc_proto_layout_classified() {
1575        // Hand-built signals: TFLite source, NHWC input, features-first det
1576        // tensor (4+80+32=116), NHWC proto ([1, 160, 160, 32]) — the other
1577        // valid proto channel ordering rule 1 must also classify correctly.
1578        let s = ModelSignals {
1579            source: ModelSource::TfLite,
1580            inputs: vec![TensorInfo {
1581                name: "images".into(),
1582                shape: vec![1, 640, 640, 3],
1583                dtype: DType::Float32,
1584                quantization: None,
1585            }],
1586            outputs: vec![
1587                TensorInfo {
1588                    name: "output0".into(),
1589                    shape: vec![1, 116, 8400],
1590                    dtype: DType::Float32,
1591                    quantization: None,
1592                },
1593                TensorInfo {
1594                    name: "output1".into(),
1595                    shape: vec![1, 160, 160, 32],
1596                    dtype: DType::Float32,
1597                    quantization: None,
1598                },
1599            ],
1600            metadata: synthetic_metadata(80, "segment", "False"),
1601        };
1602        let r = infer_ultralytics_schema(&s).unwrap();
1603        let proto = r
1604            .schema
1605            .outputs
1606            .iter()
1607            .find(|o| o.type_ == Some(LogicalType::Protos))
1608            .unwrap();
1609        assert_eq!(proto.dshape[3].0, DimName::NumProtos); // NHWC
1610        assert_eq!(proto.dshape[3].1, 32);
1611    }
1612
1613    /// The proto output is built field by field over `default_logical()`, so
1614    /// a dropped field reads as `None` rather than failing to compile. The
1615    /// decoder needs every one of these to dequantize and route the tensor.
1616    #[test]
1617    fn infer_proto_output_carries_its_name_decoder_dtype_and_quantization() {
1618        let quant = Quantization {
1619            scale: vec![0.0125],
1620            zero_point: Some(vec![-7]),
1621            axis: None,
1622            dtype: Some(DType::Int8),
1623        };
1624        let s = ModelSignals {
1625            source: ModelSource::TfLite,
1626            inputs: vec![TensorInfo {
1627                name: "images".into(),
1628                shape: vec![1, 640, 640, 3],
1629                dtype: DType::Int8,
1630                quantization: None,
1631            }],
1632            outputs: vec![
1633                TensorInfo {
1634                    name: "output0".into(),
1635                    shape: vec![1, 116, 8400],
1636                    dtype: DType::Int8,
1637                    quantization: Some(Quantization {
1638                        scale: vec![0.003_921_6],
1639                        zero_point: Some(vec![0]),
1640                        axis: None,
1641                        dtype: Some(DType::Int8),
1642                    }),
1643                },
1644                TensorInfo {
1645                    name: "protos".into(),
1646                    shape: vec![1, 160, 160, 32],
1647                    dtype: DType::Int8,
1648                    quantization: Some(quant.clone()),
1649                },
1650            ],
1651            metadata: synthetic_metadata(80, "segment", "False"),
1652        };
1653        let r = infer_ultralytics_schema(&s).unwrap();
1654        let proto = r
1655            .schema
1656            .outputs
1657            .iter()
1658            .find(|o| o.type_ == Some(LogicalType::Protos))
1659            .unwrap();
1660        assert_eq!(proto.name.as_deref(), Some("protos"));
1661        assert_eq!(proto.decoder, Some(DecoderKind::Ultralytics));
1662        assert_eq!(proto.dtype, Some(DType::Int8));
1663        assert_eq!(proto.quantization.as_ref(), Some(&quant));
1664    }
1665
1666    /// A proto tensor makes the model a segmentation model whether or not the
1667    /// metadata says `task: segment`. Requiring both would silently decode a
1668    /// segmentation export as plain detection.
1669    #[test]
1670    fn infer_treats_a_proto_tensor_without_a_task_as_segmentation() {
1671        let mut metadata = synthetic_metadata(80, "segment", "False");
1672        metadata.remove("task");
1673        let s = ModelSignals {
1674            source: ModelSource::TfLite,
1675            inputs: vec![TensorInfo {
1676                name: "images".into(),
1677                shape: vec![1, 640, 640, 3],
1678                dtype: DType::Float32,
1679                quantization: None,
1680            }],
1681            outputs: vec![
1682                TensorInfo {
1683                    name: "output0".into(),
1684                    shape: vec![1, 116, 8400],
1685                    dtype: DType::Float32,
1686                    quantization: None,
1687                },
1688                TensorInfo {
1689                    name: "output1".into(),
1690                    shape: vec![1, 160, 160, 32],
1691                    dtype: DType::Float32,
1692                    quantization: None,
1693                },
1694            ],
1695            metadata,
1696        };
1697        let r = infer_ultralytics_schema(&s).unwrap();
1698        assert!(
1699            r.schema
1700                .outputs
1701                .iter()
1702                .any(|o| o.type_ == Some(LogicalType::Protos)),
1703            "proto tensor present but no protos output was emitted"
1704        );
1705        // The task the metadata did not state is the one reported back, which
1706        // is the only place `is_segment` is observable.
1707        assert!(r.description.contains("segment"), "{}", r.description);
1708    }
1709
1710    fn proto_tensor(shape: Vec<usize>) -> TensorInfo {
1711        TensorInfo {
1712            name: "protos".into(),
1713            shape,
1714            dtype: DType::Float32,
1715            quantization: None,
1716        }
1717    }
1718
1719    /// Both halves of the rank-and-batch guard have to reject on their own.
1720    /// Requiring both instead would index past the end of a rank-3 shape.
1721    #[test]
1722    fn classify_proto_rejects_wrong_rank_or_batch() {
1723        for shape in [vec![1, 32, 160], vec![2, 32, 160, 160], vec![1, 160, 160]] {
1724            let err = classify_proto(&proto_tensor(shape.clone()), 640, 640).unwrap_err();
1725            assert!(
1726                matches!(err, InferError::UnsupportedLayout(_)),
1727                "shape {shape:?} gave {err:?}"
1728            );
1729        }
1730    }
1731
1732    /// The NHWC branch needs both spatial dims to line up. Accepting either
1733    /// one would read the proto count out of a dimension that is not it.
1734    #[test]
1735    fn classify_proto_rejects_a_shape_matching_only_one_spatial_dim() {
1736        // h/4 = w/4 = 160. s[1] matches, s[2] does not.
1737        let err = classify_proto(&proto_tensor(vec![1, 160, 7, 9]), 640, 640).unwrap_err();
1738        assert!(matches!(err, InferError::UnsupportedLayout(_)), "{err:?}");
1739
1740        // The genuine NHWC shape still classifies, with the proto count last.
1741        let (k, dshape) = classify_proto(&proto_tensor(vec![1, 160, 160, 32]), 640, 640).unwrap();
1742        assert_eq!(k, 32);
1743        assert_eq!(dshape[3], (DimName::NumProtos, 32));
1744    }
1745
1746    /// `unquote` strips a matching pair only. A single quote on one end is an
1747    /// unterminated value, not a quoted one.
1748    #[test]
1749    fn unquote_requires_a_matching_pair_of_quotes() {
1750        assert_eq!(unquote("'abc'").as_deref(), Some("abc"));
1751        assert_eq!(unquote("\"abc\"").as_deref(), Some("abc"));
1752        assert_eq!(unquote("'abc"), None);
1753        assert_eq!(unquote("abc'"), None);
1754        assert_eq!(unquote("\"abc"), None);
1755        assert_eq!(unquote("abc\""), None);
1756        assert_eq!(unquote("abc"), None);
1757    }
1758
1759    /// A comma inside a quoted value is part of the value. Without the quote
1760    /// arm the scanner never enters quoted state and splits straight through
1761    /// it.
1762    #[test]
1763    fn split_top_level_commas_keeps_quoted_commas_together() {
1764        let parts = split_top_level_commas("a='x,y', b=2");
1765        assert_eq!(parts, vec!["a='x,y'".to_string(), " b=2".to_string()]);
1766
1767        let parts = split_top_level_commas("a=\"x,y\",b=\"p,q\"");
1768        assert_eq!(
1769            parts,
1770            vec!["a=\"x,y\"".to_string(), "b=\"p,q\"".to_string()]
1771        );
1772    }
1773
1774    /// `InferError` is what a rejection reports to the caller. A `Display`
1775    /// that writes nothing turns every diagnostic into an empty string.
1776    #[test]
1777    fn infer_error_display_states_the_reason() {
1778        let err = InferError::NotUltralytics("no names key".into());
1779        let shown = err.to_string();
1780        assert!(shown.contains("no names key"), "{shown}");
1781        assert!(shown.contains("Ultralytics"), "{shown}");
1782
1783        let err = InferError::UnsupportedLayout("rank-3 proto".into());
1784        assert!(err.to_string().contains("rank-3 proto"));
1785    }
1786
1787    #[test]
1788    fn infer_segment_task_without_proto_rejected() {
1789        let s = ModelSignals {
1790            source: ModelSource::Onnx,
1791            inputs: vec![TensorInfo {
1792                name: "images".into(),
1793                shape: vec![1, 3, 640, 640],
1794                dtype: DType::Float32,
1795                quantization: None,
1796            }],
1797            outputs: vec![TensorInfo {
1798                name: "output0".into(),
1799                shape: vec![1, 84, 8400],
1800                dtype: DType::Float32,
1801                quantization: None,
1802            }],
1803            metadata: synthetic_metadata(80, "segment", "False"),
1804        };
1805        assert!(matches!(
1806            infer_ultralytics_schema(&s),
1807            Err(InferError::UnsupportedLayout(_))
1808        ));
1809    }
1810
1811    #[test]
1812    fn infer_detect_task_with_proto_rejected() {
1813        let s = ModelSignals {
1814            source: ModelSource::Onnx,
1815            inputs: vec![TensorInfo {
1816                name: "images".into(),
1817                shape: vec![1, 3, 640, 640],
1818                dtype: DType::Float32,
1819                quantization: None,
1820            }],
1821            outputs: vec![
1822                TensorInfo {
1823                    name: "output0".into(),
1824                    shape: vec![1, 84, 8400],
1825                    dtype: DType::Float32,
1826                    quantization: None,
1827                },
1828                TensorInfo {
1829                    name: "output1".into(),
1830                    shape: vec![1, 32, 160, 160],
1831                    dtype: DType::Float32,
1832                    quantization: None,
1833                },
1834            ],
1835            metadata: synthetic_metadata(80, "detect", "False"),
1836        };
1837        assert!(matches!(
1838            infer_ultralytics_schema(&s),
1839            Err(InferError::UnsupportedLayout(_))
1840        ));
1841    }
1842
1843    #[test]
1844    fn infer_int8_tflite_classifies_like_float32() {
1845        // Real Ultralytics int8 TFLite exports expose float32 I/O with NO
1846        // boundary quantization — the int8 graph is wrapped in
1847        // dequant/quant ops (Task 0 finding). The classifier must treat
1848        // this fixture exactly like its float32 counterpart, not attempt
1849        // to recover quantization that isn't there.
1850        let r = infer_ultralytics_schema(&signals_from_fixture("yolov8n_int8")).unwrap();
1851        let det = &r.schema.outputs[0];
1852        assert_eq!(det.type_, Some(LogicalType::Detection));
1853        assert_eq!(det.dtype, Some(DType::Float32));
1854        assert_eq!(det.quantization, None);
1855        assert_eq!(r.schema.decoder_version, Some(DecoderVersion::Yolov8));
1856        assert_eq!(det.normalized, Some(true));
1857    }
1858
1859    #[test]
1860    fn infer_quantized_boundary_carries_quantization() {
1861        // Hand-built signals for a model that DOES expose a quantized
1862        // boundary output (unlike any captured fixture) — the classifier
1863        // must pass `TensorInfo.quantization` through onto the schema's
1864        // detection output untouched.
1865        let quant = Quantization {
1866            scale: vec![0.02],
1867            zero_point: Some(vec![-5]),
1868            axis: None,
1869            dtype: Some(DType::Int8),
1870        };
1871        let s = ModelSignals {
1872            source: ModelSource::TfLite,
1873            inputs: vec![TensorInfo {
1874                name: "images".into(),
1875                shape: vec![1, 640, 640, 3],
1876                dtype: DType::Float32,
1877                quantization: None,
1878            }],
1879            outputs: vec![TensorInfo {
1880                name: "output0".into(),
1881                shape: vec![1, 84, 8400],
1882                dtype: DType::Int8,
1883                quantization: Some(quant.clone()),
1884            }],
1885            metadata: synthetic_metadata(80, "detect", "False"),
1886        };
1887        let r = infer_ultralytics_schema(&s).unwrap();
1888        let det = &r.schema.outputs[0];
1889        assert_eq!(det.dtype, Some(DType::Int8));
1890        assert_eq!(det.quantization, Some(quant));
1891    }
1892
1893    #[test]
1894    fn infer_rejects_batched_input() {
1895        // Every output layout here is batch-1, so a batched input would
1896        // emit an input dshape whose Batch disagrees with the detection
1897        // tensor beside it. Previously this succeeded and emitted
1898        // `(Batch, 4)`.
1899        let s = ModelSignals {
1900            source: ModelSource::Onnx,
1901            inputs: vec![TensorInfo {
1902                name: "images".into(),
1903                shape: vec![4, 3, 640, 640],
1904                dtype: DType::Float32,
1905                quantization: None,
1906            }],
1907            outputs: vec![TensorInfo {
1908                name: "output0".into(),
1909                shape: vec![1, 84, 8400],
1910                dtype: DType::Float32,
1911                quantization: None,
1912            }],
1913            metadata: synthetic_metadata(80, "detect", "False"),
1914        };
1915        match infer_ultralytics_schema(&s) {
1916            Err(InferError::UnsupportedLayout(msg)) => {
1917                assert!(msg.contains("batch 4"), "unexpected message: {msg}");
1918            }
1919            other => panic!("expected UnsupportedLayout, got {other:?}"),
1920        }
1921    }
1922
1923    #[test]
1924    fn infer_rejects_input_without_a_channel_dim() {
1925        let s = ModelSignals {
1926            source: ModelSource::Onnx,
1927            inputs: vec![TensorInfo {
1928                name: "images".into(),
1929                shape: vec![1, 4, 640, 640],
1930                dtype: DType::Float32,
1931                quantization: None,
1932            }],
1933            outputs: vec![TensorInfo {
1934                name: "output0".into(),
1935                shape: vec![1, 84, 8400],
1936                dtype: DType::Float32,
1937                quantization: None,
1938            }],
1939            metadata: synthetic_metadata(80, "detect", "False"),
1940        };
1941        match infer_ultralytics_schema(&s) {
1942            Err(InferError::UnsupportedLayout(msg)) => {
1943                assert!(msg.contains("channel dimension"), "unexpected: {msg}");
1944            }
1945            other => panic!("expected UnsupportedLayout, got {other:?}"),
1946        }
1947    }
1948
1949    #[test]
1950    fn infer_rejects_model_with_no_rank4_input() {
1951        let s = ModelSignals {
1952            source: ModelSource::Onnx,
1953            inputs: vec![TensorInfo {
1954                name: "images".into(),
1955                shape: vec![1, 3],
1956                dtype: DType::Float32,
1957                quantization: None,
1958            }],
1959            outputs: vec![TensorInfo {
1960                name: "output0".into(),
1961                shape: vec![1, 84, 8400],
1962                dtype: DType::Float32,
1963                quantization: None,
1964            }],
1965            metadata: synthetic_metadata(80, "detect", "False"),
1966        };
1967        match infer_ultralytics_schema(&s) {
1968            Err(InferError::UnsupportedLayout(msg)) => {
1969                assert!(msg.contains("no rank-4 model input"), "unexpected: {msg}");
1970            }
1971            other => panic!("expected UnsupportedLayout, got {other:?}"),
1972        }
1973    }
1974
1975    /// Builds detection signals for a 640x640 model with one rank-3 output.
1976    fn det_signals(nc: usize, out_shape: Vec<usize>, end2end: &str) -> ModelSignals {
1977        ModelSignals {
1978            source: ModelSource::Onnx,
1979            inputs: vec![TensorInfo {
1980                name: "images".into(),
1981                shape: vec![1, 3, 640, 640],
1982                dtype: DType::Float32,
1983                quantization: None,
1984            }],
1985            outputs: vec![TensorInfo {
1986                name: "output0".into(),
1987                shape: out_shape,
1988                dtype: DType::Float32,
1989                quantization: None,
1990            }],
1991            metadata: synthetic_metadata(nc, "detect", end2end),
1992        }
1993    }
1994
1995    #[test]
1996    fn infer_end_to_end_accepts_any_max_det_up_to_the_anchor_count() {
1997        // Ultralytics bakes max_det into the graph: the head's postprocess
1998        // is a TopK with k = min(max_det, anchors), so N is whatever the
1999        // export chose. The exporter clamps it to the anchor count
2000        // (exporter.py: `m.max_det = min(args.max_det, available)`), which
2001        // is the same quantity as `expected_anchors` here -- verified by
2002        // exporting yolo26n at max_det = 100/1000/2000/8400/10000, the last
2003        // two both yielding [1, 8400, 6].
2004        //
2005        // 300 is the default; 2000 is what a crowded-scene export uses and
2006        // was rejected outright before; 8400 == expected_anchors is what any
2007        // max_det >= anchors produces and is equally legitimate.
2008        for n in [1, 300, 2000, 8400] {
2009            let r = infer_ultralytics_schema(&det_signals(80, vec![1, n, 6], "True"))
2010                .unwrap_or_else(|e| panic!("max_det={n} must infer: {e}"));
2011            assert_eq!(r.schema.decoder_version, Some(DecoderVersion::Yolo26));
2012            assert_eq!(r.schema.outputs[0].type_, Some(LogicalType::Detections));
2013            assert_eq!(
2014                r.schema.outputs[0].dshape,
2015                vec![
2016                    (DimName::Batch, 1),
2017                    (DimName::NumBoxes, n),
2018                    (DimName::NumFeatures, 6),
2019                ],
2020                "max_det={n}"
2021            );
2022            // End-to-end ran NMS in-graph, so the schema adds none.
2023            assert_eq!(r.schema.nms, None, "max_det={n}");
2024        }
2025    }
2026
2027    #[test]
2028    fn infer_end_to_end_rejects_more_rows_than_anchors() {
2029        // Upstream cannot emit this: TopK's k is clamped to the anchor
2030        // count, so a longer tensor is not an Ultralytics end-to-end output
2031        // and the bound stays a real check rather than a formality.
2032        match infer_ultralytics_schema(&det_signals(80, vec![1, 8401, 6], "True")) {
2033            Err(InferError::UnsupportedLayout(msg)) => {
2034                assert!(msg.contains("N<=8400"), "unexpected message: {msg}");
2035            }
2036            other => panic!("expected UnsupportedLayout, got {other:?}"),
2037        }
2038    }
2039
2040    #[test]
2041    fn infer_legacy_nms_model_is_unaffected_by_max_det() {
2042        // The pre-NMS (NMS-required) head has no max_det in its graph at
2043        // all: it emits every anchor, and the detection count is decided
2044        // later by the decoder's own NMS. So the schema must classify it by
2045        // the anchor count regardless, and must NOT acquire a NumBoxes
2046        // dimension that looks like an end-to-end row count.
2047        let r = infer_ultralytics_schema(&det_signals(80, vec![1, 84, 8400], "False")).unwrap();
2048        assert_eq!(r.schema.decoder_version, Some(DecoderVersion::Yolov8));
2049        assert_eq!(r.schema.outputs[0].type_, Some(LogicalType::Detection));
2050        assert_eq!(
2051            r.schema.outputs[0].dshape,
2052            vec![
2053                (DimName::Batch, 1),
2054                (DimName::NumFeatures, 84),
2055                (DimName::NumBoxes, 8400),
2056            ]
2057        );
2058        // Class-aware NMS, matching Ultralytics' own agnostic=False. The
2059        // schema names the mode only; thresholds and max_det stay the
2060        // caller's, set on the builder per frame-rate/recall budget.
2061        assert_eq!(r.schema.nms, Some(NmsMode::ClassAware));
2062
2063        // Anchors-first is classified identically -- neither orientation
2064        // may be mistaken for an end-to-end [1, N, feat] tensor.
2065        let t = infer_ultralytics_schema(&det_signals(80, vec![1, 8400, 84], "False")).unwrap();
2066        assert_eq!(t.schema.decoder_version, Some(DecoderVersion::Yolov8));
2067        assert_eq!(
2068            t.schema.outputs[0].dshape,
2069            vec![
2070                (DimName::Batch, 1),
2071                (DimName::NumBoxes, 8400),
2072                (DimName::NumFeatures, 84),
2073            ]
2074        );
2075    }
2076
2077    #[test]
2078    fn infer_pre_nms_wins_when_a_shape_reads_as_both() {
2079        // With metadata stripped and nc == 2, feat == 4 + nc == 6 == the
2080        // end-to-end feature width, so [1, A, 6] is genuinely both layouts
2081        // at once. pre-NMS is tried first and wins; this is the one place a
2082        // guess remains, so it is pinned rather than left to drift.
2083        let mut m = synthetic_metadata(2, "detect", "False");
2084        m.remove("end2end");
2085        let s = ModelSignals {
2086            metadata: m,
2087            ..det_signals(2, vec![1, 8400, 6], "False")
2088        };
2089        let r = infer_ultralytics_schema(&s).unwrap();
2090        assert_eq!(r.schema.decoder_version, Some(DecoderVersion::Yolov8));
2091    }
2092
2093    #[test]
2094    fn meta_empty_names_rejected() {
2095        // nc == 0 makes feat == 4, so a [1,4,A] output would infer a schema
2096        // the decoder then refuses ("Yolo num_features 4 must be > 4").
2097        let mut m = BTreeMap::new();
2098        m.insert("names".into(), "{}".into());
2099        match UltralyticsMeta::from_metadata(&m) {
2100            Err(InferError::BadNames(msg)) => assert!(msg.contains("empty"), "{msg}"),
2101            other => panic!("expected BadNames, got {other:?}"),
2102        }
2103    }
2104
2105    #[test]
2106    fn meta_malformed_end2end_is_an_error_not_an_absent_key() {
2107        // `None` means "key absent", which selects the shape-only fallback.
2108        // A corrupt value must not collapse into that -- the model is
2109        // actively claiming something, just not something parseable.
2110        let mut m = synthetic_metadata(80, "detect", "False");
2111        m.insert("end2end".into(), "yes".into());
2112        match UltralyticsMeta::from_metadata(&m) {
2113            Err(InferError::BadMetadata(msg)) => assert!(msg.contains("end2end"), "{msg}"),
2114            other => panic!("expected BadMetadata, got {other:?}"),
2115        }
2116
2117        // Absent is still fine, and still means "fall back to shapes".
2118        m.remove("end2end");
2119        assert_eq!(UltralyticsMeta::from_metadata(&m).unwrap().end2end, None);
2120    }
2121
2122    #[test]
2123    fn meta_foreign_vendor_rejected_but_missing_provenance_allowed() {
2124        // A `names` map alone is not an Ultralytics signature; another
2125        // exporter's model would otherwise be given Ultralytics box, score
2126        // and normalization conventions.
2127        let mut m = synthetic_metadata(80, "detect", "False");
2128        m.insert("author".into(), "SomeOtherVendor".into());
2129        assert!(matches!(
2130            UltralyticsMeta::from_metadata(&m),
2131            Err(InferError::NotUltralytics(_))
2132        ));
2133
2134        // Provenance that names Ultralytics passes, in either field.
2135        m.insert("author".into(), "Ultralytics".into());
2136        m.insert("docs".into(), "https://docs.ultralytics.com".into());
2137        assert!(UltralyticsMeta::from_metadata(&m).is_ok());
2138
2139        // Absence is not a contradiction: quantizers and format converters
2140        // routinely keep `names` and drop provenance.
2141        m.remove("author");
2142        m.remove("docs");
2143        assert!(UltralyticsMeta::from_metadata(&m).is_ok());
2144    }
2145
2146    #[test]
2147    fn infer_zero_row_end_to_end_rejected() {
2148        // [1, 0, 6] satisfies every other end-to-end condition but yields a
2149        // schema the decoder refuses for a zero dimension.
2150        assert!(matches!(
2151            infer_ultralytics_schema(&det_signals(80, vec![1, 0, 6], "True")),
2152            Err(InferError::UnsupportedLayout(_))
2153        ));
2154    }
2155
2156    #[test]
2157    fn infer_duplicate_matching_outputs_are_ambiguous() {
2158        let input = TensorInfo {
2159            name: "images".into(),
2160            shape: vec![1, 3, 640, 640],
2161            dtype: DType::Float32,
2162            quantization: None,
2163        };
2164        let t = |name: &str, shape: Vec<usize>| TensorInfo {
2165            name: name.into(),
2166            shape,
2167            dtype: DType::Float32,
2168            quantization: None,
2169        };
2170
2171        // Two identical pre-NMS detection tensors: nothing distinguishes
2172        // them but the order the runtime listed them in.
2173        let pre = ModelSignals {
2174            source: ModelSource::Onnx,
2175            inputs: vec![input.clone()],
2176            outputs: vec![t("a", vec![1, 84, 8400]), t("b", vec![1, 84, 8400])],
2177            metadata: synthetic_metadata(80, "detect", "False"),
2178        };
2179        match infer_ultralytics_schema(&pre) {
2180            Err(InferError::AmbiguousLayout(msg)) => {
2181                assert!(msg.contains("`a`") && msg.contains("`b`"), "{msg}");
2182            }
2183            other => panic!("expected AmbiguousLayout, got {other:?}"),
2184        }
2185
2186        // Same for two end-to-end candidates.
2187        let e2e = ModelSignals {
2188            source: ModelSource::Onnx,
2189            inputs: vec![input],
2190            outputs: vec![t("a", vec![1, 300, 6]), t("b", vec![1, 200, 6])],
2191            metadata: synthetic_metadata(80, "detect", "True"),
2192        };
2193        assert!(matches!(
2194            infer_ultralytics_schema(&e2e),
2195            Err(InferError::AmbiguousLayout(_))
2196        ));
2197    }
2198
2199    #[test]
2200    fn infer_malformed_per_tensor_quantization_rejected() {
2201        let build = |q: Quantization| ModelSignals {
2202            source: ModelSource::TfLite,
2203            inputs: vec![TensorInfo {
2204                name: "images".into(),
2205                shape: vec![1, 640, 640, 3],
2206                dtype: DType::Float32,
2207                quantization: None,
2208            }],
2209            outputs: vec![TensorInfo {
2210                name: "output0".into(),
2211                shape: vec![1, 84, 8400],
2212                dtype: DType::Int8,
2213                quantization: Some(q),
2214            }],
2215            metadata: synthetic_metadata(80, "detect", "False"),
2216        };
2217
2218        // No scale at all: previously emitted `{"scale": []}`, which the
2219        // decoder reads as a zero scale instead of refusing.
2220        match infer_ultralytics_schema(&build(Quantization {
2221            scale: vec![],
2222            zero_point: Some(vec![]),
2223            axis: None,
2224            dtype: Some(DType::Int8),
2225        })) {
2226            Err(InferError::UnsupportedQuantization(msg)) => {
2227                assert!(msg.contains("no scale"), "{msg}");
2228            }
2229            other => panic!("expected UnsupportedQuantization, got {other:?}"),
2230        }
2231
2232        // Zero points that do not pair with the scale were silently
2233        // truncated to the first element.
2234        match infer_ultralytics_schema(&build(Quantization {
2235            scale: vec![0.02],
2236            zero_point: Some(vec![-5, -7]),
2237            axis: None,
2238            dtype: Some(DType::Int8),
2239        })) {
2240            Err(InferError::UnsupportedQuantization(msg)) => {
2241                assert!(msg.contains("zero point"), "{msg}");
2242            }
2243            other => panic!("expected UnsupportedQuantization, got {other:?}"),
2244        }
2245
2246        // Symmetric quantization (no zero points) stays valid.
2247        assert!(infer_ultralytics_schema(&build(Quantization {
2248            scale: vec![0.02],
2249            zero_point: None,
2250            axis: None,
2251            dtype: Some(DType::Int8),
2252        }))
2253        .is_ok());
2254    }
2255
2256    #[test]
2257    fn infer_pre_nms_is_independent_of_output_order() {
2258        // An export may publish more than one anchor-shaped rank-3 output:
2259        // a `[1, 4, A]` box tensor beside the `[1, 4+nc, A]` detection
2260        // tensor, or a `[1, k, A]` mask-coefficient tensor. Scanning must
2261        // not stop at the first near-miss, or the result depends on the
2262        // order the runtime happened to report its outputs in.
2263        let det = TensorInfo {
2264            name: "output0".into(),
2265            shape: vec![1, 84, 8400],
2266            dtype: DType::Float32,
2267            quantization: None,
2268        };
2269        let aux = TensorInfo {
2270            name: "boxes".into(),
2271            shape: vec![1, 4, 8400],
2272            dtype: DType::Float32,
2273            quantization: None,
2274        };
2275        let build = |outputs: Vec<TensorInfo>| ModelSignals {
2276            source: ModelSource::Onnx,
2277            inputs: vec![TensorInfo {
2278                name: "images".into(),
2279                shape: vec![1, 3, 640, 640],
2280                dtype: DType::Float32,
2281                quantization: None,
2282            }],
2283            outputs,
2284            metadata: synthetic_metadata(80, "detect", "False"),
2285        };
2286
2287        let det_first = infer_ultralytics_schema(&build(vec![det.clone(), aux.clone()]))
2288            .expect("detection tensor first must classify");
2289        let aux_first = infer_ultralytics_schema(&build(vec![aux, det]))
2290            .expect("an auxiliary anchor-shaped tensor first must not shadow the detection tensor");
2291        assert_eq!(det_first.schema.outputs[0].name, Some("output0".into()));
2292        assert_eq!(aux_first.schema.outputs[0].name, Some("output0".into()));
2293        assert_eq!(det_first.schema, aux_first.schema);
2294    }
2295
2296    #[test]
2297    fn infer_class_count_mismatch_survives_a_full_scan() {
2298        // When nothing matches, the anchor-shaped near-miss is still the
2299        // more useful diagnostic than "no output matches the layout".
2300        let s = ModelSignals {
2301            source: ModelSource::Onnx,
2302            inputs: vec![TensorInfo {
2303                name: "images".into(),
2304                shape: vec![1, 3, 640, 640],
2305                dtype: DType::Float32,
2306                quantization: None,
2307            }],
2308            outputs: vec![TensorInfo {
2309                name: "output0".into(),
2310                shape: vec![1, 84, 8400],
2311                dtype: DType::Float32,
2312                quantization: None,
2313            }],
2314            // 10 classes would need feat == 14, not the 84 present.
2315            metadata: synthetic_metadata(10, "detect", "False"),
2316        };
2317        assert!(matches!(
2318            infer_ultralytics_schema(&s),
2319            Err(InferError::ClassCountMismatch {
2320                expected: 10,
2321                found: 80
2322            })
2323        ));
2324    }
2325
2326    /// What every captured fixture must infer to, asserted field by field.
2327    ///
2328    /// `yolo11n-seg` and `yolov8n-seg_int8` previously appeared only in
2329    /// `tests/infer_builder.rs`'s "did the builder accept it" loop, which
2330    /// passes for any schema that parses -- so the only v11 segmentation
2331    /// fixture and the only int8 *segmentation* fixture had no field-level
2332    /// coverage at all. A table keeps every fixture honest at once, and
2333    /// `every_fixture_is_asserted` below makes an unasserted one impossible
2334    /// to add.
2335    struct FixtureExpectation {
2336        name: &'static str,
2337        version: DecoderVersion,
2338        det_type: LogicalType,
2339        /// Detection tensor `dshape`, in physical order.
2340        det_dshape: &'static [(DimName, usize)],
2341        normalized: bool,
2342        nms: Option<NmsMode>,
2343        /// Proto channel count and spatial dims, for segmentation fixtures.
2344        proto: Option<(usize, usize, usize)>,
2345    }
2346
2347    const A: usize = 8400; // 640x640 anchors: 80² + 40² + 20²
2348    const NC: usize = 80;
2349
2350    fn fixture_expectations() -> Vec<FixtureExpectation> {
2351        use DimName::{Batch, NumBoxes, NumFeatures};
2352        // Every captured export is features-first ([1, 4+nc, A]); no real
2353        // exporter transposes to anchors-first, despite an early draft of
2354        // these tests assuming TFLite did.
2355        const DET: &[(DimName, usize)] = &[(Batch, 1), (NumFeatures, 4 + NC), (NumBoxes, A)];
2356        const SEG: &[(DimName, usize)] = &[(Batch, 1), (NumFeatures, 4 + NC + 32), (NumBoxes, A)];
2357        // End-to-end heads emit [1, max_det, 6] (+32 mask coefficients).
2358        const E2E: &[(DimName, usize)] = &[(Batch, 1), (NumBoxes, 300), (NumFeatures, 6)];
2359        const E2E_SEG: &[(DimName, usize)] = &[(Batch, 1), (NumBoxes, 300), (NumFeatures, 6 + 32)];
2360
2361        let pre_nms = |name, det_dshape, normalized, proto| FixtureExpectation {
2362            name,
2363            version: DecoderVersion::Yolov8,
2364            det_type: LogicalType::Detection,
2365            det_dshape,
2366            normalized,
2367            nms: Some(NmsMode::ClassAware),
2368            proto,
2369        };
2370        let e2e = |name, det_dshape, normalized, proto| FixtureExpectation {
2371            name,
2372            version: DecoderVersion::Yolo26,
2373            det_type: LogicalType::Detections,
2374            det_dshape,
2375            normalized,
2376            // End-to-end heads ran NMS in-graph; the schema adds none.
2377            nms: None,
2378            proto,
2379        };
2380        // Ultralytics ONNX exports are pixel-space, TFLite exports [0,1];
2381        // proto is (k, H/4, W/4).
2382        const P: Option<(usize, usize, usize)> = Some((32, 160, 160));
2383        vec![
2384            pre_nms("yolov8n", DET, false, None),
2385            pre_nms("yolo11n", DET, false, None),
2386            pre_nms("yolov8n-seg", SEG, false, P),
2387            pre_nms("yolo11n-seg", SEG, false, P),
2388            pre_nms("yolov8n_float32", DET, true, None),
2389            pre_nms("yolov8n_int8", DET, true, None),
2390            pre_nms("yolov8n-seg_float32", SEG, true, P),
2391            pre_nms("yolov8n-seg_int8", SEG, true, P),
2392            e2e("yolo26n", E2E, false, None),
2393            e2e("yolo26n-seg", E2E_SEG, false, P),
2394            e2e("yolo26n_float32", E2E, true, None),
2395            // Same network, different exporter default. The `yolo26n` row above
2396            // was captured from ultralytics 8.4.137, whose `nms` arg defaulted to
2397            // `False`; `exporter.py`'s `model.end2end = self.args.nms is False`
2398            // then selected the one2one (NMS-free, [1, 300, 6]) head. In 8.4.153
2399            // that default is `None`, so the same expression is false and the
2400            // one2many anchor-grid head ships instead -- hence a Yolov8-shaped
2401            // pre-NMS row for a YOLO26 checkpoint. Both heads live in the weights.
2402            pre_nms("yolo26n_coreml", DET, false, None),
2403        ]
2404    }
2405
2406    #[test]
2407    fn every_captured_fixture_infers_the_expected_schema() {
2408        for want in fixture_expectations() {
2409            let name = want.name;
2410            let r = infer_ultralytics_schema(&signals_from_fixture(name))
2411                .unwrap_or_else(|e| panic!("{name}: inference failed: {e}"));
2412
2413            assert_eq!(r.labels.len(), NC, "{name}: class count");
2414            assert_eq!(
2415                r.labels[0], "person",
2416                "{name}: labels must be index-ordered"
2417            );
2418            assert_eq!(r.labels[NC - 1], "toothbrush", "{name}: labels index order");
2419            assert_eq!(
2420                r.schema.decoder_version,
2421                Some(want.version),
2422                "{name}: version"
2423            );
2424            assert_eq!(r.schema.nms, want.nms, "{name}: nms mode");
2425
2426            let det = &r.schema.outputs[0];
2427            assert_eq!(det.type_, Some(want.det_type), "{name}: detection type");
2428            assert_eq!(det.dshape, want.det_dshape, "{name}: detection dshape");
2429            assert_eq!(det.normalized, Some(want.normalized), "{name}: normalized");
2430            assert_eq!(
2431                det.decoder,
2432                Some(DecoderKind::Ultralytics),
2433                "{name}: decoder"
2434            );
2435            assert_eq!(det.encoding, Some(BoxEncoding::Direct), "{name}: encoding");
2436            assert_eq!(
2437                det.score_format,
2438                Some(ScoreFormat::PerClass),
2439                "{name}: scores"
2440            );
2441
2442            match want.proto {
2443                None => {
2444                    assert_eq!(r.schema.outputs.len(), 1, "{name}: detection-only");
2445                }
2446                Some((k, ph, pw)) => {
2447                    assert_eq!(r.schema.outputs.len(), 2, "{name}: detection + proto");
2448                    let proto = &r.schema.outputs[1];
2449                    assert_eq!(proto.type_, Some(LogicalType::Protos), "{name}: proto type");
2450                    assert_eq!(
2451                        proto.dshape,
2452                        vec![
2453                            (DimName::Batch, 1),
2454                            (DimName::NumProtos, k),
2455                            (DimName::Height, ph),
2456                            (DimName::Width, pw),
2457                        ],
2458                        "{name}: proto dshape"
2459                    );
2460                }
2461            }
2462
2463            let task = if want.proto.is_some() {
2464                "segment"
2465            } else {
2466                "detect"
2467            };
2468            assert!(
2469                r.description.contains(task) && r.description.contains("Ultralytics"),
2470                "{name}: description was `{}`",
2471                r.description
2472            );
2473        }
2474    }
2475
2476    #[test]
2477    fn every_fixture_is_asserted() {
2478        // Without this, a fixture can be added (or renamed) and silently
2479        // never asserted on -- which is exactly how `yolo11n-seg` and
2480        // `yolov8n-seg_int8` ended up with no field-level coverage.
2481        let dir = infer_fixture_dir();
2482        let mut on_disk: Vec<String> = std::fs::read_dir(&dir)
2483            .unwrap_or_else(|e| panic!("{} exists: {e}", dir.display()))
2484            .filter_map(|e| {
2485                let name = e.ok()?.file_name().into_string().ok()?;
2486                name.strip_suffix(".signals.json").map(str::to_string)
2487            })
2488            .collect();
2489        on_disk.sort();
2490
2491        let mut asserted: Vec<String> = fixture_expectations()
2492            .iter()
2493            .map(|f| f.name.to_string())
2494            .collect();
2495        asserted.sort();
2496
2497        assert_eq!(
2498            on_disk, asserted,
2499            "every *.signals.json fixture must appear in fixture_expectations()"
2500        );
2501    }
2502
2503    #[test]
2504    fn infer_other_source_refuses_rather_than_guessing_normalization() {
2505        // `normalized` follows the exporter and cannot be read off the
2506        // shapes. ONNX (pixel-space) and TFLite ([0,1]) are measured;
2507        // CoreML (pixel-space) is both source-traced and measured. An
2508        // uncharacterized container gets a
2509        // typed refusal, because guessing scales every box by the input
2510        // size and the resulting schema looks perfectly valid.
2511        let s = ModelSignals {
2512            source: ModelSource::Other,
2513            inputs: vec![TensorInfo {
2514                name: "images".into(),
2515                shape: vec![1, 3, 640, 640],
2516                dtype: DType::Float32,
2517                quantization: None,
2518            }],
2519            outputs: vec![TensorInfo {
2520                name: "output0".into(),
2521                shape: vec![1, 84, 8400],
2522                dtype: DType::Float32,
2523                quantization: None,
2524            }],
2525            metadata: synthetic_metadata(80, "detect", "False"),
2526        };
2527        assert!(matches!(
2528            infer_ultralytics_schema(&s),
2529            Err(InferError::UnknownBoxConvention)
2530        ));
2531
2532        // The same signals under a characterized source succeed, so the
2533        // refusal is about the container and nothing else.
2534        let ok = ModelSignals {
2535            source: ModelSource::Onnx,
2536            ..s
2537        };
2538        assert_eq!(
2539            infer_ultralytics_schema(&ok).unwrap().schema.outputs[0].normalized,
2540            Some(false)
2541        );
2542    }
2543
2544    #[test]
2545    fn infer_per_channel_quantization_rejected() {
2546        // Per-channel quantization on a boundary output is refused here
2547        // rather than passed through. The decoder supports per-tensor only
2548        // (`the v1 decoder only supports per-tensor quantization`), and it
2549        // refuses per-channel with OR without an `axis` -- so emitting the
2550        // schema anyway would defer the failure to `DecoderBuilder::build`,
2551        // where the message names the decoder instead of the tensor that
2552        // caused it. No captured fixture exhibits this; a future export
2553        // convention might.
2554        let s = ModelSignals {
2555            source: ModelSource::TfLite,
2556            inputs: vec![TensorInfo {
2557                name: "images".into(),
2558                shape: vec![1, 640, 640, 3],
2559                dtype: DType::Float32,
2560                quantization: None,
2561            }],
2562            outputs: vec![TensorInfo {
2563                name: "output0".into(),
2564                shape: vec![1, 84, 8400],
2565                dtype: DType::Int8,
2566                quantization: Some(Quantization {
2567                    scale: vec![0.02, 0.03],
2568                    zero_point: Some(vec![-5, -7]),
2569                    axis: Some(1),
2570                    dtype: Some(DType::Int8),
2571                }),
2572            }],
2573            metadata: synthetic_metadata(80, "detect", "False"),
2574        };
2575        match infer_ultralytics_schema(&s) {
2576            Err(InferError::UnsupportedQuantization(msg)) => {
2577                assert!(msg.contains("per-channel"), "unexpected message: {msg}");
2578                assert!(
2579                    msg.contains("output0"),
2580                    "message must name the tensor: {msg}"
2581                );
2582            }
2583            other => panic!("expected UnsupportedQuantization, got {other:?}"),
2584        }
2585    }
2586
2587    #[test]
2588    fn infer_per_channel_proto_quantization_rejected() {
2589        // The same guard runs on the proto tensor, not just the detection
2590        // output -- `check_quantized_boundary` is called for both.
2591        let s = ModelSignals {
2592            source: ModelSource::TfLite,
2593            inputs: vec![TensorInfo {
2594                name: "images".into(),
2595                shape: vec![1, 640, 640, 3],
2596                dtype: DType::Float32,
2597                quantization: None,
2598            }],
2599            outputs: vec![
2600                TensorInfo {
2601                    name: "output0".into(),
2602                    shape: vec![1, 116, 8400],
2603                    dtype: DType::Float32,
2604                    quantization: None,
2605                },
2606                TensorInfo {
2607                    name: "output1".into(),
2608                    shape: vec![1, 32, 160, 160],
2609                    dtype: DType::Int8,
2610                    quantization: Some(Quantization {
2611                        scale: vec![0.02, 0.03],
2612                        zero_point: Some(vec![-5, -7]),
2613                        axis: Some(1),
2614                        dtype: Some(DType::Int8),
2615                    }),
2616                },
2617            ],
2618            metadata: synthetic_metadata(80, "segment", "False"),
2619        };
2620        match infer_ultralytics_schema(&s) {
2621            Err(InferError::UnsupportedQuantization(msg)) => {
2622                assert!(
2623                    msg.contains("output1"),
2624                    "message must name the proto: {msg}"
2625                );
2626            }
2627            other => panic!("expected UnsupportedQuantization, got {other:?}"),
2628        }
2629    }
2630
2631    #[test]
2632    fn infer_integer_output_without_quant_rejected() {
2633        // Same synthetic signals as above but with `quantization: None` —
2634        // an integer-dtype boundary output with no quantization params is
2635        // never silently accepted (there'd be no way to dequantize it).
2636        let s = ModelSignals {
2637            source: ModelSource::TfLite,
2638            inputs: vec![TensorInfo {
2639                name: "images".into(),
2640                shape: vec![1, 640, 640, 3],
2641                dtype: DType::Float32,
2642                quantization: None,
2643            }],
2644            outputs: vec![TensorInfo {
2645                name: "output0".into(),
2646                shape: vec![1, 84, 8400],
2647                dtype: DType::Int8,
2648                quantization: None,
2649            }],
2650            metadata: synthetic_metadata(80, "detect", "False"),
2651        };
2652        match infer_ultralytics_schema(&s) {
2653            Err(InferError::UnsupportedLayout(msg)) => {
2654                assert!(
2655                    msg.contains("quantization"),
2656                    "unexpected UnsupportedLayout message: {msg}"
2657                );
2658            }
2659            other => panic!("expected UnsupportedLayout, got {other:?}"),
2660        }
2661    }
2662
2663    #[test]
2664    fn infer_pre_nms_rejects_non_unit_batch() {
2665        // A [2, 84, 8400] output has a matching anchors/feature pair but a
2666        // non-unit leading (batch) dim. It must never be silently accepted
2667        // and mislabeled batch-1 (find_e2e_candidate already requires
2668        // shape[0] == 1; find_anchors_dim must match it).
2669        let mut s = signals_from_fixture("yolov8n");
2670        s.outputs[0].shape = vec![2, 84, 8400];
2671        assert!(matches!(
2672            infer_ultralytics_schema(&s),
2673            Err(InferError::UnsupportedLayout(_))
2674        ));
2675    }
2676
2677    #[test]
2678    fn coreml_source_resolves_a_box_convention() {
2679        // CoreML must not hit the `Other` refusal: a native `.mlpackage`
2680        // export is a characterized convention, not an unknown one.
2681        let s = ModelSignals {
2682            source: ModelSource::CoreMl,
2683            inputs: vec![TensorInfo {
2684                name: "images".into(),
2685                shape: vec![1, 3, 640, 640],
2686                dtype: DType::Float16,
2687                quantization: None,
2688            }],
2689            outputs: vec![TensorInfo {
2690                name: "output0".into(),
2691                shape: vec![1, 84, 8400],
2692                dtype: DType::Float16,
2693                quantization: None,
2694            }],
2695            metadata: synthetic_metadata(80, "detect", "False"),
2696        };
2697        let r = infer_ultralytics_schema(&s)
2698            .expect("CoreML signals must infer, not raise UnknownBoxConvention");
2699        let o = &r.schema.outputs[0];
2700        assert_eq!(
2701            o.normalized,
2702            Some(false),
2703            "Ultralytics CoreML exports trace the raw model, so boxes are \
2704             pixel-space exactly as in the ONNX export"
2705        );
2706    }
2707
2708    #[test]
2709    fn coreml_nms_pipeline_export_is_refused_not_classified() {
2710        // `yolo export ... format=coreml nms=True` does not produce the
2711        // anchor-grid tensor `ModelSource::CoreMl`'s pixel-space
2712        // resolution reasons about -- it substitutes Apple's own NMS
2713        // pipeline, whose two outputs are dynamically shaped and already
2714        // post-NMS. Real measured shapes (Task 3b brief, from a genuine
2715        // `yolo26n.pt` export): `confidence` and `coordinates` are both
2716        // float32 multiArrays with an empty (fully dynamic) shape --
2717        // nothing here is the `[1, 4+nc, A]` layout this module looks
2718        // for, so inference must refuse it rather than guess. This
2719        // refusal is intentional: do not "fix" it into acceptance, since
2720        // that would apply CoreMl's pixel-space convention (or none at
2721        // all) to output Apple's pipeline has already decoded and scaled
2722        // its own way.
2723        let s = ModelSignals {
2724            source: ModelSource::CoreMl,
2725            inputs: vec![TensorInfo {
2726                name: "image".into(),
2727                shape: vec![1, 3, 640, 640],
2728                dtype: DType::Float16,
2729                quantization: None,
2730            }],
2731            outputs: vec![
2732                TensorInfo {
2733                    name: "confidence".into(),
2734                    shape: vec![],
2735                    dtype: DType::Float32,
2736                    quantization: None,
2737                },
2738                TensorInfo {
2739                    name: "coordinates".into(),
2740                    shape: vec![],
2741                    dtype: DType::Float32,
2742                    quantization: None,
2743                },
2744            ],
2745            metadata: synthetic_metadata(80, "detect", "False"),
2746        };
2747        assert!(matches!(
2748            infer_ultralytics_schema(&s),
2749            Err(InferError::UnsupportedLayout(_))
2750        ));
2751    }
2752}