From c39160760e8e7c9ce00fb746b40979721d46af08 Mon Sep 17 00:00:00 2001
From: dbogunowicz <97082108+dbogunowicz@users.noreply.github.com>
Date: Wed, 22 Mar 2023 23:49:37 +0100
Subject: [PATCH] [Cherry Pick] Support for dynamic input YOLO models (1.4.2)
 (#971)

* initial commit

* bump version to 1.4.2

---------

Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com>
---
 src/deepsparse/version.py          | 2 +-
 src/deepsparse/yolo/pipelines.py   | 6 ++++++
 src/deepsparse/yolo/utils/utils.py | 2 ++
 tests/test_benchmark.py            | 2 +-
 4 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/deepsparse/version.py b/src/deepsparse/version.py
index f33683b286..b05a9a7b1a 100644
--- a/src/deepsparse/version.py
+++ b/src/deepsparse/version.py
@@ -39,7 +39,7 @@
     from deepsparse.generated_version import is_enterprise, is_release, splash, version
 except Exception:
     # otherwise, fall back to version info in this file
-    version = "1.4.1"
+    version = "1.4.2"
     is_release = False
     is_enterprise = False
     splash = (
diff --git a/src/deepsparse/yolo/pipelines.py b/src/deepsparse/yolo/pipelines.py
index a7a0152d98..9abb910b77 100644
--- a/src/deepsparse/yolo/pipelines.py
+++ b/src/deepsparse/yolo/pipelines.py
@@ -158,6 +158,12 @@ class properties into an inference ready onnx file to be compiled by the
         model_path = model_to_path(self.model_path)
         if self._image_size is None:
             self._image_size = get_onnx_expected_image_shape(onnx.load(model_path))
+            if self._image_size == (0, 0):
+                raise ValueError(
+                    "The model does not have a static image size shape. "
+                    "Specify the expected image size by passing the"
+                    "`image_size` argument to the pipeline."
+                )
         else:
             # override model input shape to given image size
             if isinstance(self._image_size, int):
diff --git a/src/deepsparse/yolo/utils/utils.py b/src/deepsparse/yolo/utils/utils.py
index 07e7b87ec2..baa4c18721 100644
--- a/src/deepsparse/yolo/utils/utils.py
+++ b/src/deepsparse/yolo/utils/utils.py
@@ -359,6 +359,8 @@ def modify_yolo_onnx_input_shape(
     model_input = model.graph.input[0]
 
     initial_x, initial_y = get_onnx_expected_image_shape(model)
+    if initial_x == initial_y == 0:
+        initial_x, initial_y = image_shape
 
     if not (isinstance(initial_x, int) and isinstance(initial_y, int)):
         return model_path, None  # model graph does not have static integer input shape
diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py
index 8678864735..027b7596fc 100644
--- a/tests/test_benchmark.py
+++ b/tests/test_benchmark.py
@@ -75,7 +75,7 @@ def test_benchmark_help():
         ),
         (
             "zoo:cv/detection/yolov5-s/pytorch/ultralytics/coco/base-none",
-            ["-pin", "numa"],
+            ["-pin", "numa", "-shapes", "[1,3,640,640]"],
         ),
         (
             "zoo:cv/detection/yolov5-s/pytorch/ultralytics/coco/pruned-aggressive_96",