edgefirst_decoder/model_config/
container.rs

1// SPDX-FileCopyrightText: Copyright 2026 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4//! The configuration files a model carries, embedded in its container or
5//! beside it.
6
7use std::collections::BTreeMap;
8use std::fs::File;
9use std::io::{Cursor, Read, Seek};
10use std::path::Path;
11
12use super::error::invalid;
13use super::ModelConfigError;
14
15/// `edgefirst.json`, as a file name and the legacy ONNX metadata key.
16pub const EDGEFIRST_JSON: &str = "edgefirst.json";
17/// Class names, one per line.
18pub const LABELS_TXT: &str = "labels.txt";
19/// Hugging Face model config.
20pub const HF_CONFIG: &str = "config.json";
21/// Hugging Face image processor config.
22pub const HF_PREPROCESSOR: &str = "preprocessor_config.json";
23
24/// Largest sidecar file read.
25const MAX_SIDECAR: u64 = 16 << 20;
26
27/// ONNX `metadata_props` key for `edgefirst.json`.
28pub const ONNX_EDGEFIRST: &str = "edgefirst";
29/// ONNX `metadata_props` key for the class names.
30pub const ONNX_LABELS: &str = "labels";
31/// ONNX `metadata_props` key for the Hugging Face `config.json`.
32pub const ONNX_HF_CONFIG: &str = "hf_config";
33/// ONNX `metadata_props` key for the Hugging Face `preprocessor_config.json`.
34pub const ONNX_HF_PREPROCESSOR: &str = "hf_preprocessor_config";
35
36/// One configuration file and where it was found.
37#[derive(Debug, Clone, PartialEq, Eq)]
38#[non_exhaustive]
39pub struct ConfigFile {
40    /// The file name, or the ONNX metadata key it was stored under.
41    pub name: String,
42    /// The file contents.
43    pub text: String,
44}
45
46impl ConfigFile {
47    /// A file named `name` holding `text`.
48    pub fn new(name: impl Into<String>, text: impl Into<String>) -> Self {
49        Self {
50            name: name.into(),
51            text: text.into(),
52        }
53    }
54}
55
56/// The configuration files of one location: a model container or its
57/// directory.
58#[derive(Debug, Clone, Default, PartialEq, Eq)]
59#[non_exhaustive]
60pub struct ModelFiles {
61    /// `edgefirst.json`.
62    pub edgefirst_json: Option<ConfigFile>,
63    /// `labels.txt`.
64    pub labels_txt: Option<ConfigFile>,
65    /// Hugging Face `config.json`.
66    pub hf_config: Option<ConfigFile>,
67    /// Hugging Face `preprocessor_config.json`.
68    pub hf_preprocessor: Option<ConfigFile>,
69    /// Every ONNX metadata prop and ZIP text entry, by key or entry name.
70    /// Ultralytics metadata reaches schema inference through it.
71    pub metadata: BTreeMap<String, String>,
72    /// Model bytes before an appended ZIP trailer, which runtimes that check
73    /// the blob length need stripped.
74    pub payload_len: Option<u64>,
75}
76
77impl ModelFiles {
78    /// Whether nothing was found.
79    pub fn is_empty(&self) -> bool {
80        self.edgefirst_json.is_none()
81            && self.labels_txt.is_none()
82            && self.hf_config.is_none()
83            && self.hf_preprocessor.is_none()
84            && self.metadata.is_empty()
85    }
86
87    fn set(slot: &mut Option<ConfigFile>, name: &str, text: &str) {
88        if slot.is_none() {
89            *slot = Some(ConfigFile {
90                name: name.to_owned(),
91                text: text.to_owned(),
92            });
93        }
94    }
95
96    fn add_onnx_props(&mut self, props: BTreeMap<String, String>) {
97        for (key, value) in &props {
98            match key.as_str() {
99                ONNX_EDGEFIRST | EDGEFIRST_JSON => Self::set(&mut self.edgefirst_json, key, value),
100                ONNX_LABELS => Self::set(&mut self.labels_txt, key, value),
101                ONNX_HF_CONFIG => Self::set(&mut self.hf_config, key, value),
102                ONNX_HF_PREPROCESSOR => Self::set(&mut self.hf_preprocessor, key, value),
103                _ => {}
104            }
105        }
106        self.metadata.extend(props);
107    }
108
109    #[cfg(feature = "model-files")]
110    fn add_entries(&mut self, entries: Vec<(String, String)>) {
111        for (name, text) in entries {
112            let base = name.rsplit('/').next().unwrap_or(&name);
113            match base {
114                EDGEFIRST_JSON => Self::set(&mut self.edgefirst_json, base, &text),
115                LABELS_TXT => Self::set(&mut self.labels_txt, base, &text),
116                HF_CONFIG => Self::set(&mut self.hf_config, base, &text),
117                HF_PREPROCESSOR => Self::set(&mut self.hf_preprocessor, base, &text),
118                _ => {}
119            }
120            self.metadata.insert(name, text);
121        }
122    }
123}
124
125/// The configuration files embedded in the model file at `path`: ONNX
126/// `metadata_props` and, with the `model-files` feature, a ZIP trailer.
127/// Neither the graph nor the model payload is read.
128///
129/// # Errors
130///
131/// [`ModelConfigError::Io`] when the file cannot be read or is a
132/// directory, and [`ModelConfigError::Container`] when a `.onnx` file is not
133/// a valid `ModelProto`.
134pub fn read_embedded(path: impl AsRef<Path>) -> Result<ModelFiles, ModelConfigError> {
135    let path = path.as_ref();
136    let io = |source| ModelConfigError::Io {
137        path: path.to_owned(),
138        source,
139    };
140    // Windows refuses to open a directory rather than reporting one.
141    let mut file = File::open(path).map_err(|e| {
142        io(if path.is_dir() {
143            std::io::ErrorKind::IsADirectory.into()
144        } else {
145            e
146        })
147    })?;
148    let meta = file.metadata().map_err(io)?;
149    if meta.is_dir() {
150        return Err(io(std::io::ErrorKind::IsADirectory.into()));
151    }
152    let len = meta.len();
153    let onnx = path
154        .extension()
155        .is_some_and(|e| e.eq_ignore_ascii_case("onnx"));
156    read_from(&mut file, len, onnx)
157}
158
159/// [`read_embedded`] on a model already in memory. `onnx` says the bytes are
160/// an ONNX model, which makes a malformed `ModelProto` an error instead of
161/// "not ONNX".
162///
163/// # Errors
164///
165/// [`ModelConfigError::Container`] when `onnx` is set and the bytes are not
166/// a valid `ModelProto`.
167pub fn read_embedded_bytes(bytes: &[u8], onnx: bool) -> Result<ModelFiles, ModelConfigError> {
168    read_from(&mut Cursor::new(bytes), bytes.len() as u64, onnx)
169}
170
171fn read_from<R: Read + Seek>(
172    r: &mut R,
173    len: u64,
174    onnx: bool,
175) -> Result<ModelFiles, ModelConfigError> {
176    let mut files = ModelFiles::default();
177    #[cfg(feature = "model-files")]
178    if let Some(trailer) = super::ziptrailer::read_trailer(r, len)? {
179        files.payload_len = Some(trailer.payload_len);
180        files.add_entries(trailer.entries);
181    }
182    // An appended trailer is not part of the protobuf.
183    let model_len = files.payload_len.unwrap_or(len);
184    let mut first = [0u8; 1];
185    r.seek(std::io::SeekFrom::Start(0))
186        .map_err(|e| ModelConfigError::Container(e.to_string()))?;
187    let looks_onnx = model_len > 0 && r.read_exact(&mut first).is_ok() && first[0] == 0x08;
188    if onnx || looks_onnx {
189        match super::onnx::metadata_props(r, model_len) {
190            Ok(props) => files.add_onnx_props(props),
191            Err(e) if onnx => return Err(e),
192            Err(e) => log::debug!("not an ONNX model: {e}"),
193        }
194    }
195    Ok(files)
196}
197
198/// The configuration files beside a model, in `dir`: `edgefirst.json`,
199/// `labels.txt`, `config.json` and `preprocessor_config.json`.
200///
201/// # Errors
202///
203/// [`ModelConfigError::Io`] when a file exists but cannot be read, is not
204/// UTF-8 or is larger than 16 MiB.
205pub fn read_sidecars(dir: impl AsRef<Path>) -> Result<ModelFiles, ModelConfigError> {
206    let dir = dir.as_ref();
207    let mut files = ModelFiles::default();
208    for (name, slot) in [
209        (EDGEFIRST_JSON, &mut files.edgefirst_json),
210        (LABELS_TXT, &mut files.labels_txt),
211        (HF_CONFIG, &mut files.hf_config),
212        (HF_PREPROCESSOR, &mut files.hf_preprocessor),
213    ] {
214        let path = dir.join(name);
215        match read_capped(&path) {
216            Ok(text) => {
217                *slot = Some(ConfigFile {
218                    name: name.to_owned(),
219                    text,
220                })
221            }
222            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
223            Err(source) => return Err(ModelConfigError::Io { path, source }),
224        }
225    }
226    Ok(files)
227}
228
229fn read_capped(path: &Path) -> std::io::Result<String> {
230    let mut text = String::new();
231    File::open(path)?
232        .take(MAX_SIDECAR + 1)
233        .read_to_string(&mut text)?;
234    if text.len() as u64 > MAX_SIDECAR {
235        return Err(std::io::Error::new(
236            std::io::ErrorKind::FileTooLarge,
237            "larger than 16 MiB",
238        ));
239    }
240    Ok(text)
241}
242
243/// Non-empty trimmed lines of a `labels.txt`. Blank lines are skipped, so a
244/// `labels.txt` cannot name an empty class.
245pub(crate) fn parse_labels(text: &str) -> Vec<String> {
246    text.lines()
247        .map(str::trim)
248        .filter(|l| !l.is_empty())
249        .map(str::to_owned)
250        .collect()
251}
252
253/// Class names with leading and trailing whitespace removed; `file` and
254/// `field` name the source in the error for a name that is empty after
255/// trimming.
256pub(crate) fn trim_labels(
257    file: &str,
258    field: &str,
259    labels: Vec<String>,
260) -> Result<Vec<String>, ModelConfigError> {
261    let mut trimmed = 0;
262    let labels = labels
263        .into_iter()
264        .enumerate()
265        .map(|(i, name)| match name.trim() {
266            "" => Err(invalid(
267                file,
268                &format!("{field}[{i}]"),
269                "is empty after trimming whitespace",
270            )),
271            t if t.len() == name.len() => Ok(name),
272            t => {
273                trimmed += 1;
274                Ok(t.to_owned())
275            }
276        })
277        .collect::<Result<Vec<_>, _>>()?;
278    if trimmed > 0 {
279        log::debug!("{file} {field}: trimmed whitespace from {trimmed} class names");
280    }
281    Ok(labels)
282}
283
284/// The class names in `parsed`, or `None` with a warning in `notes` when they
285/// are unusable.
286pub(crate) fn usable_labels(
287    parsed: Result<Vec<String>, ModelConfigError>,
288    notes: &mut Vec<String>,
289) -> Option<Vec<String>> {
290    parsed
291        .map_err(|e| {
292            let note = format!("{e}; its class names are ignored");
293            log::warn!("{note}");
294            notes.push(note);
295        })
296        .ok()
297}
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302
303    #[test]
304    fn onnx_both_key_sets_map_to_fields() {
305        let bytes = super::super::onnx::tests::model_proto(
306            &[
307                ("edgefirst.json", "{\"legacy\":1}"),
308                ("labels", "cat\ndog"),
309                ("hf_config", "{}"),
310                ("hf_preprocessor_config", "{\"size\":1}"),
311                ("names", "{0: 'cat'}"),
312            ],
313            32,
314        );
315        let f = read_embedded_bytes(&bytes, true).unwrap();
316        assert_eq!(f.edgefirst_json.unwrap().name, "edgefirst.json");
317        assert_eq!(parse_labels(&f.labels_txt.unwrap().text), ["cat", "dog"]);
318        assert_eq!(f.hf_config.unwrap().name, "hf_config");
319        assert_eq!(f.hf_preprocessor.unwrap().text, "{\"size\":1}");
320        assert_eq!(f.metadata["names"], "{0: 'cat'}");
321
322        let bytes = super::super::onnx::tests::model_proto(&[("edgefirst", "{}")], 0);
323        let f = read_embedded_bytes(&bytes, true).unwrap();
324        assert_eq!(f.edgefirst_json.unwrap().name, "edgefirst");
325    }
326
327    #[test]
328    fn malformed_onnx_is_an_error_only_when_named_onnx() {
329        let bytes = [0x08, 0x08, 0x3a, 0x7f];
330        assert!(read_embedded_bytes(&bytes, true).is_err());
331        assert!(read_embedded_bytes(&bytes, false).unwrap().is_empty());
332    }
333
334    #[cfg(feature = "model-files")]
335    #[test]
336    fn tflite_associated_files_read() {
337        let mut payload = vec![0u8; 4];
338        payload.extend_from_slice(b"TFL3");
339        payload.resize(512, 0x5a);
340        let bytes = super::super::ziptrailer::tests::with_trailer(
341            &payload,
342            &[
343                ("config.json", "{\"model_type\":\"segformer\"}"),
344                ("preprocessor_config.json", "{}"),
345                ("metadata.json", "{\"names\":{}}"),
346            ],
347        );
348        let f = read_embedded_bytes(&bytes, false).unwrap();
349        assert_eq!(f.payload_len, Some(512));
350        assert_eq!(f.hf_config.unwrap().name, "config.json");
351        assert!(f.hf_preprocessor.is_some());
352        assert!(f.metadata.contains_key("metadata.json"));
353    }
354
355    #[cfg(feature = "model-files")]
356    #[test]
357    fn onnx_with_appended_zip_reads_both() {
358        let model = super::super::onnx::tests::model_proto(&[("labels", "a\nb")], 64);
359        let bytes =
360            super::super::ziptrailer::tests::with_trailer(&model, &[("edgefirst.json", "{}")]);
361        let f = read_embedded_bytes(&bytes, true).unwrap();
362        assert_eq!(f.payload_len, Some(model.len() as u64));
363        assert!(f.edgefirst_json.is_some());
364        assert_eq!(f.labels_txt.unwrap().name, "labels");
365    }
366
367    #[test]
368    fn a_directory_is_an_io_error() {
369        let e = read_embedded(std::env::temp_dir()).unwrap_err();
370        assert!(
371            matches!(e, ModelConfigError::Io { ref source, .. } if source.kind() == std::io::ErrorKind::IsADirectory),
372            "{e}"
373        );
374    }
375
376    #[test]
377    fn sidecars_read_from_directory() {
378        let dir = std::env::temp_dir().join(format!("ef-sidecar-{}", std::process::id()));
379        std::fs::create_dir_all(&dir).unwrap();
380        std::fs::write(dir.join("edgefirst.json"), "{}").unwrap();
381        std::fs::write(dir.join("labels.txt"), " a \n\nb\n").unwrap();
382        let f = read_sidecars(&dir).unwrap();
383        std::fs::remove_dir_all(&dir).unwrap();
384        assert_eq!(parse_labels(&f.labels_txt.unwrap().text), ["a", "b"]);
385        assert!(f.edgefirst_json.is_some());
386        assert!(f.hf_config.is_none());
387    }
388
389    #[test]
390    fn labels_trim_outer_whitespace_only() {
391        let names = ["bed ", "\u{3000}wall", "chest of drawers", "\tsky\n"];
392        let got = trim_labels("f", "id2label", names.map(String::from).to_vec()).unwrap();
393        assert_eq!(got, ["bed", "wall", "chest of drawers", "sky"]);
394    }
395
396    #[test]
397    fn label_empty_after_trim_names_source_and_index() {
398        let names = vec!["a".to_owned(), " \u{a0}".to_owned()];
399        let e = trim_labels("config.json", "id2label", names).unwrap_err();
400        assert!(
401            matches!(&e, ModelConfigError::InvalidField { file, field, .. }
402                if file == "config.json" && field == "id2label[1]"),
403            "{e}"
404        );
405    }
406}
edgefirst_decoder/model_config/container.rs

edgefirst_decoder/model_config/
container.rs
