Fix export issue in Geti (#2316)

Author: jaegukhyun

src/otx/algorithms/classification/adapters/mmcls/task.py

@@ -549,7 +549,18 @@ def hook(module, inp, outp):  # pylint: disable=unused-argument
         return eval_predictions, saliency_maps
 
     def _export_model(self, precision: ModelPrecision, export_format: ExportType, dump_features: bool):
-        self._data_cfg = None
+        self._data_cfg = ConfigDict(
+            data=ConfigDict(
+                train=ConfigDict(
+                    otx_dataset=None,
+                    labels=self._labels,
+                ),
+                test=ConfigDict(
+                    otx_dataset=None,
+                    labels=self._labels,
+                ),
+            )
+        )
         self._init_task(export=True)
 
         cfg = self.configure(False, "test", None)

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

@@ -491,7 +491,18 @@ def _export_model(
         dump_features: bool,
     ):
         """Main export function of OTX MMDetection Task."""
-        self._data_cfg = None
+        self._data_cfg = ConfigDict(
+            data=ConfigDict(
+                train=ConfigDict(
+                    otx_dataset=None,
+                    labels=self._labels,
+                ),
+                test=ConfigDict(
+                    otx_dataset=None,
+                    labels=self._labels,
+                ),
+            )
+        )
         self._init_task(export=True)
 
         cfg = self.configure(False, "test", None)

src/otx/algorithms/segmentation/adapters/mmseg/task.py

@@ -423,7 +423,18 @@ def _export_model(
     ):
         """Export function of OTX Segmentation Task."""
         # copied from OTX inference_task.py
-        self._data_cfg = None
+        self._data_cfg = ConfigDict(
+            data=ConfigDict(
+                train=ConfigDict(
+                    otx_dataset=None,
+                    labels=self._labels,
+                ),
+                test=ConfigDict(
+                    otx_dataset=None,
+                    labels=self._labels,
+                ),
+            )
+        )
         self._init_task(export=True)
 
         cfg = self.configure(False, "test", None)

tests/unit/algorithms/classification/adapters/mmcls/test_task.py

@@ -112,7 +112,8 @@ class MockExporter:
     def __init__(self, task):
         self._output_path = task._output_path
 
-    def run(self, *args, **kwargs):
+    def run(self, cfg, *args, **kwargs):
+        assert cfg.model.head.num_classes == 2
         with open(os.path.join(self._output_path, "openvino.bin"), "wb") as f:
             f.write(np.ndarray([0]))
         with open(os.path.join(self._output_path, "openvino.xml"), "wb") as f:

tests/unit/algorithms/detection/adapters/mmdet/test_task.py

@@ -114,7 +114,8 @@ class MockExporter:
     def __init__(self, task):
         self._output_path = task._output_path
 
-    def run(self, *args, **kwargs):
+    def run(self, cfg, *args, **kwargs):
+        cfg.model.bbox_head.num_classes == 3
         with open(os.path.join(self._output_path, "openvino.bin"), "wb") as f:
             f.write(np.ndarray([0]))
         with open(os.path.join(self._output_path, "openvino.xml"), "wb") as f: