Add onnx to the resulting ModelEntity in export() (#1976)

* Fix fp16 onnx export for detection

Add dumping onnx model to model entity

Fix action tests

Fix segmentaion tests

* Update changelog

CHANGELOG.md

@@ -8,6 +8,7 @@ All notable changes to this project will be documented in this file.
 
 - Add generating feature cli_report.log in output for otx training (<https://github.com/openvinotoolkit/training_extensions/pull/1959>)
 - Support multiple python versions up to 3.10 (<https://github.com/openvinotoolkit/training_extensions/pull/1978>)
+- Support export of onnx models (<https://github.com/openvinotoolkit/training_extensions/pull/1976>)
 
 ### Enhancements
 

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(os.path.join(self._output_path, 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/task.py

@@ -250,19 +250,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.py

@@ -61,6 +61,7 @@ def run(self, 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/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/task.py

@@ -225,13 +225,16 @@ def export(
 
         bin_file = outputs.get("bin")
         xml_file = outputs.get("xml")
+        onnx_file = outputs.get("onnx")
 
-        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.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/action/tasks/test_action_inference.py

@@ -115,6 +115,8 @@ def export(self):
             f.write(dummy_data)
         with open(self.work_dir + ".xml", "wb") as f:
             f.write(dummy_data)
+        with open(self.work_dir + ".onnx", "wb") as f:
+            f.write(dummy_data)
 
 
 class TestActionInferenceTask:
@@ -298,6 +300,7 @@ def test_export(self, mocker, precision: ModelPrecision, dump_features: bool) ->
         assert _model.precision[0] == precision
         assert _model.get_data("openvino.bin") is not None
         assert _model.get_data("openvino.xml") is not None
+        assert _model.get_data("model.onnx") is not None
         assert _model.get_data("confidence_threshold") is not None
         assert _model.precision == self.cls_task._precision
         assert _model.optimization_methods == self.cls_task._optimization_methods

tests/unit/algorithms/segmentation/test_task.py

@@ -39,7 +39,7 @@ def _explain_model(*args, **kwargs):
         pass
 
     def _export_model(*args, **kwargs):
-        return {"outputs": {"bin": f"/tmp/model.xml", "xml": f"/tmp/model.bin"}}
+        return {"outputs": {"bin": f"/tmp/model.xml", "xml": f"/tmp/model.bin", "onnx": f"/tmp/model.onnx"}}
 
 
 class MockModel: