Skip to main content

edgefirst_decoder/model_config/
resolve.rs

1// SPDX-FileCopyrightText: Copyright 2026 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4//! Resolution of every source into one [`ResolvedModelConfig`].
5
6use std::collections::BTreeMap;
7use std::path::Path;
8
9use edgefirst_tensor::{
10    broadcast_rgb, FitMode, Interpolation, LetterboxPlacement, Normalization, NormalizationPreset,
11    PadSpace, PreprocessSpec,
12};
13use serde_json::Value;
14
15use super::container::{parse_labels, trim_labels, usable_labels};
16use super::error::invalid;
17use super::heuristics::{hf_family_name, hf_semantic_by_names, semantic_schema, SemanticExport};
18use super::hf::{decoder_for_architecture, parse_hf_config, parse_hf_preprocessor, HfCheck};
19use super::{
20    ConfigField, ConfigSource, ModelConfigError, ModelFamily, ModelFiles, ModelOverrides,
21    ModelTask, NormalizeOverride, Provenance, ResizeOverride,
22};
23use crate::configs::DimName;
24use crate::infer::{
25    infer_ultralytics_schema, infer_ultralytics_schema_ignoring_names, InferError, InferredSchema,
26    ModelSignals,
27};
28use crate::schema::{
29    self, DecodeMode, DecoderKind, LogicalOutput, LogicalType, NormalizationConfig, SchemaV2,
30};
31use crate::DecoderBuilder;
32
33/// Everything discovery reads: the files in the model container, the files
34/// beside it, and the model's tensors as the runtime reports them.
35///
36/// ```
37/// use edgefirst_decoder::model_config::{ModelFiles, ModelSources};
38///
39/// let sources = ModelSources::new().with_sidecar(ModelFiles::default());
40/// assert!(sources.signals.is_none());
41/// ```
42#[derive(Debug, Clone, Default)]
43#[non_exhaustive]
44pub struct ModelSources {
45    /// Files embedded in the model.
46    pub embedded: ModelFiles,
47    /// Files in the model's directory.
48    pub sidecar: ModelFiles,
49    /// Input and output tensors and runtime metadata. Needed to build a
50    /// schema from a Hugging Face config, by Ultralytics inference and by
51    /// the name heuristics.
52    pub signals: Option<ModelSignals>,
53}
54
55impl ModelSources {
56    /// No sources.
57    #[must_use]
58    pub fn new() -> Self {
59        Self::default()
60    }
61
62    /// With the files embedded in the model.
63    #[must_use]
64    pub fn with_embedded(mut self, files: ModelFiles) -> Self {
65        self.embedded = files;
66        self
67    }
68
69    /// With the files in the model's directory.
70    #[must_use]
71    pub fn with_sidecar(mut self, files: ModelFiles) -> Self {
72        self.sidecar = files;
73        self
74    }
75
76    /// With the model's tensors and runtime metadata.
77    #[must_use]
78    pub fn with_signals(mut self, signals: ModelSignals) -> Self {
79        self.signals = Some(signals);
80        self
81    }
82}
83
84/// A model's configuration as resolved from every source.
85///
86/// [`preprocessing`](Self::preprocessing) and
87/// [`decode_mode`](Self::decode_mode) are authoritative; `schema` carries a
88/// copy of them for consumers that read the schema alone.
89#[derive(Debug, Clone, PartialEq)]
90#[non_exhaustive]
91pub struct ResolvedModelConfig {
92    /// The output schema, with `input.preprocessing` filled with the
93    /// resolved values when the schema has an input.
94    pub schema: SchemaV2,
95    /// Class names in channel order, with outer whitespace trimmed; empty
96    /// when no source names them.
97    pub labels: Vec<String>,
98    /// The family the configuration identifies.
99    pub family: ModelFamily,
100    /// What the outputs decode to.
101    pub task: ModelTask,
102    /// The semantic decode mode, for semantic models.
103    pub decode_mode: Option<DecodeMode>,
104    /// The input preprocessing of the reference pipeline. Its
105    /// `interpolation` may be one the image crate runs as bilinear
106    /// ([`Interpolation::supported`]); a note says so.
107    pub preprocessing: PreprocessSpec,
108    /// Where each value came from.
109    pub provenance: BTreeMap<ConfigField, Provenance>,
110    /// A one-line summary of the model.
111    pub description: String,
112    /// Deviations and skipped files worth logging.
113    pub notes: Vec<String>,
114    /// Hugging Face `do_reduce_labels`: the reference dataset drops its
115    /// background class and shifts the rest down. Exposed, not applied.
116    pub reduce_labels: Option<bool>,
117}
118
119impl ResolvedModelConfig {
120    /// Where `field` came from.
121    #[must_use]
122    pub fn source(&self, field: ConfigField) -> Option<&Provenance> {
123        self.provenance.get(&field)
124    }
125
126    /// One `key=value [source]` line per resolved value, for logging. Lines
127    /// whose key is an override (`normalize`, `resize`, `task`, `decode`)
128    /// print the value in the override grammar, e.g.
129    /// `resize=letterbox-topleft [preprocessor_config.json]` or
130    /// `normalize=imagenet [guess: SegFormer default]`.
131    #[must_use]
132    pub fn describe(&self) -> Vec<String> {
133        let p = &self.preprocessing;
134        let list = |v: &[f32; 3]| v.map(|x| x.to_string()).join(",");
135        let value = |field: ConfigField| -> Option<String> {
136            Some(match field {
137                ConfigField::Normalization => normalize_override(&p.normalization).to_string(),
138                ConfigField::Resize => resize_override(p).to_string(),
139                ConfigField::Placement => p.placement.as_str().to_owned(),
140                ConfigField::Pad => format!("{} {}", list(&p.pad), p.pad_space.as_str()),
141                ConfigField::Interpolation => p.interpolation.as_str().to_owned(),
142                ConfigField::Task => self.task.as_str().to_owned(),
143                ConfigField::DecodeMode => self.decode_mode?.as_str().to_owned(),
144                ConfigField::Labels => match self.labels.len() {
145                    0 => "unknown".to_owned(),
146                    n => format!("{n} classes"),
147                },
148                ConfigField::Size => {
149                    let (w, h) = p.size?;
150                    format!("{w}x{h}")
151                }
152            })
153        };
154        ConfigField::all()
155            .iter()
156            .filter_map(|&f| {
157                let v = value(f)?;
158                let prov = self
159                    .provenance
160                    .get(&f)
161                    .map(|p| format!(" [{p}]"))
162                    .unwrap_or_default();
163                Some(format!("{}={v}{prov}", f.as_str()))
164            })
165            .collect()
166    }
167}
168
169/// `n` as a `normalize` override: its preset, else mean/std on 0–1.
170fn normalize_override(n: &Normalization) -> NormalizeOverride {
171    match n.as_preset() {
172        Some(p) => NormalizeOverride::Preset(p),
173        None => {
174            let unit = |v: [f32; 3]| v.map(|x| (f64::from(x) / 255.0) as f32);
175            NormalizeOverride::MeanStd {
176                mean: unit(n.mean()),
177                std: unit(n.std()),
178            }
179        }
180    }
181}
182
183/// `p`'s fit as a `resize` override, with its placement when letterboxed.
184fn resize_override(p: &PreprocessSpec) -> ResizeOverride {
185    match p.fit {
186        FitMode::Letterbox => ResizeOverride::Letterbox(Some(p.placement)),
187        _ => ResizeOverride::Stretch,
188    }
189}
190
191/// A [`LogicalOutput`] with every field unset.
192pub(crate) fn empty_logical() -> LogicalOutput {
193    LogicalOutput {
194        name: None,
195        type_: None,
196        shape: Vec::new(),
197        dshape: Vec::new(),
198        decoder: None,
199        encoding: None,
200        score_format: None,
201        normalized: None,
202        anchors: None,
203        stride: None,
204        dtype: None,
205        quantization: None,
206        outputs: Vec::new(),
207        activation_applied: None,
208        activation_required: None,
209        semantic: Default::default(),
210    }
211}
212
213/// The Ultralytics preprocessing, which is also the documented default:
214/// unsigned, centred letterbox, pixel pad 114, bilinear.
215fn ultralytics_convention() -> PreprocessSpec {
216    PreprocessSpec::new()
217}
218
219/// The Hugging Face image processor defaults: ImageNet, stretch, bilinear,
220/// and for a letterbox a top-left placement padded with normalized zeros.
221fn hf_convention() -> PreprocessSpec {
222    let mut s = PreprocessSpec::new();
223    s.normalization = Normalization::imagenet();
224    s.fit = FitMode::Stretch;
225    s.placement = LetterboxPlacement::TopLeft;
226    s.pad = [0.0; 3];
227    s.pad_space = PadSpace::Normalized;
228    s
229}
230
231type Slot<T> = Option<(T, Provenance)>;
232
233fn fill<T>(slot: &mut Slot<T>, value: T, prov: &Provenance) {
234    if slot.is_none() {
235        *slot = Some((value, prov.clone()));
236    }
237}
238
239fn take<T>(
240    provenance: &mut BTreeMap<ConfigField, Provenance>,
241    field: ConfigField,
242    slot: Slot<T>,
243    default: T,
244) -> T {
245    match slot {
246        Some((v, p)) => {
247            provenance.insert(field, p);
248            v
249        }
250        None => default,
251    }
252}
253
254/// Names that appear more than once ignoring case, as `"bed" at 7, 9`, or
255/// `None` when every name is distinct.
256fn repeated_labels(labels: &[String]) -> Option<String> {
257    let mut seen: BTreeMap<String, Vec<usize>> = BTreeMap::new();
258    for (i, name) in labels.iter().enumerate() {
259        seen.entry(name.to_lowercase()).or_default().push(i);
260    }
261    let repeats: Vec<String> = seen
262        .values()
263        .filter(|at| at.len() > 1)
264        .map(|at| {
265            let index: Vec<String> = at.iter().map(ToString::to_string).collect();
266            format!("{:?} at {}", labels[at[0]], index.join(", "))
267        })
268        .collect();
269    (!repeats.is_empty()).then(|| repeats.join("; "))
270}
271
272fn at(source: ConfigSource, detail: impl Into<String>) -> Provenance {
273    Provenance::new(source, Some(&detail.into()))
274}
275
276#[derive(Default)]
277struct State {
278    schema: Option<(SchemaV2, Provenance)>,
279    family: Option<ModelFamily>,
280    /// Class-name candidates, highest priority first.
281    labels: Vec<(Vec<String>, Provenance)>,
282    normalization: Slot<Normalization>,
283    fit: Slot<FitMode>,
284    placement: Slot<LetterboxPlacement>,
285    pad: Slot<([f32; 3], PadSpace)>,
286    interpolation: Slot<Interpolation>,
287    size: Option<(usize, usize)>,
288    num_classes: Option<usize>,
289    hf_stride: Option<usize>,
290    declared_placement: bool,
291    reduce_labels: Option<bool>,
292    description: Option<String>,
293    infer_error: Option<InferError>,
294    notes: Vec<String>,
295}
296
297impl State {
298    fn set_schema(&mut self, schema: SchemaV2, prov: Provenance, family: ModelFamily) {
299        if self.schema.is_none() {
300            self.schema = Some((schema, prov));
301            self.family = Some(family);
302        }
303    }
304
305    /// Takes a schema laid out from the model's tensors.
306    fn adopt_export(
307        &mut self,
308        export: SemanticExport,
309        prov: Provenance,
310        family: ModelFamily,
311        description: String,
312    ) {
313        if self.schema.is_none() {
314            self.description = Some(description);
315            self.size = Some(export.size);
316            self.num_classes = Some(export.num_classes);
317            self.hf_stride = Some(export.stride);
318            self.set_schema(export.schema, prov, family);
319        }
320    }
321
322    /// Fills the preprocessing slots that `fields` names from `spec`.
323    fn fill_spec(&mut self, spec: &PreprocessSpec, prov: &Provenance, fields: &[ConfigField]) {
324        for field in fields {
325            match field {
326                ConfigField::Normalization => {
327                    fill(&mut self.normalization, spec.normalization, prov)
328                }
329                ConfigField::Resize => fill(&mut self.fit, spec.fit, prov),
330                ConfigField::Placement => fill(&mut self.placement, spec.placement, prov),
331                ConfigField::Pad => fill(&mut self.pad, (spec.pad, spec.pad_space), prov),
332                ConfigField::Interpolation => {
333                    fill(&mut self.interpolation, spec.interpolation, prov)
334                }
335                _ => {}
336            }
337        }
338    }
339
340    fn apply_schema_preprocessing(
341        &mut self,
342        pre: &schema::Preprocessing,
343        file: &str,
344        prov: &Provenance,
345    ) -> Result<(), ModelConfigError> {
346        if let Some(n) = &pre.normalization {
347            let (mean, std) = n.mean_std(3)?;
348            let n = Normalization::mean_std(&mean, &std)
349                .map_err(|e| invalid(file, "input.preprocessing.normalization", e.to_string()))?;
350            fill(&mut self.normalization, n, prov);
351        }
352        let fit = pre
353            .resize
354            .or_else(|| pre.letterbox.as_ref().map(|_| FitMode::Letterbox));
355        if let Some(fit) = fit {
356            fill(&mut self.fit, fit, prov);
357        }
358        if let Some(lb) = &pre.letterbox {
359            if let Some(p) = lb.placement {
360                if self.placement.is_none() {
361                    self.declared_placement = true;
362                }
363                fill(&mut self.placement, p, prov);
364            }
365            if let Some(pad) = &lb.pad {
366                let pad = broadcast_rgb(pad)
367                    .ok_or_else(|| invalid(file, "letterbox.pad", "needs 1 or 3 values"))?;
368                fill(&mut self.pad, (pad, lb.pad_space.unwrap_or_default()), prov);
369            }
370        }
371        if let Some(i) = pre.interpolation {
372            fill(&mut self.interpolation, i, prov);
373        }
374        Ok(())
375    }
376
377    /// Steps 1 and 2: `edgefirst.json` and `labels.txt`.
378    fn edgefirst(
379        &mut self,
380        files: &ModelFiles,
381        source: ConfigSource,
382    ) -> Result<(), ModelConfigError> {
383        if let Some(f) = &files.edgefirst_json {
384            let here = place(source, &f.name);
385            let prov = at(source, here.as_str());
386            let value: Value = serde_json::from_str(&f.text)
387                .map_err(|e| invalid(&f.name, "(document)", e.to_string()))?;
388            let validation = value.get("validation");
389            let legacy_normalization = validation
390                .and_then(|v| v.get("normalization"))
391                .and_then(Value::as_str)
392                .map(str::to_owned);
393            let legacy_resize = validation
394                .and_then(|v| v.get("preprocessing"))
395                .and_then(Value::as_str)
396                .map(str::to_owned);
397            let classes: Option<Vec<String>> = value
398                .pointer("/dataset/classes")
399                .and_then(Value::as_array)
400                .and_then(|a| a.iter().map(|c| c.as_str().map(str::to_owned)).collect());
401            let classes = classes.and_then(|c| {
402                usable_labels(trim_labels(&f.name, "dataset.classes", c), &mut self.notes)
403            });
404            let name = value
405                .get("name")
406                .and_then(Value::as_str)
407                .unwrap_or("")
408                .to_owned();
409            let schema = SchemaV2::from_json_value(value)?;
410            if let Some(pre) = schema.input.as_ref().and_then(|i| i.preprocessing.as_ref()) {
411                self.apply_schema_preprocessing(pre, &f.name, &prov)?;
412            }
413            if let Some(n) = legacy_normalization {
414                let preset = NormalizationPreset::from_str_code(&n).ok_or_else(|| {
415                    invalid(
416                        &f.name,
417                        "validation.normalization",
418                        format!("`{n}` is not a preset"),
419                    )
420                })?;
421                let p = at(source, format!("{here} validation.normalization"));
422                fill(&mut self.normalization, Normalization::preset(preset), &p);
423            }
424            if let Some(r) = legacy_resize {
425                let p = at(source, format!("{here} validation.preprocessing"));
426                match r.as_str() {
427                    "letterbox" => fill(&mut self.fit, FitMode::Letterbox, &p),
428                    "resize" | "stretch" => fill(&mut self.fit, FitMode::Stretch, &p),
429                    other => self.notes.push(format!(
430                        "{}: validation.preprocessing `{other}` ignored",
431                        f.name
432                    )),
433                }
434            }
435            if let Some(classes) = classes.filter(|c| !c.is_empty()) {
436                self.labels.push((classes, prov.clone()));
437            }
438            if schema.outputs.is_empty() {
439                self.notes.push(format!(
440                    "{here} declares no outputs; the schema comes from a later source"
441                ));
442            } else if self.schema.is_none() {
443                self.description = Some(format!("EdgeFirst model {name}").trim_end().to_owned());
444                self.set_schema(schema, prov, ModelFamily::EdgeFirst);
445            }
446        }
447        if let Some(f) = &files.labels_txt {
448            let labels = parse_labels(&f.text);
449            if !labels.is_empty() {
450                self.labels
451                    .push((labels, at(source, place(source, &f.name))));
452            }
453        }
454        Ok(())
455    }
456
457    /// Steps 3 and 4: Hugging Face `config.json` and `preprocessor_config.json`.
458    fn hugging_face(
459        &mut self,
460        files: &ModelFiles,
461        source: ConfigSource,
462        signals: Option<&ModelSignals>,
463    ) -> Result<(), ModelConfigError> {
464        let mut model = None;
465        if let Some(f) = &files.hf_config {
466            match parse_hf_config(&f.text, &f.name)? {
467                HfCheck::Valid(m) => model = Some((m, at(source, place(source, &f.name)))),
468                HfCheck::NotHf(why) => self.skipped(source, &f.name, "Hugging Face config", why),
469            }
470        }
471        let mut pre = None;
472        if let Some(f) = &files.hf_preprocessor {
473            match parse_hf_preprocessor(&f.text, &f.name, model.is_some())? {
474                HfCheck::Valid(p) => pre = Some((p, at(source, place(source, &f.name)))),
475                HfCheck::NotHf(why) => {
476                    self.skipped(source, &f.name, "Hugging Face image processor config", why)
477                }
478            }
479        }
480        if let Some((mut m, prov)) = model {
481            self.notes.append(&mut m.notes);
482            let file = files
483                .hf_config
484                .as_ref()
485                .map_or("config.json", |f| f.name.as_str());
486            if self.schema.is_none() {
487                let Some(arch) = m
488                    .architectures
489                    .iter()
490                    .find(|a| decoder_for_architecture(a).is_some())
491                else {
492                    return Err(ModelConfigError::UnsupportedArchitecture(m.model_type));
493                };
494                let signals = signals.ok_or_else(|| ModelConfigError::Unresolved {
495                    missing: "the model's input and output tensors, to lay out the outputs a \
496                              Hugging Face config describes"
497                        .into(),
498                    overrides: "the model signals".into(),
499                })?;
500                let export = semantic_schema(signals)?;
501                let classes = export.num_classes;
502                if let Some(n) = m.num_labels.filter(|&n| m.labels.is_none() && n != classes) {
503                    let note = format!("{file}: num_labels is {n} but the output has {classes}");
504                    log::warn!("{note}");
505                    self.notes.push(note);
506                }
507                let description = format!(
508                    "Hugging Face {} semantic segmentation ({arch}), {classes} classes",
509                    m.model_type
510                );
511                self.adopt_export(export, prov.clone(), ModelFamily::HuggingFace, description);
512            }
513            if let Some(labels) = m.labels {
514                self.labels.push((labels, prov));
515            }
516        }
517        if let Some((p, prov)) = pre {
518            if let Some(n) = p.normalization {
519                fill(&mut self.normalization, n, &prov);
520            }
521            if let Some(i) = p.interpolation {
522                fill(&mut self.interpolation, i, &prov);
523            }
524            if p.pad {
525                self.fill_spec(
526                    &hf_convention(),
527                    &prov,
528                    &[ConfigField::Placement, ConfigField::Pad],
529                );
530                fill(&mut self.fit, FitMode::Letterbox, &prov);
531            } else if p.size.is_some() {
532                fill(&mut self.fit, FitMode::Stretch, &prov);
533            }
534            if let (Some(want), Some(have)) = (p.size, self.size) {
535                if want != have {
536                    self.notes.push(format!(
537                        "{prov}: the reference resizes to {}x{}; the model input is {}x{}",
538                        want.0, want.1, have.0, have.1
539                    ));
540                }
541            }
542            self.reduce_labels = self.reduce_labels.or(p.reduce_labels);
543            self.notes.extend(p.notes);
544        }
545        Ok(())
546    }
547
548    fn skipped(&mut self, source: ConfigSource, name: &str, what: &str, why: &str) {
549        let note = format!("{} present but not a {what} ({why})", place(source, name));
550        log::debug!("{note}");
551        self.notes.push(note);
552    }
553
554    /// Step 5: Ultralytics metadata.
555    fn ultralytics(&mut self, signals: &ModelSignals, classes: Option<usize>) {
556        match self.infer_ultralytics(signals, classes) {
557            Ok(inferred) => {
558                let prov = Provenance::new(ConfigSource::UltralyticsMetadata, None);
559                self.fill_ultralytics_labels(inferred.labels, &prov);
560                self.description = Some(inferred.description);
561                self.set_schema(inferred.schema, prov, ModelFamily::Ultralytics);
562                let conv = Provenance::new(ConfigSource::Default, Some("Ultralytics convention"));
563                self.fill_spec(&ultralytics_convention(), &conv, PREPROCESSING);
564            }
565            Err(InferError::NotUltralytics(why)) => log::debug!("not Ultralytics: {why}"),
566            Err(e) => self.infer_error = Some(e),
567        }
568    }
569
570    /// Ultralytics class names for a schema another source supplied.
571    fn ultralytics_labels(&mut self, signals: &ModelSignals, classes: Option<usize>) {
572        if let Ok(inferred) = self.infer_ultralytics(signals, classes) {
573            let prov = Provenance::new(ConfigSource::UltralyticsMetadata, None);
574            self.fill_ultralytics_labels(inferred.labels, &prov);
575        }
576    }
577
578    /// Infers the Ultralytics schema; when `names` disagrees with the output,
579    /// lays it out from `classes` (the labels override's count) or the
580    /// output's class count and ignores `names`.
581    fn infer_ultralytics(
582        &mut self,
583        signals: &ModelSignals,
584        classes: Option<usize>,
585    ) -> Result<InferredSchema, InferError> {
586        match infer_ultralytics_schema(signals) {
587            Err(InferError::ClassCountMismatch { expected, .. }) => {
588                let (inferred, found) = infer_ultralytics_schema_ignoring_names(signals, classes)?;
589                let note = format!(
590                    "Ultralytics metadata: names lists {expected} classes but the output has \
591                     {found}; its class names are ignored"
592                );
593                log::warn!("{note}");
594                self.notes.push(note);
595                Ok(inferred)
596            }
597            r => r,
598        }
599    }
600
601    fn fill_ultralytics_labels(&mut self, labels: Vec<String>, prov: &Provenance) {
602        if labels.is_empty() {
603            return;
604        }
605        let trimmed = trim_labels("Ultralytics metadata", "names", labels);
606        if let Some(labels) = usable_labels(trimmed, &mut self.notes) {
607            self.labels.push((labels, prov.clone()));
608        }
609    }
610
611    /// Step 6: a Hugging Face export recognised by its tensor names.
612    fn heuristics(&mut self, signals: &ModelSignals) {
613        if let Some(export) = hf_semantic_by_names(signals) {
614            let prov = Provenance::new(ConfigSource::Heuristic, Some("pixel_values/logits"));
615            let description = format!(
616                "Hugging Face semantic segmentation export (recognised by tensor names), {} \
617                 classes",
618                export.num_classes
619            );
620            self.adopt_export(export, prov, ModelFamily::HuggingFace, description);
621        }
622    }
623
624    /// The defaults of a Hugging Face family's image processor.
625    fn hf_family_defaults(&mut self) {
626        let detail = hf_family_name(self.hf_stride.unwrap_or(0));
627        let guess = Provenance::new(ConfigSource::Heuristic, Some(detail));
628        self.fill_spec(
629            &hf_convention(),
630            &guess,
631            &[
632                ConfigField::Normalization,
633                ConfigField::Resize,
634                ConfigField::Interpolation,
635            ],
636        );
637    }
638
639    /// Step 7: the documented defaults.
640    fn defaults(&mut self) {
641        let d = Provenance::new(ConfigSource::Default, None);
642        self.fill_spec(&PreprocessSpec::new(), &d, PREPROCESSING);
643    }
644}
645
646/// Every preprocessing slot.
647const PREPROCESSING: &[ConfigField] = &[
648    ConfigField::Normalization,
649    ConfigField::Resize,
650    ConfigField::Placement,
651    ConfigField::Pad,
652    ConfigField::Interpolation,
653];
654
655/// Where a file was found, as provenance detail: its name beside the model,
656/// `embedded <name>` inside it.
657fn place(source: ConfigSource, name: &str) -> String {
658    match source {
659        ConfigSource::EmbeddedEdgefirst | ConfigSource::EmbeddedHf => format!("embedded {name}"),
660        _ => name.to_owned(),
661    }
662}
663
664/// The `(width, height)` of a schema's input, when its layout resolves.
665fn input_size(schema: &SchemaV2) -> Option<(usize, usize)> {
666    let input = schema.input.as_ref()?;
667    if !input.dshape.is_empty() {
668        let dim = |n| {
669            input
670                .dshape
671                .iter()
672                .position(|(d, _)| *d == n)
673                .and_then(|i| input.shape.get(i).copied())
674        };
675        return Some((dim(DimName::Width)?, dim(DimName::Height)?));
676    }
677    match (input.shape.as_slice(), input.channels()?) {
678        ([_, c, h, w], ch) if *c == ch => Some((*w, *h)),
679        ([_, h, w, _], _) => Some((*w, *h)),
680        _ => None,
681    }
682}
683
684fn task_of(schema: &SchemaV2) -> ModelTask {
685    let has = |t: LogicalType| schema.outputs.iter().any(|o| o.type_ == Some(t));
686    let semantic = has(LogicalType::Segmentation) || has(LogicalType::Masks);
687    let boxes =
688        has(LogicalType::Boxes) || has(LogicalType::Detection) || has(LogicalType::Detections);
689    if semantic && !boxes {
690        ModelTask::Semantic
691    } else if has(LogicalType::Protos) || has(LogicalType::MaskCoefs) {
692        ModelTask::Segment
693    } else {
694        ModelTask::Detect
695    }
696}
697
698fn semantic_outputs(schema: &mut SchemaV2) -> impl Iterator<Item = &mut LogicalOutput> {
699    schema.outputs.iter_mut().filter(|o| {
700        matches!(
701            o.type_,
702            Some(LogicalType::Segmentation | LogicalType::Masks)
703        )
704    })
705}
706
707fn semantic_classes(schema: &SchemaV2) -> Option<usize> {
708    schema
709        .outputs
710        .iter()
711        .find(|o| o.type_ == Some(LogicalType::Segmentation))
712        .and_then(|o| {
713            o.dshape
714                .iter()
715                .find(|(d, _)| *d == DimName::NumClasses)
716                .map(|(_, n)| *n)
717        })
718}
719
720/// The class count the decoder builder derives from `schema`.
721fn decoder_classes(schema: &SchemaV2) -> Option<usize> {
722    DecoderBuilder::new()
723        .with_schema(schema.clone())
724        .build()
725        .ok()?
726        .class_count()
727}
728
729/// The letterbox `(placement, pad)` a bare `letterbox` override means for
730/// a family.
731fn family_letterbox(
732    family: ModelFamily,
733    declared: bool,
734    current: (LetterboxPlacement, ([f32; 3], PadSpace)),
735) -> (LetterboxPlacement, ([f32; 3], PadSpace), &'static str) {
736    let of = |s: PreprocessSpec| (s.placement, (s.pad, s.pad_space));
737    let (placement, pad, why) = match family {
738        ModelFamily::Ultralytics => {
739            let (p, d) = of(ultralytics_convention());
740            (p, d, "Ultralytics convention")
741        }
742        ModelFamily::HuggingFace => {
743            let (p, d) = of(hf_convention());
744            (p, d, "Hugging Face convention")
745        }
746        ModelFamily::EdgeFirst if declared => (current.0, current.1, "as edgefirst.json declares"),
747        ModelFamily::EdgeFirst => (
748            LetterboxPlacement::Center,
749            current.1,
750            "edgefirst.json declares no placement",
751        ),
752        _ => {
753            let (p, d) = of(ultralytics_convention());
754            (p, d, "convention unknown")
755        }
756    };
757    (placement, pad, why)
758}
759
760/// Resolves a configuration from already-read sources.
761///
762/// When the signals carry no metadata, the embedded files' metadata stands
763/// in for it, so Ultralytics inference sees the ONNX props or the TFLite
764/// `metadata.json` the container holds.
765///
766/// # Errors
767///
768/// Any [`ModelConfigError`]: a malformed recognised file, an unsupported
769/// field or architecture, an override that cannot apply, or nothing to
770/// build a schema from ([`ModelConfigError::Unresolved`]).
771pub fn resolve_sources(
772    sources: ModelSources,
773    overrides: &ModelOverrides,
774) -> Result<ResolvedModelConfig, ModelConfigError> {
775    let ModelSources {
776        embedded,
777        sidecar,
778        signals,
779    } = sources;
780    let signals = signals.map(|mut s| {
781        if s.metadata.is_empty() {
782            s.metadata = embedded.metadata.clone();
783        }
784        s
785    });
786    let mut st = State::default();
787    st.edgefirst(&embedded, ConfigSource::EmbeddedEdgefirst)?;
788    st.edgefirst(&sidecar, ConfigSource::SidecarEdgefirst)?;
789    st.hugging_face(&embedded, ConfigSource::EmbeddedHf, signals.as_ref())?;
790    st.hugging_face(&sidecar, ConfigSource::SidecarHf, signals.as_ref())?;
791    let override_classes = overrides.labels.as_ref().map(Vec::len);
792    if let Some(s) = &signals {
793        if st.schema.is_none() {
794            st.ultralytics(s, override_classes);
795        } else {
796            st.ultralytics_labels(s, override_classes);
797        }
798        if st.schema.is_none() {
799            st.heuristics(s);
800        }
801    }
802    let ov = Provenance::new(ConfigSource::Override, None);
803    let mut task_overridden = false;
804    if st.schema.is_none() && overrides.task == Some(ModelTask::Semantic) {
805        if let Some(s) = &signals {
806            let export = semantic_schema(s)?;
807            let description = format!(
808                "semantic segmentation (task override), {} classes",
809                export.num_classes
810            );
811            st.adopt_export(export, ov.clone(), ModelFamily::Unknown, description);
812            task_overridden = true;
813        }
814    }
815    let Some((mut schema, schema_prov)) = st.schema.take() else {
816        if let Some(e) = st.infer_error {
817            return Err(ModelConfigError::Infer(e));
818        }
819        return Err(ModelConfigError::Unresolved {
820            missing: "a model configuration: no edgefirst.json, Hugging Face config or \
821                      Ultralytics metadata was found, and the input and output tensors do not \
822                      match a Hugging Face segmentation export"
823                .into(),
824            overrides: "an edgefirst.json beside the model, or task=semantic with the model \
825                        signals"
826                .into(),
827        });
828    };
829    let family = st.family.unwrap_or(ModelFamily::Unknown);
830    if family == ModelFamily::HuggingFace {
831        st.hf_family_defaults();
832    }
833    st.defaults();
834    let size = st.size.or_else(|| input_size(&schema));
835
836    let d = PreprocessSpec::new();
837    let mut provenance = BTreeMap::new();
838    let p = &mut provenance;
839    let mut normalization = take(
840        p,
841        ConfigField::Normalization,
842        st.normalization,
843        d.normalization,
844    );
845    let mut fit = take(p, ConfigField::Resize, st.fit, d.fit);
846    let mut placement = take(p, ConfigField::Placement, st.placement, d.placement);
847    let mut pad = take(p, ConfigField::Pad, st.pad, (d.pad, d.pad_space));
848    let interpolation = take(
849        p,
850        ConfigField::Interpolation,
851        st.interpolation,
852        d.interpolation,
853    );
854    let mut labels = Vec::new();
855    let mut notes = st.notes;
856    if size.is_some() {
857        let detail = Some("model input");
858        provenance.insert(
859            ConfigField::Size,
860            Provenance::new(schema_prov.source, detail),
861        );
862    }
863
864    let mut task = task_of(&schema);
865    if let Some(want) = overrides.task {
866        if want != task && !(task == ModelTask::Segment && want == ModelTask::Detect) {
867            return Err(ModelConfigError::Override(format!(
868                "task={want} does not apply to a {task} model"
869            )));
870        }
871        task = want;
872    }
873    let task_prov = if task_overridden || overrides.task.is_some() {
874        ov.clone()
875    } else {
876        schema_prov.clone()
877    };
878    provenance.insert(ConfigField::Task, task_prov);
879    if let Some(n) = overrides.normalize {
880        normalization = n.to_normalization()?;
881        provenance.insert(ConfigField::Normalization, ov.clone());
882    }
883    if let Some(r) = overrides.resize {
884        match r {
885            ResizeOverride::Letterbox(want) => {
886                let (p, d, why) = family_letterbox(family, st.declared_placement, (placement, pad));
887                fit = FitMode::Letterbox;
888                pad = d;
889                placement = want.unwrap_or(p);
890                notes.push(format!(
891                    "resize={r} resolved to {} with {} pad {:?} ({why})",
892                    ResizeOverride::Letterbox(Some(placement)),
893                    pad.1.as_str(),
894                    pad.0
895                ));
896                provenance.insert(ConfigField::Placement, ov.clone());
897                provenance.insert(ConfigField::Pad, ov.clone());
898            }
899            _ => fit = FitMode::Stretch,
900        }
901        provenance.insert(ConfigField::Resize, ov.clone());
902    }
903    let mut decode_mode = None;
904    if task == ModelTask::Semantic {
905        let declared = semantic_outputs(&mut schema).find_map(|o| o.semantic.decode_mode);
906        let modelpack =
907            semantic_outputs(&mut schema).any(|o| o.decoder == Some(DecoderKind::ModelPack));
908        decode_mode = Some(declared.unwrap_or(if modelpack {
909            DecodeMode::Fast
910        } else {
911            DecodeMode::Accurate
912        }));
913        let p = if declared.is_some() {
914            schema_prov.clone()
915        } else {
916            Provenance::new(ConfigSource::Default, None)
917        };
918        provenance.insert(ConfigField::DecodeMode, p);
919    }
920    if let Some(mode) = overrides.decode {
921        if task != ModelTask::Semantic {
922            return Err(ModelConfigError::Override(format!(
923                "decode applies to semantic models, not a {task} model"
924            )));
925        }
926        decode_mode = Some(mode);
927        for o in semantic_outputs(&mut schema) {
928            o.semantic.decode_mode = Some(mode);
929        }
930        provenance.insert(ConfigField::DecodeMode, ov.clone());
931    }
932    let classes = st
933        .num_classes
934        .or_else(|| decoder_classes(&schema))
935        .or_else(|| (task == ModelTask::Semantic).then(|| semantic_classes(&schema))?);
936    if let Some(l) = &overrides.labels {
937        if let Some(i) = l.iter().position(|n| n.trim().is_empty()) {
938            return Err(ModelConfigError::Override(format!(
939                "labels has an empty name at {i}"
940            )));
941        }
942        if let Some(expected) = classes.filter(|&n| n != l.len()) {
943            return Err(ModelConfigError::LabelCountMismatch {
944                expected,
945                found: l.len(),
946            });
947        }
948        labels = l.iter().map(|n| n.trim().to_owned()).collect();
949        provenance.insert(ConfigField::Labels, ov);
950    } else {
951        for (candidate, prov) in st.labels {
952            if let Some(n) = classes.filter(|&n| n != candidate.len()) {
953                let note = format!(
954                    "labels [{prov}] name {} classes but the model has {n}; they are ignored",
955                    candidate.len()
956                );
957                log::warn!("{note}");
958                notes.push(note);
959            } else {
960                labels = candidate;
961                provenance.insert(ConfigField::Labels, prov);
962                break;
963            }
964        }
965    }
966    if labels.is_empty() {
967        let note = "labels unknown; supply them with the labels override";
968        log::warn!("{note}");
969        notes.push(note.into());
970    }
971    if let Some(repeats) = repeated_labels(&labels) {
972        let note = format!(
973            "labels [{}] repeat names, so matching classes by name is ambiguous: {repeats}",
974            provenance
975                .get(&ConfigField::Labels)
976                .map(ToString::to_string)
977                .unwrap_or_default(),
978        );
979        log::warn!("{note}");
980        notes.push(note);
981    }
982
983    let mut preprocessing = PreprocessSpec::new();
984    preprocessing.normalization = normalization;
985    preprocessing.fit = fit;
986    preprocessing.placement = placement;
987    preprocessing.pad = pad.0;
988    preprocessing.pad_space = pad.1;
989    preprocessing.interpolation = interpolation;
990    preprocessing.size = size;
991    if interpolation.supported() != interpolation {
992        notes.push(format!(
993            "interpolation={} is not implemented by the image crate, which runs {}",
994            interpolation.as_str(),
995            interpolation.supported().as_str()
996        ));
997    }
998    if let Some(input) = schema.input.as_mut() {
999        input.preprocessing = Some(to_schema_preprocessing(&preprocessing));
1000    }
1001    let description = st
1002        .description
1003        .unwrap_or_else(|| format!("{} {} model", family.as_str(), task.as_str()));
1004    Ok(ResolvedModelConfig {
1005        schema,
1006        labels,
1007        family,
1008        task,
1009        decode_mode,
1010        preprocessing,
1011        provenance,
1012        description,
1013        notes,
1014        reduce_labels: st.reduce_labels,
1015    })
1016}
1017
1018fn to_schema_preprocessing(p: &PreprocessSpec) -> schema::Preprocessing {
1019    let normalization = match p.normalization.as_preset() {
1020        Some(preset) => NormalizationConfig::Preset(preset),
1021        None => NormalizationConfig::MeanStd {
1022            mean: p.normalization.mean().to_vec(),
1023            std: p.normalization.std().to_vec(),
1024        },
1025    };
1026    let letterbox = schema::Letterbox {
1027        placement: Some(p.placement),
1028        pad: Some(p.pad.to_vec()),
1029        pad_space: Some(p.pad_space),
1030    };
1031    schema::Preprocessing {
1032        normalization: Some(normalization),
1033        resize: Some(p.fit),
1034        interpolation: Some(p.interpolation),
1035        letterbox: (p.fit == FitMode::Letterbox).then_some(letterbox),
1036    }
1037}
1038
1039/// Resolves the configuration of the model file at `path`: what it embeds,
1040/// the files in its directory, and `signals` when given.
1041///
1042/// # Errors
1043///
1044/// [`ModelConfigError::Io`] when the model or a sidecar cannot be read, and
1045/// every error of [`resolve_sources`].
1046///
1047/// ```no_run
1048/// use edgefirst_decoder::model_config::{resolve_model_config, ModelOverrides};
1049///
1050/// let config = resolve_model_config("segformer.onnx", None, &ModelOverrides::default())?;
1051/// for line in config.describe() {
1052///     println!("{line}");
1053/// }
1054/// # Ok::<(), edgefirst_decoder::model_config::ModelConfigError>(())
1055/// ```
1056pub fn resolve_model_config(
1057    path: impl AsRef<Path>,
1058    signals: Option<ModelSignals>,
1059    overrides: &ModelOverrides,
1060) -> Result<ResolvedModelConfig, ModelConfigError> {
1061    let path = path.as_ref();
1062    let embedded = super::read_embedded(path)?;
1063    let dir = path
1064        .parent()
1065        .filter(|d| !d.as_os_str().is_empty())
1066        .unwrap_or(Path::new("."));
1067    let sidecar = super::read_sidecars(dir)?;
1068    let mut sources = ModelSources::new()
1069        .with_embedded(embedded)
1070        .with_sidecar(sidecar);
1071    sources.signals = signals;
1072    resolve_sources(sources, overrides)
1073}
1074
1075#[cfg(test)]
1076mod tests {
1077    use super::*;
1078    use crate::model_config::heuristics::tests::{segformer_signals, tensor};
1079    use crate::model_config::ConfigFile;
1080    use crate::{ModelSource, TensorInfo};
1081
1082    fn file(name: &str, text: &str) -> Option<ConfigFile> {
1083        Some(ConfigFile {
1084            name: name.into(),
1085            text: text.into(),
1086        })
1087    }
1088
1089    fn fixture(model: &str, name: &str) -> String {
1090        let path = format!(
1091            "{}/../../testdata/hf/{model}/{name}",
1092            env!("CARGO_MANIFEST_DIR")
1093        );
1094        std::fs::read_to_string(path).unwrap()
1095    }
1096
1097    fn hf_sidecar(model: &str) -> ModelFiles {
1098        ModelFiles {
1099            hf_config: file("config.json", &fixture(model, "config.json")),
1100            hf_preprocessor: file(
1101                "preprocessor_config.json",
1102                &fixture(model, "preprocessor_config.json"),
1103            ),
1104            ..Default::default()
1105        }
1106    }
1107
1108    const SEGFORMER_EF: &str = r#"{
1109      "schema_version": 2,
1110      "dataset": {"classes": ["bg", "fg"]},
1111      "input": {"shape": [1, 3, 512, 512],
1112                "dshape": [{"batch": 1}, {"num_features": 3}, {"height": 512}, {"width": 512}],
1113                "preprocessing": {"letterbox": {"placement": "topleft", "pad": [0], "pad_space": "normalized"}}},
1114      "outputs": [{"name": "logits", "type": "segmentation", "decoder": "semantic",
1115                   "shape": [1, 2, 128, 128],
1116                   "dshape": [{"batch": 1}, {"num_classes": 2}, {"height": 128}, {"width": 128}],
1117                   "dtype": "float32", "stride": 4}]
1118    }"#;
1119
1120    fn sources(
1121        embedded: ModelFiles,
1122        sidecar: ModelFiles,
1123        signals: Option<ModelSignals>,
1124    ) -> ModelSources {
1125        ModelSources {
1126            embedded,
1127            sidecar,
1128            signals,
1129        }
1130    }
1131
1132    fn resolve(s: ModelSources) -> ResolvedModelConfig {
1133        resolve_sources(s, &ModelOverrides::default()).unwrap()
1134    }
1135
1136    #[test]
1137    fn edgefirst_wins_over_hf_field_by_field() {
1138        let embedded = ModelFiles {
1139            edgefirst_json: file("edgefirst.json", SEGFORMER_EF),
1140            ..Default::default()
1141        };
1142        let c = resolve(sources(
1143            embedded,
1144            hf_sidecar("segformer-512"),
1145            Some(segformer_signals(2)),
1146        ));
1147        assert_eq!(c.family, ModelFamily::EdgeFirst);
1148        assert_eq!(c.labels, ["bg", "fg"]);
1149        assert_eq!(c.preprocessing.placement, LetterboxPlacement::TopLeft);
1150        assert_eq!(c.preprocessing.fit, FitMode::Letterbox);
1151        assert_eq!(c.preprocessing.pad_space, PadSpace::Normalized);
1152        // edgefirst.json declares no normalization: the HF processor fills it.
1153        assert_eq!(c.preprocessing.normalization, Normalization::imagenet());
1154        assert_eq!(
1155            c.source(ConfigField::Normalization).unwrap().source,
1156            ConfigSource::SidecarHf
1157        );
1158        assert_eq!(
1159            c.source(ConfigField::Placement).unwrap().to_string(),
1160            "embedded edgefirst.json"
1161        );
1162        assert_eq!(c.task, ModelTask::Semantic);
1163    }
1164
1165    #[test]
1166    fn embedded_beats_sidecar() {
1167        let embedded = ModelFiles {
1168            labels_txt: file("labels.txt", "a\nb\n"),
1169            ..Default::default()
1170        };
1171        let sidecar = ModelFiles {
1172            edgefirst_json: file("edgefirst.json", SEGFORMER_EF),
1173            labels_txt: file("labels.txt", "x\ny\n"),
1174            ..Default::default()
1175        };
1176        let c = resolve(sources(embedded, sidecar, None));
1177        assert_eq!(c.labels, ["a", "b"]);
1178        assert_eq!(
1179            c.source(ConfigField::Labels).unwrap().source,
1180            ConfigSource::EmbeddedEdgefirst
1181        );
1182        assert_eq!(
1183            c.source(ConfigField::Task).unwrap().source,
1184            ConfigSource::SidecarEdgefirst
1185        );
1186    }
1187
1188    #[test]
1189    fn overrides_apply_last() {
1190        let sources = sources(
1191            ModelFiles::default(),
1192            hf_sidecar("segformer-512"),
1193            Some(segformer_signals(150)),
1194        );
1195        let mut o = ModelOverrides::default();
1196        o.set("normalize", "signed").unwrap();
1197        o.set("resize", "letterbox").unwrap();
1198        o.set("decode", "fast").unwrap();
1199        let c = resolve_sources(sources, &o).unwrap();
1200        assert_eq!(c.preprocessing.normalization, Normalization::signed());
1201        assert_eq!(c.preprocessing.placement, LetterboxPlacement::TopLeft);
1202        assert_eq!(c.preprocessing.pad, [0.0; 3]);
1203        assert_eq!(c.decode_mode, Some(DecodeMode::Fast));
1204        assert_eq!(
1205            c.schema.outputs[0].semantic.decode_mode,
1206            Some(DecodeMode::Fast)
1207        );
1208        for f in [
1209            ConfigField::Normalization,
1210            ConfigField::Resize,
1211            ConfigField::DecodeMode,
1212        ] {
1213            assert_eq!(c.source(f).unwrap().source, ConfigSource::Override, "{f:?}");
1214        }
1215        assert!(c
1216            .notes
1217            .iter()
1218            .any(|n| n.contains("Hugging Face convention")));
1219    }
1220
1221    #[test]
1222    fn validation_fallback_fills_unset_fields() {
1223        let mut v: Value = serde_json::from_str(SEGFORMER_EF).unwrap();
1224        v["input"].as_object_mut().unwrap().remove("preprocessing");
1225        v["validation"] = serde_json::json!({"normalization": "signed", "preprocessing": "resize"});
1226        let c = resolve(sources(
1227            ModelFiles {
1228                edgefirst_json: file("edgefirst.json", &v.to_string()),
1229                ..Default::default()
1230            },
1231            ModelFiles::default(),
1232            None,
1233        ));
1234        assert_eq!(c.preprocessing.normalization, Normalization::signed());
1235        assert_eq!(c.preprocessing.fit, FitMode::Stretch);
1236        assert_eq!(
1237            c.source(ConfigField::Resize).unwrap().to_string(),
1238            "embedded edgefirst.json validation.preprocessing"
1239        );
1240        assert_eq!(
1241            c.source(ConfigField::Placement).unwrap().source,
1242            ConfigSource::Default
1243        );
1244    }
1245
1246    #[test]
1247    fn segformer_config_resolves_semantic_schema() {
1248        let c = resolve(sources(
1249            ModelFiles::default(),
1250            hf_sidecar("segformer-512"),
1251            Some(segformer_signals(150)),
1252        ));
1253        assert_eq!(c.family, ModelFamily::HuggingFace);
1254        assert_eq!(c.task, ModelTask::Semantic);
1255        assert_eq!(c.labels.len(), 150);
1256        assert_eq!(c.preprocessing.fit, FitMode::Stretch);
1257        assert_eq!(c.preprocessing.size, Some((512, 512)));
1258        assert_eq!(c.reduce_labels, Some(true));
1259        c.schema.validate().unwrap();
1260        crate::semantic::SemanticConfig::from_logical(&c.schema.outputs[0]).unwrap();
1261        let pre = c.schema.input.unwrap().preprocessing.unwrap();
1262        assert_eq!(pre.resize, Some(FitMode::Stretch));
1263        assert!(c.description.contains("SegformerForSemanticSegmentation"));
1264    }
1265
1266    #[test]
1267    fn heuristic_segformer_dev_io() {
1268        let c = resolve(sources(
1269            ModelFiles::default(),
1270            ModelFiles::default(),
1271            Some(segformer_signals(150)),
1272        ));
1273        assert_eq!(c.family, ModelFamily::HuggingFace);
1274        for f in [
1275            ConfigField::Normalization,
1276            ConfigField::Resize,
1277            ConfigField::Interpolation,
1278        ] {
1279            let p = c.source(f).unwrap();
1280            assert_eq!(p.source, ConfigSource::Heuristic, "{f:?}");
1281            assert_eq!(p.to_string(), "guess: SegFormer default");
1282        }
1283        assert_eq!(c.preprocessing.normalization, Normalization::imagenet());
1284        assert!(c.labels.is_empty());
1285        assert!(c.notes.iter().any(|n| n.contains("labels unknown")));
1286    }
1287
1288    #[test]
1289    fn bare_letterbox_follows_family() {
1290        let mut o = ModelOverrides::default();
1291        o.set("resize", "letterbox").unwrap();
1292        // Unknown family: a semantic schema built from the task override.
1293        o.set("task", "semantic").unwrap();
1294        let mut s = segformer_signals(19);
1295        s.outputs[0].name = "out".into();
1296        let c =
1297            resolve_sources(sources(Default::default(), Default::default(), Some(s)), &o).unwrap();
1298        assert_eq!(c.family, ModelFamily::Unknown);
1299        assert_eq!(c.preprocessing.placement, LetterboxPlacement::Center);
1300        assert_eq!(
1301            (c.preprocessing.pad, c.preprocessing.pad_space),
1302            ([114.0; 3], PadSpace::Pixel)
1303        );
1304
1305        let mut o = ModelOverrides::default();
1306        o.set("resize", "letterbox-center").unwrap();
1307        let c = resolve_sources(
1308            sources(
1309                Default::default(),
1310                Default::default(),
1311                Some(segformer_signals(19)),
1312            ),
1313            &o,
1314        )
1315        .unwrap();
1316        assert_eq!(c.preprocessing.placement, LetterboxPlacement::Center);
1317        assert_eq!(
1318            c.preprocessing.pad_space,
1319            PadSpace::Normalized,
1320            "keeps the HF pad"
1321        );
1322    }
1323
1324    fn ultralytics_signals() -> ModelSignals {
1325        let mut metadata = BTreeMap::new();
1326        metadata.insert("author".into(), "Ultralytics".into());
1327        metadata.insert("task".into(), "detect".into());
1328        metadata.insert("names".into(), "{0: 'person', 1: 'car'}".into());
1329        ModelSignals {
1330            source: ModelSource::Onnx,
1331            inputs: vec![tensor("images", &[1, 3, 640, 640])],
1332            outputs: vec![TensorInfo {
1333                name: "output0".into(),
1334                shape: vec![1, 6, 8400],
1335                dtype: schema::DType::Float32,
1336                quantization: None,
1337            }],
1338            metadata,
1339        }
1340    }
1341
1342    #[test]
1343    fn ultralytics_signals_resolve() {
1344        let c = resolve(sources(
1345            Default::default(),
1346            Default::default(),
1347            Some(ultralytics_signals()),
1348        ));
1349        assert_eq!(c.family, ModelFamily::Ultralytics);
1350        assert_eq!(c.task, ModelTask::Detect);
1351        assert_eq!(c.labels, ["person", "car"]);
1352        assert_eq!(c.preprocessing.pad, [114.0; 3]);
1353        assert_eq!(
1354            c.source(ConfigField::Labels).unwrap().to_string(),
1355            "ultralytics metadata"
1356        );
1357        assert_eq!(
1358            c.source(ConfigField::Pad).unwrap().to_string(),
1359            "default: Ultralytics convention"
1360        );
1361
1362        // Metadata embedded in the container reaches inference.
1363        let mut s = ultralytics_signals();
1364        let embedded = ModelFiles {
1365            metadata: std::mem::take(&mut s.metadata),
1366            ..Default::default()
1367        };
1368        let c = resolve(sources(embedded, Default::default(), Some(s)));
1369        assert_eq!(c.family, ModelFamily::Ultralytics);
1370
1371        let mut o = ModelOverrides::default();
1372        o.set("decode", "fast").unwrap();
1373        let e = resolve_sources(
1374            sources(
1375                Default::default(),
1376                Default::default(),
1377                Some(ultralytics_signals()),
1378            ),
1379            &o,
1380        );
1381        assert!(matches!(e, Err(ModelConfigError::Override(_))));
1382    }
1383
1384    #[test]
1385    fn ultralytics_names_label_an_edgefirst_schema() {
1386        let ef = r#"{"schema_version": 2,
1387          "outputs": [{"name": "output0", "type": "detection", "decoder": "ultralytics",
1388                       "shape": [1, 6, 8400], "dtype": "float32"}]}"#;
1389        let embedded = ModelFiles {
1390            edgefirst_json: file("edgefirst.json", ef),
1391            ..Default::default()
1392        };
1393        let c = resolve(sources(
1394            embedded,
1395            Default::default(),
1396            Some(ultralytics_signals()),
1397        ));
1398        assert_eq!(c.family, ModelFamily::EdgeFirst);
1399        assert_eq!(c.labels, ["person", "car"]);
1400        assert_eq!(
1401            c.source(ConfigField::Labels).unwrap().source,
1402            ConfigSource::UltralyticsMetadata
1403        );
1404    }
1405
1406    #[test]
1407    fn unresolved_names_missing_piece_and_overrides() {
1408        let mut s = segformer_signals(19);
1409        s.inputs[0].name = "input".into();
1410        let e = resolve_sources(
1411            sources(Default::default(), Default::default(), Some(s)),
1412            &ModelOverrides::default(),
1413        )
1414        .unwrap_err();
1415        let text = e.to_string();
1416        assert!(matches!(e, ModelConfigError::Unresolved { .. }), "{text}");
1417        assert!(text.contains("task=semantic"), "{text}");
1418
1419        let e = resolve_sources(
1420            sources(Default::default(), hf_sidecar("segformer-512"), None),
1421            &ModelOverrides::default(),
1422        )
1423        .unwrap_err();
1424        assert!(matches!(e, ModelConfigError::Unresolved { .. }), "{e}");
1425    }
1426
1427    #[test]
1428    fn hf_labels_that_miscount_the_output_are_ignored() {
1429        let c = resolve(sources(
1430            Default::default(),
1431            hf_sidecar("segformer-512"),
1432            Some(segformer_signals(19)),
1433        ));
1434        assert!(c.labels.is_empty());
1435        assert!(
1436            c.notes
1437                .iter()
1438                .any(|n| n.contains("name 150 classes but the model has 19")),
1439            "{:?}",
1440            c.notes
1441        );
1442    }
1443
1444    #[test]
1445    fn foreign_config_json_is_noted_and_skipped() {
1446        let sidecar = ModelFiles {
1447            hf_config: file("config.json", r#"{"learning_rate": 0.1}"#),
1448            ..Default::default()
1449        };
1450        let c = resolve(sources(
1451            Default::default(),
1452            sidecar,
1453            Some(segformer_signals(19)),
1454        ));
1455        assert!(c
1456            .notes
1457            .iter()
1458            .any(|n| n.contains("config.json present but not a Hugging Face config")));
1459        assert_eq!(
1460            c.source(ConfigField::Task).unwrap().source,
1461            ConfigSource::Heuristic
1462        );
1463    }
1464
1465    #[test]
1466    fn task_override_conflict_is_error() {
1467        let mut o = ModelOverrides::default();
1468        o.set("task", "detect").unwrap();
1469        let e = resolve_sources(
1470            sources(
1471                Default::default(),
1472                Default::default(),
1473                Some(segformer_signals(19)),
1474            ),
1475            &o,
1476        );
1477        assert!(matches!(e, Err(ModelConfigError::Override(_))));
1478        let mut o = ModelOverrides::default();
1479        o.set("labels", "a,b").unwrap();
1480        let e = resolve_sources(
1481            sources(
1482                Default::default(),
1483                Default::default(),
1484                Some(segformer_signals(19)),
1485            ),
1486            &o,
1487        );
1488        assert!(
1489            matches!(
1490                e,
1491                Err(ModelConfigError::LabelCountMismatch {
1492                    expected: 19,
1493                    found: 2
1494                })
1495            ),
1496            "label count is checked: {e:?}"
1497        );
1498    }
1499
1500    #[test]
1501    fn describe_lines_match_format() {
1502        let c = resolve(sources(
1503            Default::default(),
1504            hf_sidecar("segformer-512"),
1505            Some(segformer_signals(150)),
1506        ));
1507        let lines = c.describe();
1508        assert_eq!(
1509            lines[0], "normalize=imagenet [preprocessor_config.json]",
1510            "{lines:?}"
1511        );
1512        assert_eq!(lines[1], "resize=stretch [preprocessor_config.json]");
1513        assert!(
1514            lines.contains(&"size=512x512 [model input]".to_string()),
1515            "{lines:?}"
1516        );
1517        assert!(
1518            lines.contains(&"task=semantic [config.json]".to_string()),
1519            "{lines:?}"
1520        );
1521        assert!(lines.contains(&"labels=150 classes [config.json]".to_string()));
1522        assert!(lines.iter().any(|l| l.starts_with("decode=accurate")));
1523
1524        let mut o = ModelOverrides::default();
1525        o.set("resize", "letterbox-topleft").unwrap();
1526        let c = resolve_sources(
1527            sources(
1528                Default::default(),
1529                Default::default(),
1530                Some(segformer_signals(150)),
1531            ),
1532            &o,
1533        )
1534        .unwrap();
1535        let lines = c.describe();
1536        assert!(
1537            lines.contains(&"placement=topleft [override]".to_string()),
1538            "{lines:?}"
1539        );
1540        assert!(lines.contains(&"normalize=imagenet [guess: SegFormer default]".to_string()));
1541        assert!(lines.contains(&"resize=letterbox-topleft [override]".to_string()));
1542    }
1543
1544    /// Every override key `describe()` prints parses back to the same value.
1545    #[test]
1546    fn describe_round_trips_through_overrides() {
1547        let mut sidecar = hf_sidecar("segformer-512");
1548        sidecar.hf_preprocessor = file(
1549            "preprocessor_config.json",
1550            r#"{"image_processor_type": "SegformerImageProcessor", "do_pad": true,
1551                "image_mean": [0.25, 0.5, 0.75], "image_std": [0.5, 0.25, 0.125]}"#,
1552        );
1553        let c = resolve(sources(
1554            Default::default(),
1555            sidecar,
1556            Some(segformer_signals(150)),
1557        ));
1558        let mut o = ModelOverrides::new();
1559        for line in c.describe() {
1560            let (key, rest) = line.split_once('=').unwrap();
1561            let value = rest.split(" [").next().unwrap();
1562            if ["normalize", "resize", "task", "decode"].contains(&key) {
1563                o.set(key, value).unwrap_or_else(|e| panic!("{line}: {e}"));
1564            }
1565        }
1566        let again = resolve_sources(
1567            sources(
1568                Default::default(),
1569                hf_sidecar("segformer-512"),
1570                Some(segformer_signals(150)),
1571            ),
1572            &o,
1573        )
1574        .unwrap();
1575        assert_eq!(again.preprocessing.fit, c.preprocessing.fit);
1576        assert_eq!(again.preprocessing.placement, c.preprocessing.placement);
1577        let (a, b) = (
1578            again.preprocessing.normalization,
1579            c.preprocessing.normalization,
1580        );
1581        for (x, y) in a
1582            .mean()
1583            .iter()
1584            .chain(&a.std())
1585            .zip(b.mean().iter().chain(&b.std()))
1586        {
1587            assert!((x - y).abs() < 1e-3, "{a:?} vs {b:?}");
1588        }
1589    }
1590
1591    /// An `edgefirst.json` with no outputs supplies preprocessing but not the
1592    /// schema.
1593    #[test]
1594    fn edgefirst_without_outputs_leaves_the_schema_to_hf() {
1595        let sidecar = ModelFiles {
1596            edgefirst_json: file(
1597                "edgefirst.json",
1598                r#"{"schema_version": 2, "input": {"shape": [1, 3, 512, 512], "preprocessing": {"normalization": "signed"}}}"#,
1599            ),
1600            ..hf_sidecar("segformer-512")
1601        };
1602        let c = resolve(sources(
1603            Default::default(),
1604            sidecar,
1605            Some(segformer_signals(150)),
1606        ));
1607        assert_eq!(c.task, ModelTask::Semantic);
1608        assert_eq!(c.family, ModelFamily::HuggingFace);
1609        assert_eq!(c.preprocessing.normalization, Normalization::signed());
1610        assert!(
1611            c.notes.iter().any(|n| n.contains("declares no outputs")),
1612            "{:?}",
1613            c.notes
1614        );
1615        let empty = ModelFiles {
1616            edgefirst_json: file("edgefirst.json", "{}"),
1617            ..hf_sidecar("segformer-512")
1618        };
1619        let c = resolve(sources(
1620            Default::default(),
1621            empty,
1622            Some(segformer_signals(150)),
1623        ));
1624        assert_eq!(c.task, ModelTask::Semantic);
1625    }
1626
1627    /// Discovered labels that miscount the classes are noted and the next
1628    /// source is used; an override that miscounts is an error.
1629    #[test]
1630    fn labels_txt_count_mismatch_falls_back() {
1631        let sidecar = ModelFiles {
1632            labels_txt: file("labels.txt", "a\nb\n"),
1633            ..hf_sidecar("segformer-512")
1634        };
1635        let c = resolve(sources(
1636            Default::default(),
1637            sidecar,
1638            Some(segformer_signals(150)),
1639        ));
1640        assert_eq!(c.labels.len(), 150);
1641        assert_eq!(
1642            c.provenance[&ConfigField::Labels].source,
1643            ConfigSource::SidecarHf
1644        );
1645        assert!(
1646            c.notes
1647                .iter()
1648                .any(|n| n.contains("name 2 classes but the model has 150; they are ignored")),
1649            "{:?}",
1650            c.notes
1651        );
1652        let mut o = ModelOverrides::new();
1653        o.labels = Some(vec!["a".into(), " ".into()]);
1654        let e = resolve_sources(
1655            sources(
1656                Default::default(),
1657                hf_sidecar("segformer-512"),
1658                Some(segformer_signals(150)),
1659            ),
1660            &o,
1661        );
1662        assert!(
1663            matches!(&e, Err(ModelConfigError::Override(m)) if m.contains("empty name at 1")),
1664            "an empty name is not a count mismatch: {e:?}"
1665        );
1666    }
1667
1668    #[test]
1669    fn legacy_segformer_resolves_trimmed_labels() {
1670        let c = resolve(sources(
1671            Default::default(),
1672            hf_sidecar("segformer-legacy-512"),
1673            Some(segformer_signals(150)),
1674        ));
1675        assert_eq!(c.labels.len(), 150);
1676        assert_eq!(c.labels[7], "bed");
1677        assert_eq!(c.labels[33], "chest of drawers");
1678        assert_eq!(c.preprocessing.fit, FitMode::Stretch);
1679        assert_eq!(c.preprocessing.size, Some((512, 512)));
1680        assert_eq!(c.reduce_labels, Some(true));
1681        assert!(
1682            !c.notes.iter().any(|n| n.contains("repeat")),
1683            "{:?}",
1684            c.notes
1685        );
1686    }
1687
1688    fn edgefirst_with_classes(classes: &str) -> ModelSources {
1689        let ef = SEGFORMER_EF.replace(r#"["bg", "fg"]"#, classes);
1690        let embedded = ModelFiles {
1691            edgefirst_json: file("edgefirst.json", &ef),
1692            ..Default::default()
1693        };
1694        sources(embedded, Default::default(), None)
1695    }
1696
1697    fn ultralytics_with_names(names: &str) -> ModelSources {
1698        let mut s = ultralytics_signals();
1699        s.metadata.insert("names".into(), names.into());
1700        sources(Default::default(), Default::default(), Some(s))
1701    }
1702
1703    /// Asserts that the labels were ignored with a note naming `field`.
1704    fn assert_ignored_labels(c: &ResolvedModelConfig, field: &str) {
1705        let ignored = format!("field `{field}` is empty after trimming whitespace; its class");
1706        assert!(
1707            c.notes.iter().any(|n| n.contains(&ignored)),
1708            "{:?}",
1709            c.notes
1710        );
1711        assert!(
1712            c.notes.iter().any(|n| n.contains("labels unknown")),
1713            "{:?}",
1714            c.notes
1715        );
1716        assert!(c.labels.is_empty());
1717    }
1718
1719    #[test]
1720    fn labels_are_trimmed_from_every_source() {
1721        let none = ModelOverrides::default();
1722        let c = resolve(edgefirst_with_classes(r#"[" bg", "fg\t"]"#));
1723        assert_eq!(c.labels, ["bg", "fg"]);
1724        let c = resolve(edgefirst_with_classes(r#"["bg", " "]"#));
1725        assert_ignored_labels(&c, "dataset.classes[1]");
1726
1727        let c = resolve(ultralytics_with_names("{0: ' person', 1: 'car '}"));
1728        assert_eq!(c.labels, ["person", "car"]);
1729        let c = resolve(ultralytics_with_names("{0: ' ', 1: 'car'}"));
1730        assert_ignored_labels(&c, "names[0]");
1731
1732        let mut s = edgefirst_with_classes(r#"["bg", " "]"#);
1733        s.sidecar.labels_txt = file("labels.txt", "a\nb\n");
1734        let c = resolve_sources(s, &none).unwrap();
1735        assert_eq!(c.labels, ["a", "b"]);
1736
1737        let mut o = ModelOverrides::new();
1738        o.labels = Some(vec![" a".into(), "b c ".into()]);
1739        let c = resolve_sources(ultralytics_with_names("{0: 'x', 1: 'y'}"), &o).unwrap();
1740        assert_eq!(c.labels, ["a", "b c"]);
1741    }
1742
1743    #[test]
1744    fn miscounted_edgefirst_classes_fall_back_to_labels_txt() {
1745        let mut s = edgefirst_with_classes(r#"["bg"]"#);
1746        s.sidecar.labels_txt = file("labels.txt", "a\nb\n");
1747        let c = resolve(s);
1748        assert_eq!(c.labels, ["a", "b"]);
1749        assert!(
1750            c.notes
1751                .iter()
1752                .any(|n| n.contains("name 1 classes but the model has 2")),
1753            "{:?}",
1754            c.notes
1755        );
1756    }
1757
1758    #[test]
1759    fn miscounted_ultralytics_names_are_ignored() {
1760        let c = resolve(ultralytics_with_names("{0: 'person'}"));
1761        assert_eq!(c.family, ModelFamily::Ultralytics);
1762        assert!(c.labels.is_empty());
1763        assert!(
1764            c.notes
1765                .iter()
1766                .any(|n| n.contains("names lists 1 classes but the output has 2")),
1767            "{:?}",
1768            c.notes
1769        );
1770        let mut o = ModelOverrides::new();
1771        o.labels = Some(vec!["a".into(), "b".into()]);
1772        let c = resolve_sources(ultralytics_with_names("{0: 'person'}"), &o).unwrap();
1773        assert_eq!(c.labels, ["a", "b"]);
1774        o.labels = Some(vec!["a".into()]);
1775        let e = resolve_sources(ultralytics_with_names("{0: 'person'}"), &o).unwrap_err();
1776        assert!(
1777            matches!(
1778                e,
1779                ModelConfigError::LabelCountMismatch {
1780                    expected: 2,
1781                    found: 1
1782                }
1783            ),
1784            "{e}"
1785        );
1786    }
1787
1788    /// An `edgefirst.json` naming three classes for `outputs`, beside a
1789    /// two-line `labels.txt`.
1790    fn detection_with_classes(outputs: &str) -> ResolvedModelConfig {
1791        let ef = format!(
1792            r#"{{"schema_version": 2, "dataset": {{"classes": ["a", "b", "c"]}},
1793                 "outputs": {outputs}}}"#
1794        );
1795        let embedded = ModelFiles {
1796            edgefirst_json: file("edgefirst.json", &ef),
1797            ..Default::default()
1798        };
1799        let sidecar = ModelFiles {
1800            labels_txt: file("labels.txt", "x\ny\n"),
1801            ..Default::default()
1802        };
1803        resolve(sources(embedded, sidecar, None))
1804    }
1805
1806    #[test]
1807    fn detection_labels_fall_back_on_a_class_count_mismatch() {
1808        let ultralytics = r#"[{"name": "output0", "type": "detection", "decoder": "ultralytics",
1809            "shape": [1, 6, 8400],
1810            "dshape": [{"batch": 1}, {"num_features": 6}, {"num_boxes": 8400}]}]"#;
1811        let segdet = r#"[{"name": "output0", "type": "detection", "decoder": "ultralytics",
1812            "shape": [1, 38, 8400],
1813            "dshape": [{"batch": 1}, {"num_features": 38}, {"num_boxes": 8400}]},
1814            {"name": "protos", "type": "protos", "decoder": "ultralytics",
1815            "shape": [1, 32, 160, 160],
1816            "dshape": [{"batch": 1}, {"num_protos": 32}, {"height": 160}, {"width": 160}]}]"#;
1817        let modelpack = r#"[{"name": "det", "type": "detection", "decoder": "modelpack",
1818            "anchors": [[0.1, 0.1], [0.2, 0.2], [0.3, 0.3]],
1819            "shape": [1, 40, 40, 21],
1820            "dshape": [{"batch": 1}, {"height": 40}, {"width": 40},
1821                       {"num_anchors_x_features": 21}]}]"#;
1822        let scores = r#"[{"name": "boxes", "type": "boxes", "decoder": "ultralytics",
1823            "shape": [1, 4, 8400],
1824            "dshape": [{"batch": 1}, {"box_coords": 4}, {"num_boxes": 8400}]},
1825            {"name": "scores", "type": "scores", "decoder": "ultralytics",
1826            "shape": [1, 2, 8400],
1827            "dshape": [{"batch": 1}, {"num_classes": 2}, {"num_boxes": 8400}]}]"#;
1828        for outputs in [ultralytics, segdet, modelpack, scores] {
1829            let c = detection_with_classes(outputs);
1830            assert_eq!(c.labels, ["x", "y"], "{outputs}");
1831            assert!(
1832                c.notes
1833                    .iter()
1834                    .any(|n| n.contains("name 3 classes but the model has 2; they are ignored")),
1835                "{outputs}: {:?}",
1836                c.notes
1837            );
1838        }
1839    }
1840
1841    fn label_count_error(outputs: &str, labels: &[&str]) -> Option<(usize, usize)> {
1842        let ef = format!(r#"{{"schema_version": 2, "outputs": {outputs}}}"#);
1843        let embedded = ModelFiles {
1844            edgefirst_json: file("edgefirst.json", &ef),
1845            ..Default::default()
1846        };
1847        let mut o = ModelOverrides::new();
1848        o.labels = Some(labels.iter().map(|&l| l.to_owned()).collect());
1849        match resolve_sources(sources(embedded, Default::default(), None), &o) {
1850            Err(ModelConfigError::LabelCountMismatch { expected, found }) => {
1851                Some((expected, found))
1852            }
1853            Ok(_) => None,
1854            Err(e) => panic!("{e}"),
1855        }
1856    }
1857
1858    /// The count follows the decoder builder: an omitted `dshape` uses the
1859    /// canonical shape and an omitted `decoder` means Ultralytics.
1860    #[test]
1861    fn class_count_without_dshape_or_decoder() {
1862        let no_dshape = r#"[{"name": "output0", "type": "detection", "decoder": "ultralytics",
1863            "shape": [1, 6, 8400]}]"#;
1864        assert_eq!(label_count_error(no_dshape, &["a"]), Some((2, 1)));
1865        assert_eq!(label_count_error(no_dshape, &["a", "b"]), None);
1866        let no_decoder = r#"[{"name": "output0", "type": "detection",
1867            "shape": [1, 6, 8400],
1868            "dshape": [{"batch": 1}, {"num_features": 6}, {"num_boxes": 8400}]}]"#;
1869        assert_eq!(label_count_error(no_decoder, &["a"]), Some((2, 1)));
1870    }
1871
1872    /// Separate mask coefficients are not part of the detection features.
1873    #[test]
1874    fn two_way_segdet_counts_detection_features_only() {
1875        let c = detection_with_classes(
1876            r#"[{"name": "output0", "type": "detection", "decoder": "ultralytics",
1877                "shape": [1, 6, 8400],
1878                "dshape": [{"batch": 1}, {"num_features": 6}, {"num_boxes": 8400}]},
1879               {"name": "coefs", "type": "mask_coefs", "decoder": "ultralytics",
1880                "shape": [1, 32, 8400],
1881                "dshape": [{"batch": 1}, {"num_protos": 32}, {"num_boxes": 8400}]},
1882               {"name": "protos", "type": "protos", "decoder": "ultralytics",
1883                "shape": [1, 32, 160, 160],
1884                "dshape": [{"batch": 1}, {"num_protos": 32}, {"height": 160}, {"width": 160}]}]"#,
1885        );
1886        assert_eq!(c.labels, ["x", "y"], "{:?}", c.notes);
1887    }
1888
1889    /// Legacy `detection` outputs that the builder classifies as end-to-end
1890    /// carry no class count either.
1891    #[test]
1892    fn legacy_end_to_end_detection_labels_are_not_counted() {
1893        for version in [r#""decoder_version": "yolo26","#, ""] {
1894            let outputs = r#"[{"name": "output0", "type": "detection", "decoder": "ultralytics",
1895                "shape": [1, 300, 6],
1896                "dshape": [{"batch": 1}, {"num_boxes": 300}, {"num_features": 6}]}]"#;
1897            let ef = format!(
1898                r#"{{"schema_version": 2, {version} "dataset": {{"classes": ["a", "b", "c"]}},
1899                     "outputs": {outputs}}}"#
1900            );
1901            let embedded = ModelFiles {
1902                edgefirst_json: file("edgefirst.json", &ef),
1903                ..Default::default()
1904            };
1905            let c = resolve(sources(embedded, Default::default(), None));
1906            assert_eq!(c.labels, ["a", "b", "c"], "{version}: {:?}", c.notes);
1907        }
1908    }
1909
1910    /// A ModelPack multitask model counts its detection classes, not the
1911    /// segmentation channels that add background, whatever the output order.
1912    #[test]
1913    fn modelpack_multitask_counts_detection_classes_in_any_order() {
1914        let path = format!(
1915            "{}/../../testdata/decoder/modelpack_multitask_logical.json",
1916            env!("CARGO_MANIFEST_DIR")
1917        );
1918        let mut ef: Value = serde_json::from_str(&std::fs::read_to_string(path).unwrap()).unwrap();
1919        for _ in 0..2 {
1920            let embedded = ModelFiles {
1921                edgefirst_json: file("edgefirst.json", &ef.to_string()),
1922                ..Default::default()
1923            };
1924            let c = resolve(sources(embedded, Default::default(), None));
1925            assert_eq!(
1926                c.labels,
1927                ["cat", "dog", "person", "car", "truck"],
1928                "{:?}",
1929                c.notes
1930            );
1931            let outputs = ef["outputs"].as_array_mut().unwrap();
1932            outputs.rotate_right(1);
1933        }
1934    }
1935
1936    fn ultralytics_with_outputs(names: &str, shapes: &[&[usize]]) -> ModelSources {
1937        let mut s = ultralytics_signals();
1938        s.metadata.insert("names".into(), names.into());
1939        s.metadata.insert("end2end".into(), "False".into());
1940        s.outputs = shapes
1941            .iter()
1942            .enumerate()
1943            .map(|(i, shape)| tensor(&format!("output{i}"), shape))
1944            .collect();
1945        sources(Default::default(), Default::default(), Some(s))
1946    }
1947
1948    /// Recovery from stale `names` takes the one output that lays out, in
1949    /// either order, and refuses a choice between two.
1950    #[test]
1951    fn stale_ultralytics_names_recover_only_from_one_candidate() {
1952        let boxes: &[usize] = &[1, 4, 8400];
1953        let detections: &[usize] = &[1, 6, 8400];
1954        for shapes in [[boxes, detections], [detections, boxes]] {
1955            let c = resolve(ultralytics_with_outputs("{0: 'person'}", &shapes));
1956            assert_eq!(c.family, ModelFamily::Ultralytics);
1957            assert!(c.labels.is_empty());
1958            assert!(
1959                c.notes
1960                    .iter()
1961                    .any(|n| n.contains("names lists 1 classes but the output has 2")),
1962                "{:?}",
1963                c.notes
1964            );
1965        }
1966        let aux: &[usize] = &[1, 32, 8400];
1967        let detections: &[usize] = &[1, 84, 8400];
1968        for shapes in [[aux, detections], [detections, aux]] {
1969            let e = resolve_sources(
1970                ultralytics_with_outputs("{0: 'a', 1: 'b'}", &shapes),
1971                &ModelOverrides::default(),
1972            )
1973            .unwrap_err();
1974            assert!(
1975                matches!(&e, ModelConfigError::Infer(InferError::AmbiguousLayout(m))
1976                    if m.contains("labels override") && m.contains("edgefirst.json")),
1977                "{e}"
1978            );
1979            let mut o = ModelOverrides::new();
1980            o.labels = Some((0..80).map(|i| format!("c{i}")).collect());
1981            let c =
1982                resolve_sources(ultralytics_with_outputs("{0: 'a', 1: 'b'}", &shapes), &o).unwrap();
1983            assert_eq!(c.labels.len(), 80);
1984            assert_eq!(c.schema.outputs[0].shape, detections, "{shapes:?}");
1985        }
1986    }
1987
1988    /// An empty override states no count, so recovery finds the real one and
1989    /// the override is then rejected against it.
1990    #[test]
1991    fn empty_labels_override_does_not_pick_the_layout() {
1992        let boxes: &[usize] = &[1, 4, 8400];
1993        let detections: &[usize] = &[1, 84, 8400];
1994        let mut o = ModelOverrides::new();
1995        o.labels = Some(Vec::new());
1996        let e = resolve_sources(
1997            ultralytics_with_outputs("{0: 'a'}", &[boxes, detections]),
1998            &o,
1999        )
2000        .unwrap_err();
2001        assert!(
2002            matches!(
2003                e,
2004                ModelConfigError::LabelCountMismatch {
2005                    expected: 80,
2006                    found: 0
2007                }
2008            ),
2009            "{e}"
2010        );
2011    }
2012
2013    /// Two identical candidates for one count are ambiguous, not skipped in
2014    /// favour of another count; a labels override naming the count settles it.
2015    #[test]
2016    fn ambiguity_within_one_count_is_reported() {
2017        let small: &[usize] = &[1, 6, 8400];
2018        let large: &[usize] = &[1, 84, 8400];
2019        let shapes = [small, small, large];
2020        let e = resolve_sources(
2021            ultralytics_with_outputs("{0: 'a'}", &shapes),
2022            &ModelOverrides::default(),
2023        )
2024        .unwrap_err();
2025        assert!(
2026            matches!(&e, ModelConfigError::Infer(InferError::AmbiguousLayout(m))
2027                if m.contains("labels override")),
2028            "{e}"
2029        );
2030        let mut o = ModelOverrides::new();
2031        o.labels = Some((0..80).map(|i| format!("c{i}")).collect());
2032        let c = resolve_sources(ultralytics_with_outputs("{0: 'a'}", &shapes), &o).unwrap();
2033        assert_eq!(c.schema.outputs[0].shape, large);
2034    }
2035
2036    /// ModelPack segmentation-only models are counted from their channel
2037    /// axis, with or without a `dshape`, less the trailing background channel.
2038    #[test]
2039    fn modelpack_segmentation_is_counted_without_background() {
2040        for outputs in [
2041            r#"[{"name": "seg", "type": "segmentation", "decoder": "modelpack",
2042                "shape": [1, 160, 160, 5]}]"#,
2043            r#"[{"name": "seg", "type": "segmentation", "decoder": "modelpack",
2044                "shape": [1, 160, 160, 5],
2045                "dshape": [{"batch": 1}, {"height": 160}, {"width": 160}, {"num_classes": 5}]}]"#,
2046        ] {
2047            assert_eq!(
2048                label_count_error(outputs, &["a", "b"]),
2049                Some((4, 2)),
2050                "{outputs}"
2051            );
2052            assert_eq!(
2053                label_count_error(outputs, &["a", "b", "c", "d"]),
2054                None,
2055                "{outputs}"
2056            );
2057        }
2058    }
2059
2060    /// End-to-end detections carry no class count, so the first source wins.
2061    #[test]
2062    fn end_to_end_detection_labels_are_not_counted() {
2063        let c = detection_with_classes(
2064            r#"[{"name": "output0", "type": "detections", "decoder": "ultralytics",
2065                "shape": [1, 300, 6],
2066                "dshape": [{"batch": 1}, {"num_boxes": 300}, {"num_features": 6}]}]"#,
2067        );
2068        assert_eq!(c.labels, ["a", "b", "c"]);
2069    }
2070
2071    #[test]
2072    fn repeated_labels_are_noted() {
2073        let c = resolve(ultralytics_with_names("{0: 'bed', 1: 'Bed '}"));
2074        assert_eq!(c.labels, ["bed", "Bed"]);
2075        assert!(
2076            c.notes
2077                .iter()
2078                .any(|n| n.contains("ultralytics metadata") && n.contains(r#""bed" at 0, 1"#)),
2079            "{:?}",
2080            c.notes
2081        );
2082    }
2083
2084    /// A reference size that differs in one axis is noted.
2085    #[test]
2086    fn width_only_size_mismatch_is_noted() {
2087        let mut sidecar = hf_sidecar("segformer-512");
2088        sidecar.hf_preprocessor = file(
2089            "preprocessor_config.json",
2090            r#"{"image_processor_type": "SegformerImageProcessor", "size": {"height": 512, "width": 256}}"#,
2091        );
2092        let c = resolve(sources(
2093            Default::default(),
2094            sidecar,
2095            Some(segformer_signals(150)),
2096        ));
2097        assert!(
2098            c.notes.iter().any(|n| n.contains("resizes to 256x512")),
2099            "{:?}",
2100            c.notes
2101        );
2102    }
2103
2104    #[test]
2105    fn bicubic_plan_falls_back_with_note() {
2106        let sidecar = ModelFiles {
2107            hf_config: file("config.json", &fixture("segformer-512", "config.json")),
2108            hf_preprocessor: file(
2109                "preprocessor_config.json",
2110                r#"{"image_processor_type": "SegformerImageProcessor", "resample": 3,
2111                    "image_mean": [0.5], "image_std": [0.5], "size": {"height": 512, "width": 512}}"#,
2112            ),
2113            ..Default::default()
2114        };
2115        let c = resolve(sources(
2116            Default::default(),
2117            sidecar,
2118            Some(segformer_signals(150)),
2119        ));
2120        assert_eq!(c.preprocessing.interpolation, Interpolation::Bicubic);
2121        assert!(c.notes.iter().any(|n| n.contains("interpolation=bicubic")));
2122    }
2123
2124    #[test]
2125    fn resolve_model_config_reads_sidecars() {
2126        let dir = std::env::temp_dir().join(format!("ef-resolve-{}", std::process::id()));
2127        std::fs::create_dir_all(&dir).unwrap();
2128        std::fs::write(dir.join("edgefirst.json"), SEGFORMER_EF).unwrap();
2129        let model = dir.join("model.bin");
2130        std::fs::write(&model, [0u8; 64]).unwrap();
2131        let c = resolve_model_config(&model, None, &ModelOverrides::default());
2132        std::fs::remove_dir_all(&dir).unwrap();
2133        let c = c.unwrap();
2134        assert_eq!(
2135            c.source(ConfigField::Task).unwrap().source,
2136            ConfigSource::SidecarEdgefirst
2137        );
2138        assert_eq!(c.preprocessing.size, Some((512, 512)));
2139    }
2140}