Add dumping onnx model to model entity

Author: sovrasov

otx/algorithms/action/tasks/inference.py

@@ -362,10 +362,13 @@ def export(
             exporter.export()
             bin_file = [f for f in os.listdir(self._output_path) if f.endswith(".bin")][0]
             xml_file = [f for f in os.listdir(self._output_path) if f.endswith(".xml")][0]
+            onnx_file = [f for f in os.listdir(self._output_path) if f.endswith(".onnx")][0]
             with open(os.path.join(self._output_path, bin_file), "rb") as f:
                 output_model.set_data("openvino.bin", f.read())
             with open(os.path.join(self._output_path, xml_file), "rb") as f:
                 output_model.set_data("openvino.xml", f.read())
+            with open(onnx_file, "rb") as file:
+                output_model.set_data("model.onnx", file.read())
             output_model.set_data(
                 "confidence_threshold", np.array([self.confidence_threshold], dtype=np.float32).tobytes()
             )

otx/algorithms/anomaly/tasks/inference.py

@@ -292,6 +292,8 @@ def export(
             output_model.set_data("openvino.bin", file.read())
         with open(xml_file, "rb") as file:
             output_model.set_data("openvino.xml", file.read())
+        with open(onnx_path, "rb") as file:
+            output_model.set_data("model.onnx", file.read())
 
         output_model.precision = self.precision
         output_model.optimization_methods = self.optimization_methods

otx/algorithms/classification/tasks/inference.py

@@ -249,19 +249,22 @@ def export(
 
         bin_file = outputs.get("bin")
         xml_file = outputs.get("xml")
+        onnx_file = outputs.get("onnx")
 
         inference_config = get_cls_inferencer_configuration(self._task_environment.label_schema)
         deploy_cfg = get_cls_deploy_config(self._task_environment.label_schema, inference_config)
         ir_extra_data = get_cls_model_api_configuration(self._task_environment.label_schema, inference_config)
         ir_extra_data[("otx_config",)] = json.dumps(deploy_cfg, ensure_ascii=False)
         embed_ir_model_data(xml_file, ir_extra_data)
 
-        if xml_file is None or bin_file is None:
-            raise RuntimeError("invalid status of exporting. bin and xml should not be None")
+        if xml_file is None or bin_file is None or onnx_file is None:
+            raise RuntimeError("invalid status of exporting. bin and xml or onnx should not be None")
         with open(bin_file, "rb") as f:
             output_model.set_data("openvino.bin", f.read())
         with open(xml_file, "rb") as f:
             output_model.set_data("openvino.xml", f.read())
+        with open(onnx_file, "rb") as f:
+            output_model.set_data("model.onnx", f.read())
         output_model.precision = self._precision
         output_model.has_xai = dump_features
         output_model.set_data(

otx/algorithms/common/adapters/mmcv/tasks/exporter_mixin.py

@@ -72,6 +72,7 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs):  # noqa: C901
             "outputs": {
                 "bin": os.path.join(cfg.work_dir, f"{model_name}.bin"),
                 "xml": os.path.join(cfg.work_dir, f"{model_name}.xml"),
+                "onnx": os.path.join(cfg.work_dir, f"{model_name}.onnx"),
                 "partitioned": [
                     {
                         "bin": os.path.join(cfg.work_dir, name.replace(".xml", ".bin")),

otx/algorithms/detection/adapters/mmdet/task.py

@@ -513,18 +513,21 @@ def export(
 
         bin_file = outputs.get("bin")
         xml_file = outputs.get("xml")
+        onnx_file = outputs.get("onnx")
 
         ir_extra_data = get_det_model_api_configuration(
             self._task_environment.label_schema, self._task_type, self.confidence_threshold
         )
         embed_ir_model_data(xml_file, ir_extra_data)
 
-        if xml_file is None or bin_file is None:
-            raise RuntimeError("invalid status of exporting. bin and xml should not be None")
+        if xml_file is None or bin_file is None or onnx_file is None:
+            raise RuntimeError("invalid status of exporting. bin and xml or onnx should not be None")
         with open(bin_file, "rb") as f:
             output_model.set_data("openvino.bin", f.read())
         with open(xml_file, "rb") as f:
             output_model.set_data("openvino.xml", f.read())
+        with open(onnx_file, "rb") as f:
+            output_model.set_data("model.onnx", f.read())
         output_model.set_data(
             "confidence_threshold",
             np.array([self.confidence_threshold], dtype=np.float32).tobytes(),

otx/algorithms/segmentation/tasks/inference.py

@@ -175,12 +175,16 @@ def export(
 
         bin_file = outputs.get("bin")
         xml_file = outputs.get("xml")
-        if xml_file is None or bin_file is None:
-            raise RuntimeError("invalid status of exporting. bin and xml should not be None")
+        onnx_file = outputs.get("onnx")
+
+        if xml_file is None or bin_file is None or onnx_file is None:
+            raise RuntimeError("invalid status of exporting. bin and xml or onnx should not be None")
         with open(bin_file, "rb") as f:
             output_model.set_data("openvino.bin", f.read())
         with open(xml_file, "rb") as f:
             output_model.set_data("openvino.xml", f.read())
+        with open(onnx_file, "rb") as f:
+            output_model.set_data("model.onnx", f.read())
         output_model.precision = self._precision
         output_model.optimization_methods = self._optimization_methods
         output_model.has_xai = dump_features

tests/test_suite/run_test_command.py

@@ -233,6 +233,7 @@ def otx_export_testing(template, root, dump_features=False, half_precision=False
     path_to_xml = os.path.join(save_path, "openvino.xml")
     assert os.path.exists(path_to_xml)
     assert os.path.exists(os.path.join(save_path, "openvino.bin"))
+    assert os.path.exists(os.path.join(save_path, "model.onnx"))
     assert os.path.exists(os.path.join(save_path, "label_schema.json"))
 
     if dump_features:

tests/unit/algorithms/classification/tasks/test_classification_inference.py

@@ -131,16 +131,25 @@ def test_export_with_model_files(self, mocker, precision: ModelPrecision, dump_f
             f.write(b"foo")
         with open(f"{self.output_path}/model.bin", "wb") as f:
             f.write(b"bar")
+        with open(f"{self.output_path}/model.onnx", "wb") as f:
+            f.write(b"foo")
 
         mocker.patch("otx.algorithms.classification.tasks.inference.embed_ir_model_data", return_value=None)
-        fake_output = {"outputs": {"bin": f"{self.output_path}/model.xml", "xml": f"{self.output_path}/model.bin"}}
+        fake_output = {
+            "outputs": {
+                "bin": f"{self.output_path}/model.xml",
+                "xml": f"{self.output_path}/model.bin",
+                "onnx": f"{self.output_path}/model.onnx",
+            }
+        }
         mock_run_task = mocker.patch.object(BaseTask, "_run_task", return_value=fake_output)
         self.cls_inference_task.export(ExportType.OPENVINO, self.model, precision, dump_features)
 
         mock_run_task.assert_called_once()
         assert self.model.has_xai == dump_features
         assert self.model.get_data("openvino.bin")
         assert self.model.get_data("openvino.xml")
+        assert self.model.get_data("model.onnx")
 
     @e2e_pytest_unit
     def test_unload(self, mocker):