From 738575252d3e26a1981393d86e85fc9dac452134 Mon Sep 17 00:00:00 2001
From: jaegukhyun <jaeguk.hyun@intel.com>
Date: Wed, 12 Apr 2023 10:55:51 +0900
Subject: [PATCH] Reflect changes from
 https://github.com/openvinotoolkit/training_extensions/pull/1976

---
 otx/algorithms/detection/task.py              |  7 +++--
 .../detection/adapters/mmdet/test_task.py     | 26 +++++++++----------
 2 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/otx/algorithms/detection/task.py b/otx/algorithms/detection/task.py
index f70ed11c468..ceab022c69a 100644
--- a/otx/algorithms/detection/task.py
+++ b/otx/algorithms/detection/task.py
@@ -274,18 +274,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(),
diff --git a/tests/unit/algorithms/detection/adapters/mmdet/test_task.py b/tests/unit/algorithms/detection/adapters/mmdet/test_task.py
index 336b2c89ddd..9645fd7c9e8 100644
--- a/tests/unit/algorithms/detection/adapters/mmdet/test_task.py
+++ b/tests/unit/algorithms/detection/adapters/mmdet/test_task.py
@@ -121,11 +121,14 @@ def run(self, *args, **kwargs):
             f.write(np.ndarray([0]))
         with open(os.path.join(self._output_path, "openvino.xml"), "wb") as f:
             f.write(np.ndarray([0]))
+        with open(os.path.join(self._output_path, "model.onnx"), "wb") as f:
+            f.write(np.ndarray([0]))
 
         return {
             "outputs": {
                 "bin": os.path.join(self._output_path, "openvino.bin"),
                 "xml": os.path.join(self._output_path, "openvino.xml"),
+                "onnx": os.path.join(self._output_path, "model.onnx"),
             }
         }
 
@@ -267,26 +270,27 @@ def test_infer(self, mocker) -> None:
             assert output.get_annotations()[-1].get_labels()[0].probability == 0.7
 
     @e2e_pytest_unit
-    def test_evaluate(self) -> None:
-        """Test evaluate function.
+    def test_det_evaluate(self) -> None:
+        """Test evaluate function for detection."""
 
-        <Steps>
-            1. Create model entity
-            2. Create result set entity
-            3. Run evaluate function with same dataset, this should give 100% accuracy
-            4. Run evaluate function with empty dataset, this should give 0% accuracy
-            5. Do 1 - 4 for action detection
-        """
         _config = ModelConfiguration(DetectionConfig(), self.det_label_schema)
         _model = ModelEntity(self.det_dataset, _config)
         resultset = ResultSetEntity(_model, self.det_dataset, self.det_dataset)
         self.det_task.evaluate(resultset)
         assert resultset.performance.score.value == 1.0
 
+    @e2e_pytest_unit
+    def test_det_evaluate_with_empty_annotations(self) -> None:
+        """Test evaluate function for detection with empty predictions."""
+
         resultset = ResultSetEntity(_model, self.det_dataset, self.det_dataset.with_empty_annotations())
         self.det_task.evaluate(resultset)
         assert resultset.performance.score.value == 0.0
 
+    @e2e_pytest_unit
+    def test_iseg_evaluate(self) -> None:
+        """Test evaluate function for instance segmentation."""
+
         _config = ModelConfiguration(DetectionConfig(), self.iseg_label_schema)
         _model = ModelEntity(self.iseg_dataset, _config)
         resultset = ResultSetEntity(_model, self.iseg_dataset, self.iseg_dataset)
@@ -315,10 +319,6 @@ def test_export(self, mocker, precision: ModelPrecision) -> None:
             return_value=True,
         )
 
-        with open(os.path.join(self.det_task._output_path, "openvino.xml"), "wb") as f:
-            f.write(np.ndarray([0]))
-        with open(os.path.join(self.det_task._output_path, "openvino.bin"), "wb") as f:
-            f.write(np.ndarray([0]))
         self.det_task.export(ExportType.OPENVINO, _model, precision, False)
 
         assert _model.model_format == ModelFormat.OPENVINO