fix: 🐛 Prioritize model metadata for image size when not provided (#37)

Signed-off-by: Onuralp SEZER <onuralp@ultralytics.com>

Author: onuralpszr

README.md

@@ -103,7 +103,7 @@ cargo run --release -- predict --model <model.onnx> --source <source>
 | `--source`  | `-s`  | Input source (image, video, webcam index, or URL) | `Task dependent Ultralytics URL assets` |
 | `--conf`    |       | Confidence threshold                              | `0.25`                                  |
 | `--iou`     |       | IoU threshold for NMS                             | `0.45`                                  |
-| `--imgsz`   |       | Inference image size                              | `640`                                   |
+| `--imgsz`   |       | Inference image size                              | `Model metadata`                        |
 | `--half`    |       | Use FP16 half-precision inference                 | `false`                                 |
 | `--save`    |       | Save annotated images to runs/<task>/predict      | `false`                                 |
 | `--show`    |       | Display results in a window                       | `false`                                 |

src/main.rs

@@ -75,7 +75,7 @@ fn run_prediction(args: &[String]) {
     let mut source_path: Option<&String> = None;
     let mut conf_threshold = 0.25_f32;
     let mut iou_threshold = 0.45_f32;
-    let mut imgsz: usize = 640;
+    let mut imgsz: Option<usize> = None;
     let mut save = false;
     let mut half = false;
     let mut verbose = true;
@@ -144,8 +144,16 @@ fn run_prediction(args: &[String]) {
             }
             "--imgsz" => {
                 if i + 1 < args.len() {
-                    imgsz = args[i + 1].parse().unwrap_or(640);
-                    i += 2;
+                    if let Ok(val) = args[i + 1].parse::<usize>() {
+                        imgsz = Some(val);
+                        i += 2;
+                    } else {
+                        eprintln!(
+                            "Error: --imgsz requires a positive integer, found '{}'",
+                            args[i + 1]
+                        );
+                        process::exit(1);
+                    }
                 } else {
                     eprintln!("Error: --imgsz requires a value");
                     process::exit(1);
@@ -201,8 +209,10 @@ fn run_prediction(args: &[String]) {
         .with_iou(iou_threshold)
         .with_half(half);
 
-    // Apply imgsz (default 640)
-    config = config.with_imgsz(imgsz, imgsz);
+    // Apply imgsz if specified
+    if let Some(sz) = imgsz {
+        config = config.with_imgsz(sz, sz);
+    }
 
     let mut model = match YOLOModel::load_with_config(&model_path, config) {
         Ok(m) => m,

src/metadata.rs

@@ -35,7 +35,7 @@ pub struct ModelMetadata {
     /// Batch size the model was exported with.
     pub batch: usize,
     /// Input image size as (height, width).
-    pub imgsz: (usize, usize),
+    pub imgsz: Option<(usize, usize)>,
     /// Number of input channels (typically 3 for RGB).
     pub channels: usize,
     /// Whether the model uses FP16 (half precision).
@@ -171,7 +171,7 @@ impl ModelMetadata {
 
     /// Parse the imgsz field which can be a YAML list.
     #[allow(clippy::unnecessary_wraps)]
-    fn parse_imgsz(yaml_str: &str, imgsz_line: &str) -> Result<(usize, usize)> {
+    fn parse_imgsz(yaml_str: &str, imgsz_line: &str) -> Result<Option<(usize, usize)>> {
         // Check if imgsz is on a single line like "imgsz: [640, 640]"
         if let Some(bracket_start) = imgsz_line.find('[')
             && let Some(bracket_end) = imgsz_line.find(']')
@@ -181,7 +181,7 @@ impl ModelMetadata {
                 .filter_map(|s| s.trim().parse().ok())
                 .collect();
             if values.len() >= 2 {
-                return Ok((values[0], values[1]));
+                return Ok(Some((values[0], values[1])));
             }
         }
 
@@ -210,10 +210,9 @@ impl ModelMetadata {
         }
 
         if imgsz_values.len() >= 2 {
-            Ok((imgsz_values[0], imgsz_values[1]))
+            Ok(Some((imgsz_values[0], imgsz_values[1])))
         } else {
-            // Default to 640x640
-            Ok((640, 640))
+            Ok(None)
         }
     }
 
@@ -334,7 +333,7 @@ impl Default for ModelMetadata {
             task: Task::Detect,
             stride: 32,
             batch: 1,
-            imgsz: (640, 640),
+            imgsz: None,
             channels: 3,
             half: false,
             names: HashMap::new(),
@@ -374,7 +373,7 @@ channels: 3
         assert_eq!(metadata.task, Task::Detect);
         assert_eq!(metadata.stride, 32);
         assert_eq!(metadata.batch, 1);
-        assert_eq!(metadata.imgsz, (640, 640));
+        assert_eq!(metadata.imgsz, Some((640, 640)));
         assert_eq!(metadata.channels, 3);
         assert_eq!(metadata.num_classes(), 4);
         assert_eq!(metadata.class_name(0), Some("person"));
@@ -387,14 +386,14 @@ channels: 3
     fn test_parse_inline_imgsz() {
         let yaml = "task: detect\nimgsz: [640, 640]\nstride: 32";
         let metadata = ModelMetadata::from_yaml_str(yaml).unwrap();
-        assert_eq!(metadata.imgsz, (640, 640));
+        assert_eq!(metadata.imgsz, Some((640, 640)));
     }
 
     #[test]
     fn test_default_metadata() {
         let metadata = ModelMetadata::default();
         assert_eq!(metadata.task, Task::Detect);
         assert_eq!(metadata.stride, 32);
-        assert_eq!(metadata.imgsz, (640, 640));
+        assert_eq!(metadata.imgsz, None);
     }
 }

src/model.rs

@@ -208,6 +208,17 @@ impl YOLOModel {
             )
         });
 
+        // Check for dynamic input dimensions
+        // Dimensions are typically [-1, 3, -1, -1] for dynamic batch/height/width
+        // ort 2.0 returns Option<i64> (None means dynamic/unknown)
+        let is_dynamic = input_info.is_some_and(|i| {
+            if let ValueType::Tensor { shape, .. } = i.dtype() {
+                shape.iter().any(|d| *d == -1 || *d == 0)
+            } else {
+                false
+            }
+        });
+
         // Use FP16 input if model tensor type requires it
         let fp16_input = model_input_fp16;
 
@@ -217,14 +228,53 @@ impl YOLOModel {
             .map(|o| o.name().to_string())
             .collect();
 
-        // Update config with model metadata if not overridden
-        // Sync half flag with model's half metadata (for display purposes)
+        // Resolve image size
+        // Priority:
+        // 1. User config
+        // 2. Model metadata
+        // 3. Dynamic default (1024 for OBB, 640 for others)
+        // 4. Static input shape
+        // 5. Hard default (640)
+        let resolved_imgsz = if let Some(sz) = config.imgsz {
+            sz
+        } else if let Some(sz) = metadata.imgsz {
+            sz
+        } else if is_dynamic {
+            // Dynamic input without metadata -> apply robust defaults
+            match metadata.task {
+                Task::Obb => (1024, 1024),
+                _ => (640, 640),
+            }
+        } else {
+            // Static input without metadata -> try to read from tensor shape
+            // Typically [1, 3, H, W]
+            input_info
+                .and_then(|i| {
+                    if let ValueType::Tensor { shape, .. } = i.dtype() {
+                        if shape.len() == 4 && shape[2] > 0 && shape[3] > 0 {
+                            #[allow(clippy::cast_sign_loss)]
+                            Some((shape[2] as usize, shape[3] as usize))
+                        } else {
+                            None
+                        }
+                    } else {
+                        None
+                    }
+                })
+                .unwrap_or((640, 640))
+        };
+
+        // Update config with resolved values
         let config = InferenceConfig {
-            imgsz: config.imgsz.or(Some(metadata.imgsz)),
+            imgsz: Some(resolved_imgsz),
             half: config.half || metadata.half, // Use half if user requested OR model was exported with half
             ..config
         };
 
+        // Ensure metadata reflects the resolved size
+        let mut metadata = metadata;
+        metadata.imgsz = Some(resolved_imgsz);
+
         let mut model = Self {
             session,
             metadata,
@@ -253,7 +303,11 @@ impl YOLOModel {
             return Ok(());
         }
 
-        let target_size = self.config.imgsz.unwrap_or(self.metadata.imgsz);
+        let target_size = self
+            .config
+            .imgsz
+            .or(self.metadata.imgsz)
+            .unwrap_or((640, 640));
 
         // Sanity check to prevent huge allocations from invalid imgsz
         if target_size.0 > Self::MAX_IMGSZ || target_size.1 > Self::MAX_IMGSZ {
@@ -379,7 +433,11 @@ impl YOLOModel {
     /// Vector of Results.
     pub fn predict_image(&mut self, image: &DynamicImage, path: String) -> Result<Vec<Results>> {
         // Get target size from config or metadata
-        let target_size = self.config.imgsz.unwrap_or(self.metadata.imgsz);
+        let target_size = self
+            .config
+            .imgsz
+            .or(self.metadata.imgsz)
+            .unwrap_or((640, 640));
 
         // Preprocess - generate FP16 tensor if model expects FP16 input
         let start_preprocess = Instant::now();
@@ -595,8 +653,8 @@ impl YOLOModel {
 
     /// Get the model's input size.
     #[must_use]
-    pub const fn imgsz(&self) -> (usize, usize) {
-        self.metadata.imgsz
+    pub fn imgsz(&self) -> (usize, usize) {
+        self.metadata.imgsz.unwrap_or((640, 640))
     }
 
     /// Get the model's stride.