1use std::collections::{BTreeMap, BTreeSet};
13use std::fmt;
14
15use crate::configs::DimName;
16use crate::schema::{
17 BoxEncoding, DType, DecoderKind, DecoderVersion, InputSpec, LogicalOutput, LogicalType,
18 NmsMode, Quantization, SchemaV2, ScoreFormat,
19};
20
21#[derive(Debug, Clone, PartialEq)]
23pub struct TensorInfo {
24 pub name: String,
26 pub shape: Vec<usize>,
28 pub dtype: DType,
30 pub quantization: Option<Quantization>,
32}
33
34edgefirst_tensor::ef_vocabulary! {
35 pub enum ModelSource {
41 Onnx = 0, "onnx", ONNX,
44 TfLite = 1, "tflite", TFLITE,
47 CoreMl = 3, "coreml", COREML,
67 Other = 2, "other", OTHER,
75 }
76 #[doc(hidden)]
77 pub mod model_source_code;
78}
79
80#[derive(Debug, Clone, PartialEq)]
83pub struct ModelSignals {
84 pub source: ModelSource,
86 pub inputs: Vec<TensorInfo>,
88 pub outputs: Vec<TensorInfo>,
90 pub metadata: BTreeMap<String, String>,
94}
95
96#[derive(Debug, Clone, PartialEq)]
98pub struct InferredSchema {
99 pub schema: SchemaV2,
101 pub labels: Vec<String>,
103 pub description: String,
105}
106
107#[derive(Debug)]
109pub enum InferError {
110 NotUltralytics(String),
112 BadNames(String),
114 BadMetadata(String),
118 UnsupportedTask(String),
120 ClassCountMismatch { expected: usize, found: usize },
123 UnsupportedLayout(String),
126 AmbiguousLayout(String),
130 UnsupportedQuantization(String),
134 UnknownBoxConvention,
140}
141
142impl fmt::Display for InferError {
143 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
144 match self {
145 InferError::NotUltralytics(s) => {
146 write!(f, "model metadata carries no Ultralytics signature: {s}")
147 }
148 InferError::BadNames(s) => {
149 write!(f, "cannot parse Ultralytics `names` metadata: {s}")
150 }
151 InferError::BadMetadata(s) => {
152 write!(f, "malformed Ultralytics metadata: {s}")
153 }
154 InferError::UnsupportedTask(s) => write!(
155 f,
156 "unsupported Ultralytics task `{s}` (supported: detect, segment)"
157 ),
158 InferError::ClassCountMismatch { expected, found } => write!(
159 f,
160 "class-count mismatch: names give {expected} classes but output features imply {found}"
161 ),
162 InferError::UnsupportedLayout(s) => write!(f, "unsupported output layout: {s}"),
163 InferError::AmbiguousLayout(s) => write!(f, "ambiguous output layout: {s}"),
164 InferError::UnsupportedQuantization(s) => {
165 write!(f, "unsupported quantization: {s}")
166 }
167 InferError::UnknownBoxConvention => write!(
168 f,
169 "box coordinate convention is unknown for this container; only \
170 onnx (pixel-space), tflite ([0,1]), and coreml (pixel-space) \
171 are characterized"
172 ),
173 }
174 }
175}
176
177impl std::error::Error for InferError {}
178
179#[derive(Debug, Clone, Copy, PartialEq, Eq)]
181enum Task {
182 Detect,
183 Segment,
184}
185
186impl Task {
187 fn parse(s: &str) -> Result<Self, InferError> {
188 match s.trim().to_ascii_lowercase().as_str() {
189 "detect" => Ok(Task::Detect),
190 "segment" => Ok(Task::Segment),
191 other => Err(InferError::UnsupportedTask(other.to_string())),
192 }
193 }
194}
195
196#[derive(Debug, Clone, PartialEq)]
199struct UltralyticsMeta {
200 names: Vec<String>,
201 task: Option<Task>,
202 end2end: Option<bool>,
206}
207
208fn reject_foreign_vendor(author: Option<&String>, docs: Option<&String>) -> Result<(), InferError> {
224 let names_ultralytics = |v: &String| v.to_ascii_lowercase().contains("ultralytics");
225 for (field, value) in [("author", author), ("docs", docs)] {
226 if let Some(v) = value {
227 if !v.trim().is_empty() && !names_ultralytics(v) {
228 return Err(InferError::NotUltralytics(format!(
229 "metadata `{field}` is `{v}`, not Ultralytics"
230 )));
231 }
232 }
233 }
234 Ok(())
235}
236
237impl UltralyticsMeta {
238 fn from_metadata(meta: &BTreeMap<String, String>) -> Result<Self, InferError> {
244 if let Some(names_raw) = meta.get("names") {
245 reject_foreign_vendor(meta.get("author"), meta.get("docs"))?;
246 return Self::from_flat_props(meta, names_raw);
247 }
248
249 for value in meta.values() {
252 if let Ok(serde_json::Value::Object(obj)) = serde_json::from_str(value) {
253 if let Some(names) = obj.get("names") {
254 let author = obj
255 .get("author")
256 .and_then(|v| v.as_str())
257 .map(str::to_string);
258 let docs = obj.get("docs").and_then(|v| v.as_str()).map(str::to_string);
259 reject_foreign_vendor(author.as_ref(), docs.as_ref())?;
260 return Self::from_json_envelope(&obj, names);
261 }
262 }
263 }
264
265 Err(InferError::NotUltralytics(
266 "no `names` metadata field or metadata.json envelope found".into(),
267 ))
268 }
269
270 fn from_flat_props(
271 meta: &BTreeMap<String, String>,
272 names_raw: &str,
273 ) -> Result<Self, InferError> {
274 let names = parse_names(names_raw)?;
275 let task = meta.get("task").map(|t| Task::parse(t)).transpose()?;
276 let end2end = meta.get("end2end").map(|s| parse_bool_str(s)).transpose()?;
277 Ok(UltralyticsMeta {
278 names,
279 task,
280 end2end,
281 })
282 }
283
284 fn from_json_envelope(
288 obj: &serde_json::Map<String, serde_json::Value>,
289 names_value: &serde_json::Value,
290 ) -> Result<Self, InferError> {
291 let names = names_from_json_value(names_value)?;
292 let task = obj
293 .get("task")
294 .and_then(|v| v.as_str())
295 .map(Task::parse)
296 .transpose()?;
297 let end2end = match obj.get("end2end") {
300 None | Some(serde_json::Value::Null) => None,
301 Some(serde_json::Value::Bool(b)) => Some(*b),
302 Some(other) => {
303 return Err(InferError::BadMetadata(format!(
304 "`end2end` is {other}, expected a boolean"
305 )))
306 }
307 };
308 Ok(UltralyticsMeta {
309 names,
310 task,
311 end2end,
312 })
313 }
314}
315
316fn parse_bool_str(s: &str) -> Result<bool, InferError> {
324 match s.trim().to_ascii_lowercase().as_str() {
325 "true" => Ok(true),
326 "false" => Ok(false),
327 other => Err(InferError::BadMetadata(format!(
328 "`end2end` is `{other}`, expected true or false"
329 ))),
330 }
331}
332
333fn parse_names(s: &str) -> Result<Vec<String>, InferError> {
338 let trimmed = s.trim();
339
340 if let Ok(map) = serde_json::from_str::<BTreeMap<String, String>>(trimmed) {
341 return names_from_str_keyed_map(map);
342 }
343
344 let body = trimmed
345 .strip_prefix('{')
346 .and_then(|s| s.strip_suffix('}'))
347 .ok_or_else(|| InferError::BadNames(format!("not a dict: {trimmed}")))?;
348
349 let mut indexed: BTreeMap<usize, String> = BTreeMap::new();
350 for item in split_top_level_commas(body) {
351 let item = item.trim();
352 if item.is_empty() {
353 continue;
354 }
355 let (key, value) = item
356 .split_once(':')
357 .ok_or_else(|| InferError::BadNames(format!("malformed entry `{item}`")))?;
358 let idx: usize = key
359 .trim()
360 .parse()
361 .map_err(|_| InferError::BadNames(format!("non-numeric key `{key}`")))?;
362 let value = unquote(value.trim())
363 .ok_or_else(|| InferError::BadNames(format!("unquoted value `{value}`")))?;
364 indexed.insert(idx, value);
365 }
366 contiguous_values(indexed)
367}
368
369fn names_from_json_value(v: &serde_json::Value) -> Result<Vec<String>, InferError> {
372 let obj = v
373 .as_object()
374 .ok_or_else(|| InferError::BadNames("`names` is not a JSON object".into()))?;
375 let mut map = BTreeMap::new();
376 for (k, val) in obj {
377 let s = val
378 .as_str()
379 .ok_or_else(|| InferError::BadNames(format!("non-string name for key `{k}`")))?;
380 map.insert(k.clone(), s.to_string());
381 }
382 names_from_str_keyed_map(map)
383}
384
385fn names_from_str_keyed_map(map: BTreeMap<String, String>) -> Result<Vec<String>, InferError> {
389 let mut indexed: BTreeMap<usize, String> = BTreeMap::new();
390 for (k, v) in map {
391 let idx: usize = k
392 .parse()
393 .map_err(|_| InferError::BadNames(format!("non-numeric key `{k}`")))?;
394 indexed.insert(idx, v);
395 }
396 contiguous_values(indexed)
397}
398
399fn contiguous_values(indexed: BTreeMap<usize, String>) -> Result<Vec<String>, InferError> {
402 if indexed.is_empty() {
407 return Err(InferError::BadNames("class map is empty".into()));
408 }
409 let n = indexed.len();
410 for i in 0..n {
411 if !indexed.contains_key(&i) {
412 return Err(InferError::BadNames("non-contiguous class indices".into()));
413 }
414 }
415 Ok(indexed.into_values().collect())
416}
417
418fn split_top_level_commas(s: &str) -> Vec<String> {
421 let mut parts = Vec::new();
422 let mut current = String::new();
423 let mut quote: Option<char> = None;
424 for c in s.chars() {
425 match quote {
426 Some(q) => {
427 current.push(c);
428 if c == q {
429 quote = None;
430 }
431 }
432 None => match c {
433 '\'' | '"' => {
434 quote = Some(c);
435 current.push(c);
436 }
437 ',' => parts.push(std::mem::take(&mut current)),
438 _ => current.push(c),
439 },
440 }
441 }
442 if !current.is_empty() {
443 parts.push(current);
444 }
445 parts
446}
447
448fn unquote(s: &str) -> Option<String> {
451 let bytes = s.as_bytes();
452 if s.len() >= 2 {
453 let first = bytes[0];
454 let last = bytes[s.len() - 1];
455 if (first == b'\'' && last == b'\'') || (first == b'"' && last == b'"') {
456 return Some(s[1..s.len() - 1].to_string());
457 }
458 }
459 None
460}
461
462#[allow(clippy::type_complexity)]
473fn classify_input_layout(
474 input: &TensorInfo,
475) -> Result<(usize, usize, Vec<(DimName, usize)>), InferError> {
476 debug_assert_eq!(
477 input.shape.len(),
478 4,
479 "caller selects the input by rank-4; see infer_ultralytics_schema"
480 );
481 if input.shape[0] != 1 {
482 return Err(InferError::UnsupportedLayout(format!(
483 "model input `{}` has batch {} but every supported output layout \
484 is batch-1; shape {:?}",
485 input.name, input.shape[0], input.shape
486 )));
487 }
488 if input.shape[1] == 3 {
489 let (h, w) = (input.shape[2], input.shape[3]);
490 Ok((
491 h,
492 w,
493 vec![
494 (DimName::Batch, input.shape[0]),
495 (DimName::NumFeatures, 3),
496 (DimName::Height, h),
497 (DimName::Width, w),
498 ],
499 ))
500 } else if input.shape[3] == 3 {
501 let (h, w) = (input.shape[1], input.shape[2]);
502 Ok((
503 h,
504 w,
505 vec![
506 (DimName::Batch, input.shape[0]),
507 (DimName::Height, h),
508 (DimName::Width, w),
509 (DimName::NumFeatures, 3),
510 ],
511 ))
512 } else {
513 Err(InferError::UnsupportedLayout(format!(
514 "cannot locate the channel dimension (=3) in input shape {:?}",
515 input.shape
516 )))
517 }
518}
519
520fn find_anchors_dim(shape: &[usize], expected_anchors: usize) -> Option<(usize, usize)> {
527 if shape.len() != 3 || shape[0] != 1 {
528 return None;
529 }
530 if shape[2] == expected_anchors {
531 Some((shape[1], 1)) } else if shape[1] == expected_anchors {
533 Some((shape[2], 2)) } else {
535 None
536 }
537}
538
539#[allow(clippy::type_complexity)]
555fn classify_pre_nms<'a>(
556 rank3_outputs: &[&'a TensorInfo],
557 expected_anchors: usize,
558 feat: usize,
559 nc: usize,
560) -> Result<(&'a TensorInfo, Vec<(DimName, usize)>), InferError> {
561 let mut mismatch: Option<InferError> = None;
562 let mut matched: Option<(&'a TensorInfo, Vec<(DimName, usize)>)> = None;
563 for out in rank3_outputs {
564 let Some((other_dim, feat_idx)) = find_anchors_dim(&out.shape, expected_anchors) else {
565 continue;
566 };
567 if other_dim == feat {
568 let dshape = if feat_idx == 1 {
569 vec![
570 (DimName::Batch, 1),
571 (DimName::NumFeatures, feat),
572 (DimName::NumBoxes, expected_anchors),
573 ]
574 } else {
575 vec![
576 (DimName::Batch, 1),
577 (DimName::NumBoxes, expected_anchors),
578 (DimName::NumFeatures, feat),
579 ]
580 };
581 if let Some((first, _)) = &matched {
584 return Err(InferError::AmbiguousLayout(format!(
585 "outputs `{}` and `{}` both match the pre-NMS detection layout \
586 ({:?}); cannot tell which carries the detections",
587 first.name, out.name, out.shape
588 )));
589 }
590 matched = Some((out, dshape));
591 continue;
592 }
593 mismatch.get_or_insert_with(|| InferError::ClassCountMismatch {
594 expected: nc,
595 found: other_dim.saturating_sub(feat.saturating_sub(nc)),
598 });
599 }
600 if let Some(found) = matched {
601 return Ok(found);
602 }
603 Err(mismatch.unwrap_or_else(|| {
604 InferError::UnsupportedLayout(format!(
605 "no output matches the Ultralytics pre-NMS layout (expected a dim == \
606 {expected_anchors} anchors and a dim == {feat} features [4 + {nc} classes]); \
607 outputs seen: {:?}",
608 rank3_outputs.iter().map(|o| &o.shape).collect::<Vec<_>>()
609 ))
610 }))
611}
612
613fn find_e2e_candidate<'a>(
628 rank3_outputs: &[&'a TensorInfo],
629 expected_anchors: usize,
630 feat: usize,
631) -> Result<Option<(&'a TensorInfo, usize)>, InferError> {
632 let matches: Vec<(&'a TensorInfo, usize)> = rank3_outputs
633 .iter()
634 .filter_map(|o| {
635 let s = &o.shape;
636 debug_assert_eq!(s.len(), 3, "caller filters to rank-3 outputs");
637 (s[0] == 1 && s[2] == feat && s[1] > 0 && s[1] <= expected_anchors)
640 .then_some((*o, s[1]))
641 })
642 .collect();
643 if matches.len() > 1 {
647 return Err(InferError::AmbiguousLayout(format!(
648 "{} outputs match the end-to-end layout [1, N, {feat}]: {:?}; \
649 cannot tell which carries the detections",
650 matches.len(),
651 matches.iter().map(|(o, _)| &o.name).collect::<Vec<_>>()
652 )));
653 }
654 Ok(matches.into_iter().next())
655}
656
657fn classify_proto(
670 proto: &TensorInfo,
671 h: usize,
672 w: usize,
673) -> Result<(usize, Vec<(DimName, usize)>), InferError> {
674 let (eh, ew) = (h / 4, w / 4);
675 let s = &proto.shape;
676 if s.len() != 4 || s[0] != 1 {
677 return Err(InferError::UnsupportedLayout(format!(
678 "expected a rank-4 proto tensor with batch 1, found shape {s:?}"
679 )));
680 }
681 if s[2] == eh && s[3] == ew {
682 let k = s[1];
683 Ok((
684 k,
685 vec![
686 (DimName::Batch, 1),
687 (DimName::NumProtos, k),
688 (DimName::Height, eh),
689 (DimName::Width, ew),
690 ],
691 ))
692 } else if s[1] == eh && s[2] == ew {
693 let k = s[3];
694 Ok((
695 k,
696 vec![
697 (DimName::Batch, 1),
698 (DimName::Height, eh),
699 (DimName::Width, ew),
700 (DimName::NumProtos, k),
701 ],
702 ))
703 } else {
704 Err(InferError::UnsupportedLayout(format!(
705 "proto tensor shape {s:?} doesn't match the expected ({eh}, {ew}) spatial dims"
706 )))
707 }
708}
709
710fn default_logical() -> LogicalOutput {
713 LogicalOutput {
714 name: None,
715 type_: None,
716 shape: Vec::new(),
717 dshape: Vec::new(),
718 decoder: None,
719 encoding: None,
720 score_format: None,
721 normalized: None,
722 anchors: None,
723 stride: None,
724 dtype: None,
725 quantization: None,
726 outputs: Vec::new(),
727 activation_applied: None,
728 activation_required: None,
729 semantic: Default::default(),
730 }
731}
732
733fn check_quantized_boundary(tensor: &TensorInfo) -> Result<(), InferError> {
749 if tensor.dtype.is_integer() && tensor.quantization.is_none() {
750 return Err(InferError::UnsupportedLayout(format!(
751 "output `{}` has integer dtype {:?} but no quantization params",
752 tensor.name, tensor.dtype
753 )));
754 }
755 if let Some(q) = &tensor.quantization {
756 if q.scale.len() > 1 {
757 return Err(InferError::UnsupportedQuantization(format!(
758 "output `{}` carries per-channel quantization ({} scales); \
759 the decoder supports per-tensor quantization only",
760 tensor.name,
761 q.scale.len()
762 )));
763 }
764 if q.scale.is_empty() {
768 return Err(InferError::UnsupportedQuantization(format!(
769 "output `{}` is quantized but carries no scale",
770 tensor.name
771 )));
772 }
773 if let Some(zp) = &q.zero_point {
777 if zp.len() != q.scale.len() {
778 return Err(InferError::UnsupportedQuantization(format!(
779 "output `{}` has {} scale(s) but {} zero point(s)",
780 tensor.name,
781 q.scale.len(),
782 zp.len()
783 )));
784 }
785 }
786 }
787 Ok(())
788}
789
790pub fn infer_ultralytics_schema(signals: &ModelSignals) -> Result<InferredSchema, InferError> {
867 infer_from_meta(signals, UltralyticsMeta::from_metadata(&signals.metadata)?)
868}
869
870pub(crate) fn infer_ultralytics_schema_ignoring_names(
875 signals: &ModelSignals,
876 classes: Option<usize>,
877) -> Result<(InferredSchema, usize), InferError> {
878 let meta = UltralyticsMeta::from_metadata(&signals.metadata)?;
879 let with_count = |nc: usize| {
880 let names = (0..nc).map(|i| format!("class{i}")).collect();
881 let candidate = UltralyticsMeta {
882 names,
883 task: meta.task,
884 end2end: meta.end2end,
885 };
886 infer_from_meta(signals, candidate)
887 };
888 const REMEDY: &str = "supply a labels override naming every class, or describe the \
889 outputs in an edgefirst.json beside the model";
890 let stated = classes
891 .filter(|&nc| nc > 0)
892 .and_then(|nc| with_count(nc).ok());
893 let mut found: Option<InferredSchema> = stated;
894 if found.is_none() {
895 for nc in candidate_class_counts(signals)? {
896 let inferred = match with_count(nc) {
897 Ok(inferred) => inferred,
898 Err(InferError::AmbiguousLayout(why)) => {
899 return Err(InferError::AmbiguousLayout(format!("{why}; {REMEDY}")))
900 }
901 Err(_) => continue,
902 };
903 if let Some(first) = &found {
904 return Err(InferError::AmbiguousLayout(format!(
905 "`names` miscounts the classes and the outputs allow both {} and {nc}; \
906 {REMEDY}",
907 first.labels.len()
908 )));
909 }
910 found = Some(inferred);
911 }
912 }
913 let mut inferred = found.ok_or_else(|| {
914 InferError::UnsupportedLayout("no output carries a usable class count".into())
915 })?;
916 let nc = std::mem::take(&mut inferred.labels).len();
917 Ok((inferred, nc))
918}
919
920fn candidate_class_counts(signals: &ModelSignals) -> Result<BTreeSet<usize>, InferError> {
922 let input = signals
923 .inputs
924 .iter()
925 .find(|t| t.shape.len() == 4)
926 .ok_or_else(|| InferError::UnsupportedLayout("no rank-4 model input found".into()))?;
927 let (h, w, _) = classify_input_layout(input)?;
928 let k = match signals.outputs.iter().find(|o| o.shape.len() == 4) {
929 Some(proto) => classify_proto(proto, h, w)?.0,
930 None => 0,
931 };
932 Ok(signals
933 .outputs
934 .iter()
935 .filter_map(|o| find_anchors_dim(&o.shape, anchor_count(h, w)))
936 .filter_map(|(features, _)| features.checked_sub(4 + k))
937 .filter(|&nc| nc > 0)
938 .collect())
939}
940
941fn anchor_count(h: usize, w: usize) -> usize {
943 (h / 8) * (w / 8) + (h / 16) * (w / 16) + (h / 32) * (w / 32)
944}
945
946fn infer_from_meta(
947 signals: &ModelSignals,
948 meta: UltralyticsMeta,
949) -> Result<InferredSchema, InferError> {
950 let proto = signals.outputs.iter().find(|o| o.shape.len() == 4);
951 match (meta.task, proto.is_some()) {
952 (Some(Task::Segment), false) => {
953 return Err(InferError::UnsupportedLayout(
954 "segment task but no proto tensor".into(),
955 ));
956 }
957 (Some(Task::Detect), true) => {
958 return Err(InferError::UnsupportedLayout(
959 "detect task but proto tensor present".into(),
960 ));
961 }
962 _ => {}
963 }
964 let is_segment = meta.task == Some(Task::Segment) || proto.is_some();
965
966 let input = signals
967 .inputs
968 .iter()
969 .find(|t| t.shape.len() == 4)
970 .ok_or_else(|| {
971 InferError::UnsupportedLayout(format!(
972 "no rank-4 model input found; inputs seen: {:?}",
973 signals.inputs.iter().map(|t| &t.shape).collect::<Vec<_>>()
974 ))
975 })?;
976 let (h, w, input_dshape) = classify_input_layout(input)?;
977
978 let expected_anchors = anchor_count(h, w);
979 let nc = meta.names.len();
980
981 let (k, proto_dshape) = match proto {
989 Some(proto) => {
990 let (k, dshape) = classify_proto(proto, h, w)?;
991 (k, Some(dshape))
992 }
993 None => (0, None),
994 };
995 let feat = 4 + nc + k;
996 let e2e_feat = 6 + k;
997
998 let rank3_outputs: Vec<&TensorInfo> = signals
999 .outputs
1000 .iter()
1001 .filter(|o| o.shape.len() == 3)
1002 .collect();
1003
1004 let (out, dshape, version) = match meta.end2end {
1008 Some(true) => {
1009 let (out, n) = find_e2e_candidate(&rank3_outputs, expected_anchors, e2e_feat)?
1010 .ok_or_else(|| {
1011 InferError::UnsupportedLayout(format!(
1012 "metadata declares end2end=true but no output matches the YOLO26 \
1013 end-to-end layout [1, N<={expected_anchors}, {e2e_feat}]; \
1014 outputs seen: {:?}",
1015 rank3_outputs.iter().map(|o| &o.shape).collect::<Vec<_>>()
1016 ))
1017 })?;
1018 (
1019 out,
1020 vec![
1021 (DimName::Batch, 1),
1022 (DimName::NumBoxes, n),
1023 (DimName::NumFeatures, e2e_feat),
1024 ],
1025 DecoderVersion::Yolo26,
1026 )
1027 }
1028 Some(false) => {
1029 let (out, dshape) = classify_pre_nms(&rank3_outputs, expected_anchors, feat, nc)?;
1030 (out, dshape, DecoderVersion::Yolov8)
1031 }
1032 None => match classify_pre_nms(&rank3_outputs, expected_anchors, feat, nc) {
1033 Ok((out, dshape)) => (out, dshape, DecoderVersion::Yolov8),
1034 Err(err @ InferError::ClassCountMismatch { .. }) => return Err(err),
1035 Err(_) => {
1036 let (out, n) = find_e2e_candidate(&rank3_outputs, expected_anchors, e2e_feat)?
1037 .ok_or_else(|| {
1038 InferError::UnsupportedLayout(format!(
1039 "no output matches a known Ultralytics detection layout \
1040 (pre-NMS or end-to-end); outputs seen: {:?}",
1041 rank3_outputs.iter().map(|o| &o.shape).collect::<Vec<_>>()
1042 ))
1043 })?;
1044 (
1045 out,
1046 vec![
1047 (DimName::Batch, 1),
1048 (DimName::NumBoxes, n),
1049 (DimName::NumFeatures, e2e_feat),
1050 ],
1051 DecoderVersion::Yolo26,
1052 )
1053 }
1054 },
1055 };
1056
1057 check_quantized_boundary(out)?;
1058 if let Some(proto) = proto {
1059 check_quantized_boundary(proto)?;
1060 }
1061
1062 let normalized = match signals.source {
1079 ModelSource::TfLite => true,
1080 ModelSource::Onnx | ModelSource::CoreMl => false,
1081 ModelSource::Other => {
1082 return Err(InferError::UnknownBoxConvention);
1083 }
1084 };
1085 let det = LogicalOutput {
1086 name: Some(out.name.clone()),
1087 type_: Some(match version {
1094 DecoderVersion::Yolo26 => LogicalType::Detections,
1095 _ => LogicalType::Detection,
1096 }),
1097 shape: out.shape.clone(),
1098 dshape,
1099 decoder: Some(DecoderKind::Ultralytics),
1100 encoding: Some(BoxEncoding::Direct), score_format: Some(ScoreFormat::PerClass),
1102 normalized: Some(normalized),
1103 anchors: None,
1104 stride: None,
1105 dtype: Some(out.dtype),
1106 quantization: out.quantization.clone(),
1107 outputs: Vec::new(),
1108 activation_applied: None,
1109 activation_required: None,
1110 semantic: Default::default(),
1111 };
1112
1113 let mut outputs = vec![det];
1114 if let (Some(proto), Some(proto_dshape)) = (proto, proto_dshape) {
1115 outputs.push(LogicalOutput {
1116 name: Some(proto.name.clone()),
1117 type_: Some(LogicalType::Protos),
1118 shape: proto.shape.clone(),
1119 dshape: proto_dshape,
1120 decoder: Some(DecoderKind::Ultralytics),
1121 dtype: Some(proto.dtype),
1122 quantization: proto.quantization.clone(),
1123 ..default_logical()
1124 });
1125 }
1126
1127 let schema = SchemaV2 {
1128 schema_version: 2,
1129 input: Some(InputSpec {
1130 shape: input.shape.clone(),
1131 dshape: input_dshape,
1132 cameraadaptor: Some("rgb".into()),
1133 preprocessing: None,
1134 }),
1135 outputs,
1136 nms: match version {
1143 DecoderVersion::Yolo26 => None,
1144 _ => Some(NmsMode::ClassAware),
1145 },
1146 nms_multi_label: None,
1147 decoder_version: Some(version),
1148 };
1149
1150 let version_label = if version == DecoderVersion::Yolo26 {
1151 "YOLO26"
1152 } else {
1153 "YOLOv8/11"
1154 };
1155 let task_label = if is_segment { "segment" } else { "detect" };
1156 let description = format!("Ultralytics {version_label} {task_label}, {nc} classes");
1157
1158 Ok(InferredSchema {
1159 schema,
1160 labels: meta.names,
1161 description,
1162 })
1163}
1164
1165#[cfg(test)]
1166mod tests {
1167 use super::*;
1168
1169 fn infer_fixture_dir() -> std::path::PathBuf {
1181 std::env::var("EDGEFIRST_TESTDATA_DIR")
1182 .map(std::path::PathBuf::from)
1183 .unwrap_or_else(|_| {
1184 std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("testdata")
1185 })
1186 .join("infer")
1187 }
1188
1189 #[test]
1190 fn parse_names_python_dict_repr() {
1191 let s = "{0: 'person', 1: 'bicycle', 2: \"fire hydrant\"}";
1192 assert_eq!(
1193 parse_names(s).unwrap(),
1194 vec!["person", "bicycle", "fire hydrant"]
1195 );
1196 }
1197
1198 #[test]
1199 fn parse_names_json_object() {
1200 let s = r#"{"0": "person", "1": "bicycle"}"#;
1201 assert_eq!(parse_names(s).unwrap(), vec!["person", "bicycle"]);
1202 }
1203
1204 #[test]
1205 fn parse_names_orders_by_index_not_insertion() {
1206 let s = "{1: 'b', 0: 'a'}";
1207 assert_eq!(parse_names(s).unwrap(), vec!["a", "b"]);
1208 }
1209
1210 #[test]
1211 fn parse_names_rejects_garbage() {
1212 assert!(matches!(
1213 parse_names("not a dict"),
1214 Err(InferError::BadNames(_))
1215 ));
1216 }
1217
1218 #[test]
1219 fn parse_names_rejects_non_contiguous_indices() {
1220 match parse_names("{0: 'a', 2: 'b'}") {
1222 Err(InferError::BadNames(msg)) => {
1223 assert!(
1224 msg.contains("non-contiguous"),
1225 "unexpected BadNames message: {msg}"
1226 );
1227 }
1228 other => panic!("expected BadNames, got {other:?}"),
1229 }
1230
1231 match parse_names(r#"{"0": "a", "2": "b"}"#) {
1233 Err(InferError::BadNames(msg)) => {
1234 assert!(
1235 msg.contains("non-contiguous"),
1236 "unexpected BadNames message: {msg}"
1237 );
1238 }
1239 other => panic!("expected BadNames, got {other:?}"),
1240 }
1241 }
1242
1243 #[test]
1244 fn meta_from_onnx_props() {
1245 let mut m = BTreeMap::new();
1246 m.insert("names".into(), "{0: 'person'}".into());
1247 m.insert("task".into(), "detect".into());
1248 let u = UltralyticsMeta::from_metadata(&m).unwrap();
1249 assert_eq!(u.names, vec!["person"]);
1250 assert_eq!(u.task, Some(Task::Detect));
1251 }
1252
1253 #[test]
1254 fn meta_task_pose_rejected() {
1255 let mut m = BTreeMap::new();
1256 m.insert("names".into(), "{0: 'person'}".into());
1257 m.insert("task".into(), "pose".into());
1258 assert!(matches!(
1259 UltralyticsMeta::from_metadata(&m),
1260 Err(InferError::UnsupportedTask(_))
1261 ));
1262 }
1263
1264 #[test]
1265 fn meta_missing_names_is_not_ultralytics() {
1266 assert!(matches!(
1267 UltralyticsMeta::from_metadata(&BTreeMap::new()),
1268 Err(InferError::NotUltralytics(_))
1269 ));
1270 }
1271
1272 #[test]
1273 fn meta_end2end_true_or_false_from_onnx_flat_props() {
1274 let mut m = BTreeMap::new();
1275 m.insert("names".into(), "{0: 'person'}".into());
1276 m.insert("end2end".into(), "True".into());
1277 assert_eq!(
1278 UltralyticsMeta::from_metadata(&m).unwrap().end2end,
1279 Some(true)
1280 );
1281
1282 m.insert("end2end".into(), "False".into());
1283 assert_eq!(
1284 UltralyticsMeta::from_metadata(&m).unwrap().end2end,
1285 Some(false)
1286 );
1287 }
1288
1289 fn load_fixture_metadata(name: &str) -> BTreeMap<String, String> {
1294 let path = infer_fixture_dir().join(format!("{name}.signals.json"));
1295 let content = std::fs::read_to_string(&path)
1296 .unwrap_or_else(|e| panic!("failed to read fixture {}: {e}", path.display()));
1297 let json: serde_json::Value =
1298 serde_json::from_str(&content).expect("fixture is valid JSON");
1299 json["metadata"]
1300 .as_object()
1301 .expect("fixture has a metadata object")
1302 .iter()
1303 .map(|(k, v)| {
1304 (
1305 k.clone(),
1306 v.as_str().expect("metadata value is a string").to_string(),
1307 )
1308 })
1309 .collect()
1310 }
1311
1312 #[test]
1313 fn meta_from_real_tflite_metadata_json_envelope() {
1314 let meta = load_fixture_metadata("yolov8n_float32");
1315 let u = UltralyticsMeta::from_metadata(&meta).unwrap();
1316 assert_eq!(u.names.len(), 80);
1317 assert_eq!(u.task, Some(Task::Detect));
1318 }
1319
1320 #[test]
1321 fn meta_end2end_true_on_real_yolo26_tflite_fixture() {
1322 let meta = load_fixture_metadata("yolo26n_float32");
1323 let u = UltralyticsMeta::from_metadata(&meta).unwrap();
1324 assert_eq!(u.end2end, Some(true));
1325 }
1326
1327 #[test]
1328 fn meta_end2end_false_on_real_yolov8_onnx_fixture() {
1329 let meta = load_fixture_metadata("yolov8n");
1330 let u = UltralyticsMeta::from_metadata(&meta).unwrap();
1331 assert_eq!(u.end2end, Some(false));
1332 }
1333
1334 fn signals_from_fixture(name: &str) -> ModelSignals {
1338 let path = infer_fixture_dir().join(format!("{name}.signals.json"));
1339 let content = std::fs::read_to_string(&path)
1340 .unwrap_or_else(|e| panic!("failed to read fixture {}: {e}", path.display()));
1341 let json: serde_json::Value =
1342 serde_json::from_str(&content).expect("fixture is valid JSON");
1343
1344 let source = match json["source"]
1345 .as_str()
1346 .expect("fixture `source` is a string")
1347 {
1348 "onnx" => ModelSource::Onnx,
1349 "tflite" => ModelSource::TfLite,
1350 "coreml" => ModelSource::CoreMl,
1351 other => panic!("fixture {name}: unknown source `{other}`"),
1352 };
1353
1354 fn parse_dtype(s: &str) -> DType {
1355 match s {
1356 "float32" => DType::Float32,
1357 "float16" => DType::Float16,
1358 "int8" => DType::Int8,
1359 "uint8" => DType::Uint8,
1360 "int16" => DType::Int16,
1361 "uint16" => DType::Uint16,
1362 "int32" => DType::Int32,
1363 "uint32" => DType::Uint32,
1364 other => panic!("unknown dtype `{other}`"),
1365 }
1366 }
1367
1368 fn parse_tensor(v: &serde_json::Value) -> TensorInfo {
1369 TensorInfo {
1370 name: v["name"]
1371 .as_str()
1372 .expect("tensor `name` is a string")
1373 .to_string(),
1374 shape: v["shape"]
1375 .as_array()
1376 .expect("tensor `shape` is an array")
1377 .iter()
1378 .map(|d| d.as_u64().expect("shape dim is a number") as usize)
1379 .collect(),
1380 dtype: parse_dtype(v["dtype"].as_str().expect("tensor `dtype` is a string")),
1381 quantization: None,
1385 }
1386 }
1387
1388 let inputs = json["inputs"]
1389 .as_array()
1390 .expect("fixture `inputs` is an array")
1391 .iter()
1392 .map(parse_tensor)
1393 .collect();
1394 let outputs = json["outputs"]
1395 .as_array()
1396 .expect("fixture `outputs` is an array")
1397 .iter()
1398 .map(parse_tensor)
1399 .collect();
1400 let metadata = json["metadata"]
1401 .as_object()
1402 .expect("fixture `metadata` is an object")
1403 .iter()
1404 .map(|(k, v)| {
1405 (
1406 k.clone(),
1407 v.as_str().expect("metadata value is a string").to_string(),
1408 )
1409 })
1410 .collect();
1411
1412 ModelSignals {
1413 source,
1414 inputs,
1415 outputs,
1416 metadata,
1417 }
1418 }
1419
1420 fn synthetic_metadata(nc: usize, task: &str, end2end: &str) -> BTreeMap<String, String> {
1424 let names = (0..nc)
1425 .map(|i| format!("{i}: 'c{i}'"))
1426 .collect::<Vec<_>>()
1427 .join(", ");
1428 let mut m = BTreeMap::new();
1429 m.insert("names".into(), format!("{{{names}}}"));
1430 m.insert("task".into(), task.into());
1431 m.insert("end2end".into(), end2end.into());
1432 m
1433 }
1434
1435 #[test]
1436 fn infer_transposed_det_layout_classified() {
1437 let s = ModelSignals {
1443 source: ModelSource::TfLite,
1444 inputs: vec![TensorInfo {
1445 name: "images".into(),
1446 shape: vec![1, 640, 640, 3],
1447 dtype: DType::Float32,
1448 quantization: None,
1449 }],
1450 outputs: vec![TensorInfo {
1451 name: "output0".into(),
1452 shape: vec![1, 8400, 84],
1453 dtype: DType::Float32,
1454 quantization: None,
1455 }],
1456 metadata: synthetic_metadata(80, "detect", "False"),
1457 };
1458 let r = infer_ultralytics_schema(&s).unwrap();
1459 let o = &r.schema.outputs[0];
1460 assert_eq!(
1461 o.dshape,
1462 vec![
1463 (DimName::Batch, 1),
1464 (DimName::NumBoxes, 8400),
1465 (DimName::NumFeatures, 84),
1466 ]
1467 );
1468 assert_eq!(r.schema.decoder_version, Some(DecoderVersion::Yolov8));
1469 }
1470
1471 #[test]
1472 fn infer_yolo26n_onnx_end_to_end() {
1473 let s = signals_from_fixture("yolo26n");
1474 let r = infer_ultralytics_schema(&s).unwrap();
1475 assert_eq!(r.schema.decoder_version, Some(DecoderVersion::Yolo26));
1476 assert_eq!(r.schema.outputs[0].shape[2], 6);
1477 }
1478
1479 #[test]
1480 fn infer_class_count_mismatch_rejected() {
1481 let mut s = signals_from_fixture("yolov8n");
1482 s.metadata.insert("names".into(), "{0: 'only-one'}".into()); assert!(matches!(
1484 infer_ultralytics_schema(&s),
1485 Err(InferError::ClassCountMismatch {
1486 expected: 1,
1487 found: 80
1488 })
1489 ));
1490 }
1491
1492 #[test]
1493 fn infer_seg_class_count_mismatch_rejected() {
1494 let mut s = signals_from_fixture("yolov8n-seg");
1503 s.metadata.insert("names".into(), "{0: 'only-one'}".into()); assert!(matches!(
1505 infer_ultralytics_schema(&s),
1506 Err(InferError::ClassCountMismatch {
1507 expected: 1,
1508 found: 80
1509 })
1510 ));
1511 }
1512
1513 #[test]
1514 fn infer_end2end_flag_shape_conflict_rejected() {
1515 let mut s = signals_from_fixture("yolov8n");
1520 s.metadata.insert("end2end".into(), "True".into());
1521 assert!(matches!(
1522 infer_ultralytics_schema(&s),
1523 Err(InferError::UnsupportedLayout(_))
1524 ));
1525 }
1526
1527 #[test]
1528 fn infer_yolov8n_seg_onnx() {
1529 let r = infer_ultralytics_schema(&signals_from_fixture("yolov8n-seg")).unwrap();
1530 assert_eq!(r.schema.outputs.len(), 2);
1531 let proto = r
1532 .schema
1533 .outputs
1534 .iter()
1535 .find(|o| o.type_ == Some(LogicalType::Protos))
1536 .unwrap();
1537 assert_eq!(proto.dshape[1].0, DimName::NumProtos); let det = r
1539 .schema
1540 .outputs
1541 .iter()
1542 .find(|o| o.type_ == Some(LogicalType::Detection))
1543 .unwrap();
1544 assert!(det.shape.contains(&116)); }
1546
1547 #[test]
1548 fn infer_yolov8n_seg_tflite_proto_nchw() {
1549 let r = infer_ultralytics_schema(&signals_from_fixture("yolov8n-seg_float32")).unwrap();
1556 let proto = r
1557 .schema
1558 .outputs
1559 .iter()
1560 .find(|o| o.type_ == Some(LogicalType::Protos))
1561 .unwrap();
1562 assert_eq!(proto.dshape[1].0, DimName::NumProtos); assert_eq!(proto.dshape[1].1, 32);
1564 }
1565
1566 #[test]
1567 fn infer_yolo26n_seg_end_to_end() {
1568 let r = infer_ultralytics_schema(&signals_from_fixture("yolo26n-seg")).unwrap();
1569 assert_eq!(r.schema.decoder_version, Some(DecoderVersion::Yolo26));
1570 assert_eq!(r.schema.outputs.len(), 2);
1571 }
1572
1573 #[test]
1574 fn infer_nhwc_proto_layout_classified() {
1575 let s = ModelSignals {
1579 source: ModelSource::TfLite,
1580 inputs: vec![TensorInfo {
1581 name: "images".into(),
1582 shape: vec![1, 640, 640, 3],
1583 dtype: DType::Float32,
1584 quantization: None,
1585 }],
1586 outputs: vec![
1587 TensorInfo {
1588 name: "output0".into(),
1589 shape: vec![1, 116, 8400],
1590 dtype: DType::Float32,
1591 quantization: None,
1592 },
1593 TensorInfo {
1594 name: "output1".into(),
1595 shape: vec![1, 160, 160, 32],
1596 dtype: DType::Float32,
1597 quantization: None,
1598 },
1599 ],
1600 metadata: synthetic_metadata(80, "segment", "False"),
1601 };
1602 let r = infer_ultralytics_schema(&s).unwrap();
1603 let proto = r
1604 .schema
1605 .outputs
1606 .iter()
1607 .find(|o| o.type_ == Some(LogicalType::Protos))
1608 .unwrap();
1609 assert_eq!(proto.dshape[3].0, DimName::NumProtos); assert_eq!(proto.dshape[3].1, 32);
1611 }
1612
1613 #[test]
1617 fn infer_proto_output_carries_its_name_decoder_dtype_and_quantization() {
1618 let quant = Quantization {
1619 scale: vec![0.0125],
1620 zero_point: Some(vec![-7]),
1621 axis: None,
1622 dtype: Some(DType::Int8),
1623 };
1624 let s = ModelSignals {
1625 source: ModelSource::TfLite,
1626 inputs: vec![TensorInfo {
1627 name: "images".into(),
1628 shape: vec![1, 640, 640, 3],
1629 dtype: DType::Int8,
1630 quantization: None,
1631 }],
1632 outputs: vec![
1633 TensorInfo {
1634 name: "output0".into(),
1635 shape: vec![1, 116, 8400],
1636 dtype: DType::Int8,
1637 quantization: Some(Quantization {
1638 scale: vec![0.003_921_6],
1639 zero_point: Some(vec![0]),
1640 axis: None,
1641 dtype: Some(DType::Int8),
1642 }),
1643 },
1644 TensorInfo {
1645 name: "protos".into(),
1646 shape: vec![1, 160, 160, 32],
1647 dtype: DType::Int8,
1648 quantization: Some(quant.clone()),
1649 },
1650 ],
1651 metadata: synthetic_metadata(80, "segment", "False"),
1652 };
1653 let r = infer_ultralytics_schema(&s).unwrap();
1654 let proto = r
1655 .schema
1656 .outputs
1657 .iter()
1658 .find(|o| o.type_ == Some(LogicalType::Protos))
1659 .unwrap();
1660 assert_eq!(proto.name.as_deref(), Some("protos"));
1661 assert_eq!(proto.decoder, Some(DecoderKind::Ultralytics));
1662 assert_eq!(proto.dtype, Some(DType::Int8));
1663 assert_eq!(proto.quantization.as_ref(), Some(&quant));
1664 }
1665
1666 #[test]
1670 fn infer_treats_a_proto_tensor_without_a_task_as_segmentation() {
1671 let mut metadata = synthetic_metadata(80, "segment", "False");
1672 metadata.remove("task");
1673 let s = ModelSignals {
1674 source: ModelSource::TfLite,
1675 inputs: vec![TensorInfo {
1676 name: "images".into(),
1677 shape: vec![1, 640, 640, 3],
1678 dtype: DType::Float32,
1679 quantization: None,
1680 }],
1681 outputs: vec![
1682 TensorInfo {
1683 name: "output0".into(),
1684 shape: vec![1, 116, 8400],
1685 dtype: DType::Float32,
1686 quantization: None,
1687 },
1688 TensorInfo {
1689 name: "output1".into(),
1690 shape: vec![1, 160, 160, 32],
1691 dtype: DType::Float32,
1692 quantization: None,
1693 },
1694 ],
1695 metadata,
1696 };
1697 let r = infer_ultralytics_schema(&s).unwrap();
1698 assert!(
1699 r.schema
1700 .outputs
1701 .iter()
1702 .any(|o| o.type_ == Some(LogicalType::Protos)),
1703 "proto tensor present but no protos output was emitted"
1704 );
1705 assert!(r.description.contains("segment"), "{}", r.description);
1708 }
1709
1710 fn proto_tensor(shape: Vec<usize>) -> TensorInfo {
1711 TensorInfo {
1712 name: "protos".into(),
1713 shape,
1714 dtype: DType::Float32,
1715 quantization: None,
1716 }
1717 }
1718
1719 #[test]
1722 fn classify_proto_rejects_wrong_rank_or_batch() {
1723 for shape in [vec![1, 32, 160], vec![2, 32, 160, 160], vec![1, 160, 160]] {
1724 let err = classify_proto(&proto_tensor(shape.clone()), 640, 640).unwrap_err();
1725 assert!(
1726 matches!(err, InferError::UnsupportedLayout(_)),
1727 "shape {shape:?} gave {err:?}"
1728 );
1729 }
1730 }
1731
1732 #[test]
1735 fn classify_proto_rejects_a_shape_matching_only_one_spatial_dim() {
1736 let err = classify_proto(&proto_tensor(vec![1, 160, 7, 9]), 640, 640).unwrap_err();
1738 assert!(matches!(err, InferError::UnsupportedLayout(_)), "{err:?}");
1739
1740 let (k, dshape) = classify_proto(&proto_tensor(vec![1, 160, 160, 32]), 640, 640).unwrap();
1742 assert_eq!(k, 32);
1743 assert_eq!(dshape[3], (DimName::NumProtos, 32));
1744 }
1745
1746 #[test]
1749 fn unquote_requires_a_matching_pair_of_quotes() {
1750 assert_eq!(unquote("'abc'").as_deref(), Some("abc"));
1751 assert_eq!(unquote("\"abc\"").as_deref(), Some("abc"));
1752 assert_eq!(unquote("'abc"), None);
1753 assert_eq!(unquote("abc'"), None);
1754 assert_eq!(unquote("\"abc"), None);
1755 assert_eq!(unquote("abc\""), None);
1756 assert_eq!(unquote("abc"), None);
1757 }
1758
1759 #[test]
1763 fn split_top_level_commas_keeps_quoted_commas_together() {
1764 let parts = split_top_level_commas("a='x,y', b=2");
1765 assert_eq!(parts, vec!["a='x,y'".to_string(), " b=2".to_string()]);
1766
1767 let parts = split_top_level_commas("a=\"x,y\",b=\"p,q\"");
1768 assert_eq!(
1769 parts,
1770 vec!["a=\"x,y\"".to_string(), "b=\"p,q\"".to_string()]
1771 );
1772 }
1773
1774 #[test]
1777 fn infer_error_display_states_the_reason() {
1778 let err = InferError::NotUltralytics("no names key".into());
1779 let shown = err.to_string();
1780 assert!(shown.contains("no names key"), "{shown}");
1781 assert!(shown.contains("Ultralytics"), "{shown}");
1782
1783 let err = InferError::UnsupportedLayout("rank-3 proto".into());
1784 assert!(err.to_string().contains("rank-3 proto"));
1785 }
1786
1787 #[test]
1788 fn infer_segment_task_without_proto_rejected() {
1789 let s = ModelSignals {
1790 source: ModelSource::Onnx,
1791 inputs: vec![TensorInfo {
1792 name: "images".into(),
1793 shape: vec![1, 3, 640, 640],
1794 dtype: DType::Float32,
1795 quantization: None,
1796 }],
1797 outputs: vec![TensorInfo {
1798 name: "output0".into(),
1799 shape: vec![1, 84, 8400],
1800 dtype: DType::Float32,
1801 quantization: None,
1802 }],
1803 metadata: synthetic_metadata(80, "segment", "False"),
1804 };
1805 assert!(matches!(
1806 infer_ultralytics_schema(&s),
1807 Err(InferError::UnsupportedLayout(_))
1808 ));
1809 }
1810
1811 #[test]
1812 fn infer_detect_task_with_proto_rejected() {
1813 let s = ModelSignals {
1814 source: ModelSource::Onnx,
1815 inputs: vec![TensorInfo {
1816 name: "images".into(),
1817 shape: vec![1, 3, 640, 640],
1818 dtype: DType::Float32,
1819 quantization: None,
1820 }],
1821 outputs: vec![
1822 TensorInfo {
1823 name: "output0".into(),
1824 shape: vec![1, 84, 8400],
1825 dtype: DType::Float32,
1826 quantization: None,
1827 },
1828 TensorInfo {
1829 name: "output1".into(),
1830 shape: vec![1, 32, 160, 160],
1831 dtype: DType::Float32,
1832 quantization: None,
1833 },
1834 ],
1835 metadata: synthetic_metadata(80, "detect", "False"),
1836 };
1837 assert!(matches!(
1838 infer_ultralytics_schema(&s),
1839 Err(InferError::UnsupportedLayout(_))
1840 ));
1841 }
1842
1843 #[test]
1844 fn infer_int8_tflite_classifies_like_float32() {
1845 let r = infer_ultralytics_schema(&signals_from_fixture("yolov8n_int8")).unwrap();
1851 let det = &r.schema.outputs[0];
1852 assert_eq!(det.type_, Some(LogicalType::Detection));
1853 assert_eq!(det.dtype, Some(DType::Float32));
1854 assert_eq!(det.quantization, None);
1855 assert_eq!(r.schema.decoder_version, Some(DecoderVersion::Yolov8));
1856 assert_eq!(det.normalized, Some(true));
1857 }
1858
1859 #[test]
1860 fn infer_quantized_boundary_carries_quantization() {
1861 let quant = Quantization {
1866 scale: vec![0.02],
1867 zero_point: Some(vec![-5]),
1868 axis: None,
1869 dtype: Some(DType::Int8),
1870 };
1871 let s = ModelSignals {
1872 source: ModelSource::TfLite,
1873 inputs: vec![TensorInfo {
1874 name: "images".into(),
1875 shape: vec![1, 640, 640, 3],
1876 dtype: DType::Float32,
1877 quantization: None,
1878 }],
1879 outputs: vec![TensorInfo {
1880 name: "output0".into(),
1881 shape: vec![1, 84, 8400],
1882 dtype: DType::Int8,
1883 quantization: Some(quant.clone()),
1884 }],
1885 metadata: synthetic_metadata(80, "detect", "False"),
1886 };
1887 let r = infer_ultralytics_schema(&s).unwrap();
1888 let det = &r.schema.outputs[0];
1889 assert_eq!(det.dtype, Some(DType::Int8));
1890 assert_eq!(det.quantization, Some(quant));
1891 }
1892
1893 #[test]
1894 fn infer_rejects_batched_input() {
1895 let s = ModelSignals {
1900 source: ModelSource::Onnx,
1901 inputs: vec![TensorInfo {
1902 name: "images".into(),
1903 shape: vec![4, 3, 640, 640],
1904 dtype: DType::Float32,
1905 quantization: None,
1906 }],
1907 outputs: vec![TensorInfo {
1908 name: "output0".into(),
1909 shape: vec![1, 84, 8400],
1910 dtype: DType::Float32,
1911 quantization: None,
1912 }],
1913 metadata: synthetic_metadata(80, "detect", "False"),
1914 };
1915 match infer_ultralytics_schema(&s) {
1916 Err(InferError::UnsupportedLayout(msg)) => {
1917 assert!(msg.contains("batch 4"), "unexpected message: {msg}");
1918 }
1919 other => panic!("expected UnsupportedLayout, got {other:?}"),
1920 }
1921 }
1922
1923 #[test]
1924 fn infer_rejects_input_without_a_channel_dim() {
1925 let s = ModelSignals {
1926 source: ModelSource::Onnx,
1927 inputs: vec![TensorInfo {
1928 name: "images".into(),
1929 shape: vec![1, 4, 640, 640],
1930 dtype: DType::Float32,
1931 quantization: None,
1932 }],
1933 outputs: vec![TensorInfo {
1934 name: "output0".into(),
1935 shape: vec![1, 84, 8400],
1936 dtype: DType::Float32,
1937 quantization: None,
1938 }],
1939 metadata: synthetic_metadata(80, "detect", "False"),
1940 };
1941 match infer_ultralytics_schema(&s) {
1942 Err(InferError::UnsupportedLayout(msg)) => {
1943 assert!(msg.contains("channel dimension"), "unexpected: {msg}");
1944 }
1945 other => panic!("expected UnsupportedLayout, got {other:?}"),
1946 }
1947 }
1948
1949 #[test]
1950 fn infer_rejects_model_with_no_rank4_input() {
1951 let s = ModelSignals {
1952 source: ModelSource::Onnx,
1953 inputs: vec![TensorInfo {
1954 name: "images".into(),
1955 shape: vec![1, 3],
1956 dtype: DType::Float32,
1957 quantization: None,
1958 }],
1959 outputs: vec![TensorInfo {
1960 name: "output0".into(),
1961 shape: vec![1, 84, 8400],
1962 dtype: DType::Float32,
1963 quantization: None,
1964 }],
1965 metadata: synthetic_metadata(80, "detect", "False"),
1966 };
1967 match infer_ultralytics_schema(&s) {
1968 Err(InferError::UnsupportedLayout(msg)) => {
1969 assert!(msg.contains("no rank-4 model input"), "unexpected: {msg}");
1970 }
1971 other => panic!("expected UnsupportedLayout, got {other:?}"),
1972 }
1973 }
1974
1975 fn det_signals(nc: usize, out_shape: Vec<usize>, end2end: &str) -> ModelSignals {
1977 ModelSignals {
1978 source: ModelSource::Onnx,
1979 inputs: vec![TensorInfo {
1980 name: "images".into(),
1981 shape: vec![1, 3, 640, 640],
1982 dtype: DType::Float32,
1983 quantization: None,
1984 }],
1985 outputs: vec![TensorInfo {
1986 name: "output0".into(),
1987 shape: out_shape,
1988 dtype: DType::Float32,
1989 quantization: None,
1990 }],
1991 metadata: synthetic_metadata(nc, "detect", end2end),
1992 }
1993 }
1994
1995 #[test]
1996 fn infer_end_to_end_accepts_any_max_det_up_to_the_anchor_count() {
1997 for n in [1, 300, 2000, 8400] {
2009 let r = infer_ultralytics_schema(&det_signals(80, vec![1, n, 6], "True"))
2010 .unwrap_or_else(|e| panic!("max_det={n} must infer: {e}"));
2011 assert_eq!(r.schema.decoder_version, Some(DecoderVersion::Yolo26));
2012 assert_eq!(r.schema.outputs[0].type_, Some(LogicalType::Detections));
2013 assert_eq!(
2014 r.schema.outputs[0].dshape,
2015 vec![
2016 (DimName::Batch, 1),
2017 (DimName::NumBoxes, n),
2018 (DimName::NumFeatures, 6),
2019 ],
2020 "max_det={n}"
2021 );
2022 assert_eq!(r.schema.nms, None, "max_det={n}");
2024 }
2025 }
2026
2027 #[test]
2028 fn infer_end_to_end_rejects_more_rows_than_anchors() {
2029 match infer_ultralytics_schema(&det_signals(80, vec![1, 8401, 6], "True")) {
2033 Err(InferError::UnsupportedLayout(msg)) => {
2034 assert!(msg.contains("N<=8400"), "unexpected message: {msg}");
2035 }
2036 other => panic!("expected UnsupportedLayout, got {other:?}"),
2037 }
2038 }
2039
2040 #[test]
2041 fn infer_legacy_nms_model_is_unaffected_by_max_det() {
2042 let r = infer_ultralytics_schema(&det_signals(80, vec![1, 84, 8400], "False")).unwrap();
2048 assert_eq!(r.schema.decoder_version, Some(DecoderVersion::Yolov8));
2049 assert_eq!(r.schema.outputs[0].type_, Some(LogicalType::Detection));
2050 assert_eq!(
2051 r.schema.outputs[0].dshape,
2052 vec![
2053 (DimName::Batch, 1),
2054 (DimName::NumFeatures, 84),
2055 (DimName::NumBoxes, 8400),
2056 ]
2057 );
2058 assert_eq!(r.schema.nms, Some(NmsMode::ClassAware));
2062
2063 let t = infer_ultralytics_schema(&det_signals(80, vec![1, 8400, 84], "False")).unwrap();
2066 assert_eq!(t.schema.decoder_version, Some(DecoderVersion::Yolov8));
2067 assert_eq!(
2068 t.schema.outputs[0].dshape,
2069 vec![
2070 (DimName::Batch, 1),
2071 (DimName::NumBoxes, 8400),
2072 (DimName::NumFeatures, 84),
2073 ]
2074 );
2075 }
2076
2077 #[test]
2078 fn infer_pre_nms_wins_when_a_shape_reads_as_both() {
2079 let mut m = synthetic_metadata(2, "detect", "False");
2084 m.remove("end2end");
2085 let s = ModelSignals {
2086 metadata: m,
2087 ..det_signals(2, vec![1, 8400, 6], "False")
2088 };
2089 let r = infer_ultralytics_schema(&s).unwrap();
2090 assert_eq!(r.schema.decoder_version, Some(DecoderVersion::Yolov8));
2091 }
2092
2093 #[test]
2094 fn meta_empty_names_rejected() {
2095 let mut m = BTreeMap::new();
2098 m.insert("names".into(), "{}".into());
2099 match UltralyticsMeta::from_metadata(&m) {
2100 Err(InferError::BadNames(msg)) => assert!(msg.contains("empty"), "{msg}"),
2101 other => panic!("expected BadNames, got {other:?}"),
2102 }
2103 }
2104
2105 #[test]
2106 fn meta_malformed_end2end_is_an_error_not_an_absent_key() {
2107 let mut m = synthetic_metadata(80, "detect", "False");
2111 m.insert("end2end".into(), "yes".into());
2112 match UltralyticsMeta::from_metadata(&m) {
2113 Err(InferError::BadMetadata(msg)) => assert!(msg.contains("end2end"), "{msg}"),
2114 other => panic!("expected BadMetadata, got {other:?}"),
2115 }
2116
2117 m.remove("end2end");
2119 assert_eq!(UltralyticsMeta::from_metadata(&m).unwrap().end2end, None);
2120 }
2121
2122 #[test]
2123 fn meta_foreign_vendor_rejected_but_missing_provenance_allowed() {
2124 let mut m = synthetic_metadata(80, "detect", "False");
2128 m.insert("author".into(), "SomeOtherVendor".into());
2129 assert!(matches!(
2130 UltralyticsMeta::from_metadata(&m),
2131 Err(InferError::NotUltralytics(_))
2132 ));
2133
2134 m.insert("author".into(), "Ultralytics".into());
2136 m.insert("docs".into(), "https://docs.ultralytics.com".into());
2137 assert!(UltralyticsMeta::from_metadata(&m).is_ok());
2138
2139 m.remove("author");
2142 m.remove("docs");
2143 assert!(UltralyticsMeta::from_metadata(&m).is_ok());
2144 }
2145
2146 #[test]
2147 fn infer_zero_row_end_to_end_rejected() {
2148 assert!(matches!(
2151 infer_ultralytics_schema(&det_signals(80, vec![1, 0, 6], "True")),
2152 Err(InferError::UnsupportedLayout(_))
2153 ));
2154 }
2155
2156 #[test]
2157 fn infer_duplicate_matching_outputs_are_ambiguous() {
2158 let input = TensorInfo {
2159 name: "images".into(),
2160 shape: vec![1, 3, 640, 640],
2161 dtype: DType::Float32,
2162 quantization: None,
2163 };
2164 let t = |name: &str, shape: Vec<usize>| TensorInfo {
2165 name: name.into(),
2166 shape,
2167 dtype: DType::Float32,
2168 quantization: None,
2169 };
2170
2171 let pre = ModelSignals {
2174 source: ModelSource::Onnx,
2175 inputs: vec![input.clone()],
2176 outputs: vec![t("a", vec![1, 84, 8400]), t("b", vec![1, 84, 8400])],
2177 metadata: synthetic_metadata(80, "detect", "False"),
2178 };
2179 match infer_ultralytics_schema(&pre) {
2180 Err(InferError::AmbiguousLayout(msg)) => {
2181 assert!(msg.contains("`a`") && msg.contains("`b`"), "{msg}");
2182 }
2183 other => panic!("expected AmbiguousLayout, got {other:?}"),
2184 }
2185
2186 let e2e = ModelSignals {
2188 source: ModelSource::Onnx,
2189 inputs: vec![input],
2190 outputs: vec![t("a", vec![1, 300, 6]), t("b", vec![1, 200, 6])],
2191 metadata: synthetic_metadata(80, "detect", "True"),
2192 };
2193 assert!(matches!(
2194 infer_ultralytics_schema(&e2e),
2195 Err(InferError::AmbiguousLayout(_))
2196 ));
2197 }
2198
2199 #[test]
2200 fn infer_malformed_per_tensor_quantization_rejected() {
2201 let build = |q: Quantization| ModelSignals {
2202 source: ModelSource::TfLite,
2203 inputs: vec![TensorInfo {
2204 name: "images".into(),
2205 shape: vec![1, 640, 640, 3],
2206 dtype: DType::Float32,
2207 quantization: None,
2208 }],
2209 outputs: vec![TensorInfo {
2210 name: "output0".into(),
2211 shape: vec![1, 84, 8400],
2212 dtype: DType::Int8,
2213 quantization: Some(q),
2214 }],
2215 metadata: synthetic_metadata(80, "detect", "False"),
2216 };
2217
2218 match infer_ultralytics_schema(&build(Quantization {
2221 scale: vec![],
2222 zero_point: Some(vec![]),
2223 axis: None,
2224 dtype: Some(DType::Int8),
2225 })) {
2226 Err(InferError::UnsupportedQuantization(msg)) => {
2227 assert!(msg.contains("no scale"), "{msg}");
2228 }
2229 other => panic!("expected UnsupportedQuantization, got {other:?}"),
2230 }
2231
2232 match infer_ultralytics_schema(&build(Quantization {
2235 scale: vec![0.02],
2236 zero_point: Some(vec![-5, -7]),
2237 axis: None,
2238 dtype: Some(DType::Int8),
2239 })) {
2240 Err(InferError::UnsupportedQuantization(msg)) => {
2241 assert!(msg.contains("zero point"), "{msg}");
2242 }
2243 other => panic!("expected UnsupportedQuantization, got {other:?}"),
2244 }
2245
2246 assert!(infer_ultralytics_schema(&build(Quantization {
2248 scale: vec![0.02],
2249 zero_point: None,
2250 axis: None,
2251 dtype: Some(DType::Int8),
2252 }))
2253 .is_ok());
2254 }
2255
2256 #[test]
2257 fn infer_pre_nms_is_independent_of_output_order() {
2258 let det = TensorInfo {
2264 name: "output0".into(),
2265 shape: vec![1, 84, 8400],
2266 dtype: DType::Float32,
2267 quantization: None,
2268 };
2269 let aux = TensorInfo {
2270 name: "boxes".into(),
2271 shape: vec![1, 4, 8400],
2272 dtype: DType::Float32,
2273 quantization: None,
2274 };
2275 let build = |outputs: Vec<TensorInfo>| ModelSignals {
2276 source: ModelSource::Onnx,
2277 inputs: vec![TensorInfo {
2278 name: "images".into(),
2279 shape: vec![1, 3, 640, 640],
2280 dtype: DType::Float32,
2281 quantization: None,
2282 }],
2283 outputs,
2284 metadata: synthetic_metadata(80, "detect", "False"),
2285 };
2286
2287 let det_first = infer_ultralytics_schema(&build(vec![det.clone(), aux.clone()]))
2288 .expect("detection tensor first must classify");
2289 let aux_first = infer_ultralytics_schema(&build(vec![aux, det]))
2290 .expect("an auxiliary anchor-shaped tensor first must not shadow the detection tensor");
2291 assert_eq!(det_first.schema.outputs[0].name, Some("output0".into()));
2292 assert_eq!(aux_first.schema.outputs[0].name, Some("output0".into()));
2293 assert_eq!(det_first.schema, aux_first.schema);
2294 }
2295
2296 #[test]
2297 fn infer_class_count_mismatch_survives_a_full_scan() {
2298 let s = ModelSignals {
2301 source: ModelSource::Onnx,
2302 inputs: vec![TensorInfo {
2303 name: "images".into(),
2304 shape: vec![1, 3, 640, 640],
2305 dtype: DType::Float32,
2306 quantization: None,
2307 }],
2308 outputs: vec![TensorInfo {
2309 name: "output0".into(),
2310 shape: vec![1, 84, 8400],
2311 dtype: DType::Float32,
2312 quantization: None,
2313 }],
2314 metadata: synthetic_metadata(10, "detect", "False"),
2316 };
2317 assert!(matches!(
2318 infer_ultralytics_schema(&s),
2319 Err(InferError::ClassCountMismatch {
2320 expected: 10,
2321 found: 80
2322 })
2323 ));
2324 }
2325
2326 struct FixtureExpectation {
2336 name: &'static str,
2337 version: DecoderVersion,
2338 det_type: LogicalType,
2339 det_dshape: &'static [(DimName, usize)],
2341 normalized: bool,
2342 nms: Option<NmsMode>,
2343 proto: Option<(usize, usize, usize)>,
2345 }
2346
2347 const A: usize = 8400; const NC: usize = 80;
2349
2350 fn fixture_expectations() -> Vec<FixtureExpectation> {
2351 use DimName::{Batch, NumBoxes, NumFeatures};
2352 const DET: &[(DimName, usize)] = &[(Batch, 1), (NumFeatures, 4 + NC), (NumBoxes, A)];
2356 const SEG: &[(DimName, usize)] = &[(Batch, 1), (NumFeatures, 4 + NC + 32), (NumBoxes, A)];
2357 const E2E: &[(DimName, usize)] = &[(Batch, 1), (NumBoxes, 300), (NumFeatures, 6)];
2359 const E2E_SEG: &[(DimName, usize)] = &[(Batch, 1), (NumBoxes, 300), (NumFeatures, 6 + 32)];
2360
2361 let pre_nms = |name, det_dshape, normalized, proto| FixtureExpectation {
2362 name,
2363 version: DecoderVersion::Yolov8,
2364 det_type: LogicalType::Detection,
2365 det_dshape,
2366 normalized,
2367 nms: Some(NmsMode::ClassAware),
2368 proto,
2369 };
2370 let e2e = |name, det_dshape, normalized, proto| FixtureExpectation {
2371 name,
2372 version: DecoderVersion::Yolo26,
2373 det_type: LogicalType::Detections,
2374 det_dshape,
2375 normalized,
2376 nms: None,
2378 proto,
2379 };
2380 const P: Option<(usize, usize, usize)> = Some((32, 160, 160));
2383 vec![
2384 pre_nms("yolov8n", DET, false, None),
2385 pre_nms("yolo11n", DET, false, None),
2386 pre_nms("yolov8n-seg", SEG, false, P),
2387 pre_nms("yolo11n-seg", SEG, false, P),
2388 pre_nms("yolov8n_float32", DET, true, None),
2389 pre_nms("yolov8n_int8", DET, true, None),
2390 pre_nms("yolov8n-seg_float32", SEG, true, P),
2391 pre_nms("yolov8n-seg_int8", SEG, true, P),
2392 e2e("yolo26n", E2E, false, None),
2393 e2e("yolo26n-seg", E2E_SEG, false, P),
2394 e2e("yolo26n_float32", E2E, true, None),
2395 pre_nms("yolo26n_coreml", DET, false, None),
2403 ]
2404 }
2405
2406 #[test]
2407 fn every_captured_fixture_infers_the_expected_schema() {
2408 for want in fixture_expectations() {
2409 let name = want.name;
2410 let r = infer_ultralytics_schema(&signals_from_fixture(name))
2411 .unwrap_or_else(|e| panic!("{name}: inference failed: {e}"));
2412
2413 assert_eq!(r.labels.len(), NC, "{name}: class count");
2414 assert_eq!(
2415 r.labels[0], "person",
2416 "{name}: labels must be index-ordered"
2417 );
2418 assert_eq!(r.labels[NC - 1], "toothbrush", "{name}: labels index order");
2419 assert_eq!(
2420 r.schema.decoder_version,
2421 Some(want.version),
2422 "{name}: version"
2423 );
2424 assert_eq!(r.schema.nms, want.nms, "{name}: nms mode");
2425
2426 let det = &r.schema.outputs[0];
2427 assert_eq!(det.type_, Some(want.det_type), "{name}: detection type");
2428 assert_eq!(det.dshape, want.det_dshape, "{name}: detection dshape");
2429 assert_eq!(det.normalized, Some(want.normalized), "{name}: normalized");
2430 assert_eq!(
2431 det.decoder,
2432 Some(DecoderKind::Ultralytics),
2433 "{name}: decoder"
2434 );
2435 assert_eq!(det.encoding, Some(BoxEncoding::Direct), "{name}: encoding");
2436 assert_eq!(
2437 det.score_format,
2438 Some(ScoreFormat::PerClass),
2439 "{name}: scores"
2440 );
2441
2442 match want.proto {
2443 None => {
2444 assert_eq!(r.schema.outputs.len(), 1, "{name}: detection-only");
2445 }
2446 Some((k, ph, pw)) => {
2447 assert_eq!(r.schema.outputs.len(), 2, "{name}: detection + proto");
2448 let proto = &r.schema.outputs[1];
2449 assert_eq!(proto.type_, Some(LogicalType::Protos), "{name}: proto type");
2450 assert_eq!(
2451 proto.dshape,
2452 vec![
2453 (DimName::Batch, 1),
2454 (DimName::NumProtos, k),
2455 (DimName::Height, ph),
2456 (DimName::Width, pw),
2457 ],
2458 "{name}: proto dshape"
2459 );
2460 }
2461 }
2462
2463 let task = if want.proto.is_some() {
2464 "segment"
2465 } else {
2466 "detect"
2467 };
2468 assert!(
2469 r.description.contains(task) && r.description.contains("Ultralytics"),
2470 "{name}: description was `{}`",
2471 r.description
2472 );
2473 }
2474 }
2475
2476 #[test]
2477 fn every_fixture_is_asserted() {
2478 let dir = infer_fixture_dir();
2482 let mut on_disk: Vec<String> = std::fs::read_dir(&dir)
2483 .unwrap_or_else(|e| panic!("{} exists: {e}", dir.display()))
2484 .filter_map(|e| {
2485 let name = e.ok()?.file_name().into_string().ok()?;
2486 name.strip_suffix(".signals.json").map(str::to_string)
2487 })
2488 .collect();
2489 on_disk.sort();
2490
2491 let mut asserted: Vec<String> = fixture_expectations()
2492 .iter()
2493 .map(|f| f.name.to_string())
2494 .collect();
2495 asserted.sort();
2496
2497 assert_eq!(
2498 on_disk, asserted,
2499 "every *.signals.json fixture must appear in fixture_expectations()"
2500 );
2501 }
2502
2503 #[test]
2504 fn infer_other_source_refuses_rather_than_guessing_normalization() {
2505 let s = ModelSignals {
2512 source: ModelSource::Other,
2513 inputs: vec![TensorInfo {
2514 name: "images".into(),
2515 shape: vec![1, 3, 640, 640],
2516 dtype: DType::Float32,
2517 quantization: None,
2518 }],
2519 outputs: vec![TensorInfo {
2520 name: "output0".into(),
2521 shape: vec![1, 84, 8400],
2522 dtype: DType::Float32,
2523 quantization: None,
2524 }],
2525 metadata: synthetic_metadata(80, "detect", "False"),
2526 };
2527 assert!(matches!(
2528 infer_ultralytics_schema(&s),
2529 Err(InferError::UnknownBoxConvention)
2530 ));
2531
2532 let ok = ModelSignals {
2535 source: ModelSource::Onnx,
2536 ..s
2537 };
2538 assert_eq!(
2539 infer_ultralytics_schema(&ok).unwrap().schema.outputs[0].normalized,
2540 Some(false)
2541 );
2542 }
2543
2544 #[test]
2545 fn infer_per_channel_quantization_rejected() {
2546 let s = ModelSignals {
2555 source: ModelSource::TfLite,
2556 inputs: vec![TensorInfo {
2557 name: "images".into(),
2558 shape: vec![1, 640, 640, 3],
2559 dtype: DType::Float32,
2560 quantization: None,
2561 }],
2562 outputs: vec![TensorInfo {
2563 name: "output0".into(),
2564 shape: vec![1, 84, 8400],
2565 dtype: DType::Int8,
2566 quantization: Some(Quantization {
2567 scale: vec![0.02, 0.03],
2568 zero_point: Some(vec![-5, -7]),
2569 axis: Some(1),
2570 dtype: Some(DType::Int8),
2571 }),
2572 }],
2573 metadata: synthetic_metadata(80, "detect", "False"),
2574 };
2575 match infer_ultralytics_schema(&s) {
2576 Err(InferError::UnsupportedQuantization(msg)) => {
2577 assert!(msg.contains("per-channel"), "unexpected message: {msg}");
2578 assert!(
2579 msg.contains("output0"),
2580 "message must name the tensor: {msg}"
2581 );
2582 }
2583 other => panic!("expected UnsupportedQuantization, got {other:?}"),
2584 }
2585 }
2586
2587 #[test]
2588 fn infer_per_channel_proto_quantization_rejected() {
2589 let s = ModelSignals {
2592 source: ModelSource::TfLite,
2593 inputs: vec![TensorInfo {
2594 name: "images".into(),
2595 shape: vec![1, 640, 640, 3],
2596 dtype: DType::Float32,
2597 quantization: None,
2598 }],
2599 outputs: vec![
2600 TensorInfo {
2601 name: "output0".into(),
2602 shape: vec![1, 116, 8400],
2603 dtype: DType::Float32,
2604 quantization: None,
2605 },
2606 TensorInfo {
2607 name: "output1".into(),
2608 shape: vec![1, 32, 160, 160],
2609 dtype: DType::Int8,
2610 quantization: Some(Quantization {
2611 scale: vec![0.02, 0.03],
2612 zero_point: Some(vec![-5, -7]),
2613 axis: Some(1),
2614 dtype: Some(DType::Int8),
2615 }),
2616 },
2617 ],
2618 metadata: synthetic_metadata(80, "segment", "False"),
2619 };
2620 match infer_ultralytics_schema(&s) {
2621 Err(InferError::UnsupportedQuantization(msg)) => {
2622 assert!(
2623 msg.contains("output1"),
2624 "message must name the proto: {msg}"
2625 );
2626 }
2627 other => panic!("expected UnsupportedQuantization, got {other:?}"),
2628 }
2629 }
2630
2631 #[test]
2632 fn infer_integer_output_without_quant_rejected() {
2633 let s = ModelSignals {
2637 source: ModelSource::TfLite,
2638 inputs: vec![TensorInfo {
2639 name: "images".into(),
2640 shape: vec![1, 640, 640, 3],
2641 dtype: DType::Float32,
2642 quantization: None,
2643 }],
2644 outputs: vec![TensorInfo {
2645 name: "output0".into(),
2646 shape: vec![1, 84, 8400],
2647 dtype: DType::Int8,
2648 quantization: None,
2649 }],
2650 metadata: synthetic_metadata(80, "detect", "False"),
2651 };
2652 match infer_ultralytics_schema(&s) {
2653 Err(InferError::UnsupportedLayout(msg)) => {
2654 assert!(
2655 msg.contains("quantization"),
2656 "unexpected UnsupportedLayout message: {msg}"
2657 );
2658 }
2659 other => panic!("expected UnsupportedLayout, got {other:?}"),
2660 }
2661 }
2662
2663 #[test]
2664 fn infer_pre_nms_rejects_non_unit_batch() {
2665 let mut s = signals_from_fixture("yolov8n");
2670 s.outputs[0].shape = vec![2, 84, 8400];
2671 assert!(matches!(
2672 infer_ultralytics_schema(&s),
2673 Err(InferError::UnsupportedLayout(_))
2674 ));
2675 }
2676
2677 #[test]
2678 fn coreml_source_resolves_a_box_convention() {
2679 let s = ModelSignals {
2682 source: ModelSource::CoreMl,
2683 inputs: vec![TensorInfo {
2684 name: "images".into(),
2685 shape: vec![1, 3, 640, 640],
2686 dtype: DType::Float16,
2687 quantization: None,
2688 }],
2689 outputs: vec![TensorInfo {
2690 name: "output0".into(),
2691 shape: vec![1, 84, 8400],
2692 dtype: DType::Float16,
2693 quantization: None,
2694 }],
2695 metadata: synthetic_metadata(80, "detect", "False"),
2696 };
2697 let r = infer_ultralytics_schema(&s)
2698 .expect("CoreML signals must infer, not raise UnknownBoxConvention");
2699 let o = &r.schema.outputs[0];
2700 assert_eq!(
2701 o.normalized,
2702 Some(false),
2703 "Ultralytics CoreML exports trace the raw model, so boxes are \
2704 pixel-space exactly as in the ONNX export"
2705 );
2706 }
2707
2708 #[test]
2709 fn coreml_nms_pipeline_export_is_refused_not_classified() {
2710 let s = ModelSignals {
2724 source: ModelSource::CoreMl,
2725 inputs: vec![TensorInfo {
2726 name: "image".into(),
2727 shape: vec![1, 3, 640, 640],
2728 dtype: DType::Float16,
2729 quantization: None,
2730 }],
2731 outputs: vec![
2732 TensorInfo {
2733 name: "confidence".into(),
2734 shape: vec![],
2735 dtype: DType::Float32,
2736 quantization: None,
2737 },
2738 TensorInfo {
2739 name: "coordinates".into(),
2740 shape: vec![],
2741 dtype: DType::Float32,
2742 quantization: None,
2743 },
2744 ],
2745 metadata: synthetic_metadata(80, "detect", "False"),
2746 };
2747 assert!(matches!(
2748 infer_ultralytics_schema(&s),
2749 Err(InferError::UnsupportedLayout(_))
2750 ));
2751 }
2752}