ultralytics_inference/
metadata.rs

1// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
2
3//! ONNX model metadata parsing.
4//!
5//! This module handles parsing metadata from Ultralytics YOLO ONNX models.
6//! The metadata is stored as YAML in the ONNX model's custom metadata properties.
7
8use std::collections::HashMap;
9use std::sync::Arc;
10
11use crate::error::{InferenceError, Result};
12use crate::task::Task;
13
14/// Metadata extracted from an Ultralytics YOLO ONNX model.
15///
16/// This struct contains all the configuration information embedded in the model,
17/// including class names, input dimensions, and task type.
18#[derive(Debug, Clone)]
19pub struct ModelMetadata {
20    /// Model description (e.g., "Ultralytics `YOLO11n` model trained on coco.yaml").
21    pub description: String,
22    /// Model author.
23    pub author: String,
24    /// Export date.
25    pub date: String,
26    /// Ultralytics version used for export.
27    pub version: String,
28    /// License information.
29    pub license: String,
30    /// Documentation URL.
31    pub docs: String,
32    /// The task this model performs.
33    pub task: Task,
34    /// Model stride (typically 32 for YOLO).
35    pub stride: u32,
36    /// Batch size the model was exported with.
37    pub batch: usize,
38    /// Input image size as (height, width).
39    pub imgsz: Option<(usize, usize)>,
40    /// Number of input channels (typically 3 for RGB).
41    pub channels: usize,
42    /// Whether the model uses FP16 (half precision).
43    pub half: bool,
44    /// Class ID to class name mapping.
45    pub names: Arc<HashMap<usize, String>>,
46    /// Whether the model was exported with end-to-end NMS-free output
47    /// (YOLO26-style post-NMS output: `[B, max_det, 6+extra]`).
48    pub end2end: bool,
49    /// Pose keypoint shape as (`num_keypoints`, `dims`), e.g. (17, 3).
50    pub kpt_shape: Option<(usize, usize)>,
51}
52
53impl ModelMetadata {
54    /// Parse metadata from ONNX model custom metadata properties.
55    ///
56    /// # Arguments
57    ///
58    /// * `metadata_map` - The custom metadata from the ONNX model session.
59    ///
60    /// # Returns
61    ///
62    /// * A new `ModelMetadata` instance.
63    ///
64    /// # Errors
65    ///
66    /// Returns an error if the metadata is missing or malformed.
67    pub fn from_onnx_metadata(metadata_map: &HashMap<String, String>) -> Result<Self> {
68        // The metadata is typically stored under a single key containing YAML
69        // Try common key names used by Ultralytics
70        let yaml_str = metadata_map
71            .get("metadata")
72            .or_else(|| metadata_map.get("model_metadata"))
73            .or_else(|| {
74                // If no standard key, check if all metadata is in one value
75                metadata_map.values().find(|v| v.contains("task:"))
76            })
77            .ok_or_else(|| {
78                InferenceError::ModelLoadError(
79                    "No metadata found in ONNX model. Ensure the model was exported with Ultralytics.".to_string()
80                )
81            })?;
82
83        Self::from_yaml_str(yaml_str)
84    }
85
86    /// Parse metadata from a YAML string.
87    ///
88    /// # Arguments
89    ///
90    /// * `yaml_str` - The YAML-formatted metadata string.
91    ///
92    /// # Returns
93    ///
94    /// * A new `ModelMetadata` instance.
95    ///
96    /// # Errors
97    ///
98    /// Returns an error if the YAML is malformed or missing required fields.
99    pub fn from_yaml_str(yaml_str: &str) -> Result<Self> {
100        let mut metadata = Self::default();
101        let mut names: HashMap<usize, String> = HashMap::new();
102
103        for line in yaml_str.lines() {
104            let line = line.trim();
105            if line.is_empty() || line.starts_with('#') {
106                continue;
107            }
108
109            // Handle key: value pairs
110            if let Some((key, value)) = line.split_once(':') {
111                let key = key.trim();
112                let value = value.trim().trim_matches('\'').trim_matches('"');
113
114                match key {
115                    "description" => metadata.description = value.to_string(),
116                    "author" => metadata.author = value.to_string(),
117                    "date" => metadata.date = value.to_string(),
118                    "version" => metadata.version = value.to_string(),
119                    "license" => metadata.license = value.to_string(),
120                    "docs" => metadata.docs = value.to_string(),
121                    "task" => {
122                        metadata.task = value.parse().map_err(|e| {
123                            InferenceError::ModelLoadError(format!("Invalid task in metadata: {e}"))
124                        })?;
125                    }
126                    "stride" => {
127                        metadata.stride = value.parse().map_err(|_| {
128                            InferenceError::ModelLoadError(format!("Invalid stride value: {value}"))
129                        })?;
130                    }
131                    "batch" => {
132                        metadata.batch = value.parse().map_err(|_| {
133                            InferenceError::ModelLoadError(format!("Invalid batch value: {value}"))
134                        })?;
135                    }
136                    "channels" => {
137                        metadata.channels = value.parse().map_err(|_| {
138                            InferenceError::ModelLoadError(format!(
139                                "Invalid channels value: {value}"
140                            ))
141                        })?;
142                    }
143                    "half" => {
144                        metadata.half = value == "true" || value == "True";
145                    }
146                    "end2end" => {
147                        metadata.end2end = value == "true" || value == "True";
148                    }
149                    "args" => {
150                        // Parse args dict for half flag: {'half': True, ...}
151                        if value.contains("'half': True")
152                            || value.contains("\"half\": true")
153                            || value.contains("'half':True")
154                        {
155                            metadata.half = true;
156                        }
157                    }
158                    _ => {
159                        // Check for class name entries (numeric keys)
160                        if let Ok(class_id) = key.trim().parse::<usize>() {
161                            names.insert(class_id, value.to_string());
162                        }
163                    }
164                }
165            }
166        }
167
168        // imgsz is a two-integer list (`[640, 640]` inline or a `- 640` block),
169        // parsed with the same helper as kpt_shape.
170        metadata.imgsz = Self::parse_int_pair(yaml_str, "imgsz");
171
172        // kpt_shape is a two-integer list (inline `[17, 3]` or a `- 17` block),
173        // parsed with the same helper as imgsz.
174        metadata.kpt_shape = Self::parse_int_pair(yaml_str, "kpt_shape");
175
176        // Parse names block if not already parsed inline
177        if names.is_empty() {
178            names = Self::parse_names_block(yaml_str);
179        }
180
181        metadata.names = Arc::new(names);
182        Ok(metadata)
183    }
184
185    /// Parse a two-integer YAML value for `key`, accepting either the inline form
186    /// (`key: [a, b]` / `key: a, b`) or a multi-line block list (`key:` then
187    /// `- a` / `- b` on following lines). Returns the first two integers.
188    fn parse_int_pair(yaml_str: &str, key: &str) -> Option<(usize, usize)> {
189        let prefix = format!("{key}:");
190        let lines: Vec<&str> = yaml_str.lines().collect();
191        for (i, line) in lines.iter().enumerate() {
192            let Some(rest) = line.trim_start().strip_prefix(&prefix) else {
193                continue;
194            };
195            // Inline form: [a, b] / (a, b) / a, b
196            let rest = rest
197                .trim()
198                .trim_matches(|c| matches!(c, '[' | ']' | '(' | ')'));
199            let inline: Vec<usize> = rest
200                .split(',')
201                .filter_map(|s| s.trim().parse().ok())
202                .collect();
203            if inline.len() >= 2 {
204                return Some((inline[0], inline[1]));
205            }
206            // Multi-line block list on following `- value` lines.
207            let mut vals = Vec::new();
208            for following in lines.iter().skip(i + 1) {
209                let t = following.trim();
210                if let Some(v) = t.strip_prefix('-') {
211                    if let Ok(n) = v.trim().parse::<usize>() {
212                        vals.push(n);
213                    }
214                } else if !t.is_empty() && !t.starts_with('#') {
215                    break;
216                }
217                if vals.len() >= 2 {
218                    break;
219                }
220            }
221            if vals.len() >= 2 {
222                return Some((vals[0], vals[1]));
223            }
224        }
225        None
226    }
227
228    /// Parse the names block from YAML or Python dict format.
229    fn parse_names_block(yaml_str: &str) -> HashMap<usize, String> {
230        let mut names = HashMap::new();
231
232        // First, try to find `names: {0: 'person', 1: 'bicycle', ...}` Python dict format
233        // This is how Ultralytics stores names in ONNX metadata
234        if let Some(start) = yaml_str.find("names:") {
235            let after_names = &yaml_str[start + 6..];
236            let trimmed = after_names.trim();
237
238            // Check if it's Python dict format (starts with {)
239            if trimmed.starts_with('{')
240                && let Some(end) = trimmed.find('}')
241            {
242                let dict_str = &trimmed[1..end];
243                return Self::parse_python_dict(dict_str);
244            }
245        }
246
247        // Fall back to YAML block format
248        let lines: Vec<&str> = yaml_str.lines().collect();
249        let mut in_names_block = false;
250        let mut names_indent = 0;
251
252        for line in &lines {
253            let trimmed = line.trim();
254
255            if trimmed.starts_with("names:") {
256                in_names_block = true;
257                names_indent = line.len() - line.trim_start().len();
258                continue;
259            }
260
261            if in_names_block {
262                let current_indent = line.len() - line.trim_start().len();
263
264                // Check if we've exited the names block
265                if !trimmed.is_empty()
266                    && !trimmed.starts_with('#')
267                    && current_indent <= names_indent
268                {
269                    // Only exit if this isn't a class entry
270                    if !trimmed.chars().next().is_some_and(|c| c.is_ascii_digit()) {
271                        break;
272                    }
273                }
274
275                // Parse class entries like "0: person" or "  0: person"
276                if let Some((key, value)) = trimmed.split_once(':')
277                    && let Ok(class_id) = key.trim().parse::<usize>()
278                {
279                    let class_name = value.trim().trim_matches('\'').trim_matches('"');
280                    names.insert(class_id, class_name.to_string());
281                }
282            }
283        }
284
285        names
286    }
287
288    /// Parse a Python dict string like `0: 'person', 1: 'bicycle'`.
289    fn parse_python_dict(dict_str: &str) -> HashMap<usize, String> {
290        let mut names = HashMap::new();
291
292        // Split by comma, but be careful with quotes
293        for entry in dict_str.split(',') {
294            let entry = entry.trim();
295            if let Some((key, value)) = entry.split_once(':') {
296                let key = key.trim();
297                let value = value.trim().trim_matches('\'').trim_matches('"');
298                if let Ok(class_id) = key.parse::<usize>() {
299                    names.insert(class_id, value.to_string());
300                }
301            }
302        }
303
304        names
305    }
306
307    /// Get the number of classes in this model.
308    ///
309    /// # Returns
310    ///
311    /// * The count of classes.
312    #[must_use]
313    pub fn num_classes(&self) -> usize {
314        self.names.len()
315    }
316
317    /// Get a class name by ID.
318    ///
319    /// # Arguments
320    ///
321    /// * `class_id` - The numeric identifier for the class.
322    ///
323    /// # Returns
324    ///
325    /// * `Some` class name if found, otherwise `None`.
326    #[must_use]
327    pub fn class_name(&self, class_id: usize) -> Option<&str> {
328        self.names.get(&class_id).map(String::as_str)
329    }
330
331    /// Extract the model name from the description.
332    ///
333    /// E.g. "Ultralytics `YOLO11n` model..." -> "`YOLO11n`"
334    /// Returns `YOLO` if extraction fails.
335    #[must_use]
336    pub fn model_name(&self) -> String {
337        // Description format: "Ultralytics <MODEL> model..."
338        self.description
339            .split_whitespace()
340            .find(|&word| word.to_lowercase().starts_with("yolo"))
341            .unwrap_or("YOLO")
342            .to_string()
343    }
344}
345
346impl Default for ModelMetadata {
347    fn default() -> Self {
348        Self {
349            description: String::new(),
350            author: "Ultralytics".to_string(),
351            date: String::new(),
352            version: String::new(),
353            license: "AGPL-3.0".to_string(),
354            docs: "https://docs.ultralytics.com".to_string(),
355            task: Task::Detect,
356            stride: 32,
357            batch: 1,
358            imgsz: None,
359            channels: 3,
360            half: false,
361            names: Arc::new(HashMap::new()),
362            end2end: false,
363            kpt_shape: None,
364        }
365    }
366}
367
368#[cfg(test)]
369mod tests {
370    use super::*;
371
372    const SAMPLE_METADATA: &str = r"
373description: Ultralytics YOLO11n model trained on /usr/src/ultralytics/ultralytics/cfg/datasets/coco.yaml
374author: Ultralytics
375date: '2025-12-11T20:19:45.464021'
376version: 8.3.236
377license: AGPL-3.0 License (https://ultralytics.com/license)
378docs: https://docs.ultralytics.com
379stride: 32
380task: detect
381batch: 1
382imgsz:
383- 640
384- 640
385names:
386  0: person
387  1: bicycle
388  2: car
389  3: motorcycle
390channels: 3
391";
392
393    #[test]
394    fn test_parse_metadata() {
395        let metadata = ModelMetadata::from_yaml_str(SAMPLE_METADATA).unwrap();
396
397        assert_eq!(metadata.task, Task::Detect);
398        assert_eq!(metadata.stride, 32);
399        assert_eq!(metadata.batch, 1);
400        assert_eq!(metadata.imgsz, Some((640, 640)));
401        assert_eq!(metadata.channels, 3);
402        assert_eq!(metadata.num_classes(), 4);
403        assert_eq!(metadata.class_name(0), Some("person"));
404        assert_eq!(metadata.class_name(1), Some("bicycle"));
405        assert_eq!(metadata.class_name(2), Some("car"));
406        assert_eq!(metadata.class_name(3), Some("motorcycle"));
407    }
408
409    #[test]
410    fn test_parse_inline_imgsz() {
411        let yaml = "task: detect\nimgsz: [640, 640]\nstride: 32";
412        let metadata = ModelMetadata::from_yaml_str(yaml).unwrap();
413        assert_eq!(metadata.imgsz, Some((640, 640)));
414    }
415
416    #[test]
417    fn test_default_metadata() {
418        let metadata = ModelMetadata::default();
419        assert_eq!(metadata.task, Task::Detect);
420        assert_eq!(metadata.stride, 32);
421        assert_eq!(metadata.imgsz, None);
422    }
423
424    #[test]
425    fn test_parse_multiline_kpt_shape_and_imgsz() {
426        // Ultralytics `.tflite` metadata.yaml writes imgsz and kpt_shape as
427        // multi-line YAML block lists rather than inline `[a, b]`.
428        let yaml = "task: pose\nstride: 32\nimgsz:\n- 640\n- 640\nkpt_shape:\n- 17\n- 3\n";
429        let metadata = ModelMetadata::from_yaml_str(yaml).unwrap();
430        assert_eq!(metadata.task, Task::Pose);
431        assert_eq!(metadata.imgsz, Some((640, 640)));
432        assert_eq!(metadata.kpt_shape, Some((17, 3)));
433    }
434
435    #[test]
436    fn test_parse_inline_kpt_shape() {
437        let yaml = "task: pose\nkpt_shape: [17, 3]\nstride: 32";
438        let metadata = ModelMetadata::from_yaml_str(yaml).unwrap();
439        assert_eq!(metadata.kpt_shape, Some((17, 3)));
440    }
441
442    #[test]
443    fn test_parse_int_pair_inline_and_block() {
444        // Inline forms: `[a, b]`, `(a, b)`, `a, b`.
445        assert_eq!(
446            ModelMetadata::parse_int_pair("kpt_shape: [17, 3]", "kpt_shape"),
447            Some((17, 3))
448        );
449        assert_eq!(
450            ModelMetadata::parse_int_pair("kpt_shape: (17, 2)", "kpt_shape"),
451            Some((17, 2))
452        );
453        assert_eq!(
454            ModelMetadata::parse_int_pair("foo: 5, 6", "foo"),
455            Some((5, 6))
456        );
457        // Multi-line block list.
458        assert_eq!(
459            ModelMetadata::parse_int_pair("kpt_shape:\n- 17\n- 3", "kpt_shape"),
460            Some((17, 3))
461        );
462        // A single value, junk, or wrong key yields None.
463        assert_eq!(
464            ModelMetadata::parse_int_pair("kpt_shape: [17]", "kpt_shape"),
465            None
466        );
467        assert_eq!(ModelMetadata::parse_int_pair("foo: 5", "foo"), None);
468        assert_eq!(ModelMetadata::parse_int_pair("other: 1, 2", "foo"), None);
469    }
470
471    #[test]
472    fn test_python_dict_names() {
473        let yaml = "task: detect\nnames: {0: 'person', 1: 'bicycle', 2: 'car'}";
474        let m = ModelMetadata::from_yaml_str(yaml).unwrap();
475        assert_eq!(m.num_classes(), 3);
476        assert_eq!(m.class_name(1), Some("bicycle"));
477    }
478
479    #[test]
480    fn test_tflite_style_metadata_text() {
481        // The `.tflite` reader rebuilds this text (unquoted `names:` block, inline
482        // lists) and hands it to `from_yaml_str`, exactly like the ONNX path.
483        let text = "task: detect\nimgsz: [640, 640]\nstride: 32\nend2end: false\nnames:\n  0: person\n  1: traffic light\n  2: men's shoe";
484        let m = ModelMetadata::from_yaml_str(text).unwrap();
485        assert_eq!(m.task, Task::Detect);
486        assert_eq!(m.imgsz, Some((640, 640)));
487        assert_eq!(m.num_classes(), 3);
488        // Names with a space and an apostrophe must survive (no escaping needed).
489        assert_eq!(m.class_name(1), Some("traffic light"));
490        assert_eq!(m.class_name(2), Some("men's shoe"));
491    }
492
493    #[test]
494    fn test_names_block_with_yaml_significant_chars() {
495        // The line parser splits on the first `:` and only treats a line as a
496        // comment when it *starts* with `#`, so class names containing `:` or `#`
497        // round-trip without escaping.
498        let text = "task: detect\nnames:\n  0: ratio 16:9\n  1: #1 pick";
499        let m = ModelMetadata::from_yaml_str(text).unwrap();
500        assert_eq!(m.class_name(0), Some("ratio 16:9"));
501        assert_eq!(m.class_name(1), Some("#1 pick"));
502    }
503
504    #[test]
505    fn test_pose_kpt_shape_and_half_flags() {
506        let yaml = "task: pose\nkpt_shape: [17, 3]\nhalf: true\nend2end: True";
507        let m = ModelMetadata::from_yaml_str(yaml).unwrap();
508        assert_eq!(m.task, Task::Pose);
509        assert_eq!(m.kpt_shape, Some((17, 3)));
510        assert!(m.half);
511        assert!(m.end2end);
512
513        // half can also arrive embedded in an args dict.
514        let yaml2 = "task: detect\nargs: {'half': True, 'imgsz': 640}";
515        let m2 = ModelMetadata::from_yaml_str(yaml2).unwrap();
516        assert!(m2.half);
517    }
518
519    #[test]
520    fn test_invalid_fields_error() {
521        assert!(ModelMetadata::from_yaml_str("task: notarealtask").is_err());
522        assert!(ModelMetadata::from_yaml_str("task: detect\nstride: abc").is_err());
523        assert!(ModelMetadata::from_yaml_str("task: detect\nbatch: xyz").is_err());
524    }
525
526    #[test]
527    fn test_from_onnx_metadata_keys_and_error() {
528        // Standard "metadata" key.
529        let map = HashMap::from([("metadata".to_string(), SAMPLE_METADATA.to_string())]);
530        let m = ModelMetadata::from_onnx_metadata(&map).unwrap();
531        assert_eq!(m.task, Task::Detect);
532
533        // Fallback: any value containing "task:".
534        let map = HashMap::from([("whatever".to_string(), "task: pose\nstride: 32".to_string())]);
535        let m = ModelMetadata::from_onnx_metadata(&map).unwrap();
536        assert_eq!(m.task, Task::Pose);
537
538        // No usable metadata -> error.
539        let map = HashMap::from([("unrelated".to_string(), "no yaml here".to_string())]);
540        assert!(ModelMetadata::from_onnx_metadata(&map).is_err());
541    }
542
543    #[test]
544    fn test_model_name_extraction() {
545        let m = ModelMetadata::from_yaml_str(SAMPLE_METADATA).unwrap();
546        assert_eq!(m.model_name(), "YOLO11n");
547        // Falls back to "YOLO" when the description has no yolo token.
548        let plain = ModelMetadata::default();
549        assert_eq!(plain.model_name(), "YOLO");
550    }
551}
ultralytics_inference/metadata.rs

ultralytics_inference/
metadata.rs
