Export static shape for detection/instance-segmentation models (#4113)

* Add onnxruntime dependency and update dynamic axes in detection/inst-seg models
openvinotoolkit · Nov 19, 2024 · 0bb3f2e · 0bb3f2e
1 parent 046602d
commit 0bb3f2e
Show file tree

Hide file tree

Showing 11 changed files with 156 additions and 139 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -55,6 +55,7 @@ dev = [
     "pytest-mock",
     "pytest-csv",
     "pytest-cov",
+    "onnxruntime==1.20.0",
 ]
 
 docs = [

diff --git a/src/otx/algo/detection/atss.py b/src/otx/algo/detection/atss.py
@@ -183,7 +183,7 @@ def _exporter(self) -> OTXModelExporter:
                 "input_names": ["image"],
                 "output_names": ["boxes", "labels"],
                 "dynamic_axes": {
-                    "image": {0: "batch", 2: "height", 3: "width"},
+                    "image": {0: "batch"},
                     "boxes": {0: "batch", 1: "num_dets"},
                     "labels": {0: "batch", 1: "num_dets"},
                 },

diff --git a/src/otx/algo/detection/rtmdet.py b/src/otx/algo/detection/rtmdet.py
@@ -138,7 +138,7 @@ def _exporter(self) -> OTXModelExporter:
                 "input_names": ["image"],
                 "output_names": ["boxes", "labels"],
                 "dynamic_axes": {
-                    "image": {0: "batch", 2: "height", 3: "width"},
+                    "image": {0: "batch"},
                     "boxes": {0: "batch", 1: "num_dets"},
                     "labels": {0: "batch", 1: "num_dets"},
                 },

diff --git a/src/otx/algo/detection/ssd.py b/src/otx/algo/detection/ssd.py
@@ -346,7 +346,7 @@ def _exporter(self) -> OTXModelExporter:
                 "input_names": ["image"],
                 "output_names": ["boxes", "labels"],
                 "dynamic_axes": {
-                    "image": {0: "batch", 2: "height", 3: "width"},
+                    "image": {0: "batch"},
                     "boxes": {0: "batch", 1: "num_dets"},
                     "labels": {0: "batch", 1: "num_dets"},
                 },

diff --git a/src/otx/algo/detection/yolov9.py b/src/otx/algo/detection/yolov9.py
@@ -125,7 +125,7 @@ def _exporter(self) -> OTXModelExporter:
                 "export_params": True,
                 "opset_version": 11,
                 "dynamic_axes": {
-                    "image": {0: "batch", 2: "height", 3: "width"},
+                    "image": {0: "batch"},
                     "boxes": {0: "batch", 1: "num_dets"},
                     "labels": {0: "batch", 1: "num_dets"},
                 },

diff --git a/src/otx/algo/detection/yolox.py b/src/otx/algo/detection/yolox.py
@@ -158,7 +158,7 @@ def _exporter(self) -> OTXModelExporter:
                 "export_params": True,
                 "opset_version": 11,
                 "dynamic_axes": {
-                    "image": {0: "batch", 2: "height", 3: "width"},
+                    "image": {0: "batch"},
                     "boxes": {0: "batch", 1: "num_dets"},
                     "labels": {0: "batch", 1: "num_dets"},
                 },

diff --git a/src/otx/algo/instance_segmentation/maskdino.py b/src/otx/algo/instance_segmentation/maskdino.py
@@ -190,7 +190,7 @@ def _exporter(self) -> OTXModelExporter:
             onnx_export_configuration={
                 "input_names": ["image"],
                 "dynamic_axes": {
-                    "image": {0: "batch", 2: "height", 3: "width"},
+                    "image": {0: "batch"},
                     "boxes": {0: "batch", 1: "num_dets"},
                     "labels": {0: "batch", 1: "num_dets"},
                     "masks": {0: "batch", 1: "num_dets", 2: "height", 3: "width"},

diff --git a/src/otx/algo/instance_segmentation/maskrcnn.py b/src/otx/algo/instance_segmentation/maskrcnn.py
@@ -319,7 +319,7 @@ def _exporter(self) -> OTXModelExporter:
                 "input_names": ["image"],
                 "output_names": ["boxes", "labels", "masks"],
                 "dynamic_axes": {
-                    "image": {0: "batch", 2: "height", 3: "width"},
+                    "image": {0: "batch"},
                     "boxes": {0: "batch", 1: "num_dets"},
                     "labels": {0: "batch", 1: "num_dets"},
                     "masks": {0: "batch", 1: "num_dets", 2: "height", 3: "width"},

diff --git a/src/otx/algo/instance_segmentation/maskrcnn_tv.py b/src/otx/algo/instance_segmentation/maskrcnn_tv.py
@@ -214,7 +214,7 @@ def _exporter(self) -> OTXModelExporter:
                 "input_names": ["image"],
                 "output_names": ["boxes", "labels", "masks"],
                 "dynamic_axes": {
-                    "image": {0: "batch", 2: "height", 3: "width"},
+                    "image": {0: "batch"},
                     "boxes": {0: "batch", 1: "num_dets"},
                     "labels": {0: "batch", 1: "num_dets"},
                     "masks": {0: "batch", 1: "num_dets", 2: "height", 3: "width"},

diff --git a/src/otx/algo/instance_segmentation/rtmdet_inst.py b/src/otx/algo/instance_segmentation/rtmdet_inst.py
@@ -122,7 +122,7 @@ def _exporter(self) -> OTXModelExporter:
                 "input_names": ["image"],
                 "output_names": ["boxes", "labels", "masks"],
                 "dynamic_axes": {
-                    "image": {0: "batch", 2: "height", 3: "width"},
+                    "image": {0: "batch"},
                     "boxes": {0: "batch", 1: "num_dets"},
                     "labels": {0: "batch", 1: "num_dets"},
                     "masks": {0: "batch", 1: "num_dets", 2: "height", 3: "width"},