Reflect changes from openvinotoolkit#1976

Author: jaegukhyun

otx/algorithms/action/adapters/mmaction/task.py

@@ -421,10 +421,12 @@ def _export_model(self, precision: ModelPrecision, dump_features: bool = True):
         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]
         results = {
             "outputs": {
                 "bin": os.path.join(self._output_path, bin_file),
                 "xml": os.path.join(self._output_path, xml_file),
+                "onnx": os.path.join(self._output_path, onnx_file),
             }
         }
         return results

otx/algorithms/action/task.py

@@ -255,13 +255,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:
+        if xml_file is None or bin_file is None or onnx_file is None:
             raise RuntimeError("invalid status of exporting. bin and xml 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(),

tests/unit/algorithms/action/adapters/mmaction/test_task.py

@@ -116,14 +116,16 @@ def __iter__(self):
 class MockExporter:
     """Mock class for Exporter."""
 
-    def __init__(self, recipe_cfg, weights, deploy_cfg, work_dir, half_precision):
-        self.work_dir = work_dir
+    def __init__(self, task):
+        self.work_dir = task._output_path
 
     def export(self):
         dummy_data = np.ndarray((1, 1, 1))
-        with open(self.work_dir + ".bin", "wb") as f:
+        with open(os.path.join(self.work_dir, "openvino.bin"), "wb") as f:
             f.write(dummy_data)
-        with open(self.work_dir + ".xml", "wb") as f:
+        with open(os.path.join(self.work_dir, "openvino.xml"), "wb") as f:
+            f.write(dummy_data)
+        with open(os.path.join(self.work_dir, "model.onnx"), "wb") as f:
             f.write(dummy_data)
 
 
@@ -331,17 +333,10 @@ def test_export(self, mocker, precision: ModelPrecision) -> None:
         """
         _config = ModelConfiguration(ActionConfig(), self.cls_label_schema)
         _model = ModelEntity(self.cls_dataset, _config)
-        mocker.patch("otx.algorithms.action.adapters.mmaction.utils.Exporter", side_effect=MockExporter)
-        mocker.patch("otx.algorithms.action.adapters.mmaction.utils.export_utils.export", return_value=True)
-        mocker.patch("otx.algorithms.action.adapters.mmaction.utils.export_utils.from_onnx", return_value=True)
+        mocker.patch("otx.algorithms.action.adapters.mmaction.task.Exporter", return_value=MockExporter(self.cls_task))
         mocker.patch("torch.load", return_value={})
         mocker.patch("torch.nn.Module.load_state_dict", return_value=True)
-        mocker.patch("os.listdir", return_value=["openvino.xml", "openvino.bin"])
 
-        with open(os.path.join(self.cls_task._output_path, "openvino.xml"), "wb") as f:
-            f.write(np.ndarray([0]))
-        with open(os.path.join(self.cls_task._output_path, "openvino.bin"), "wb") as f:
-            f.write(np.ndarray([0]))
         self.cls_task.export(ExportType.OPENVINO, _model, precision, False)
 
         assert _model.model_format == ModelFormat.OPENVINO