Add configuration to segmentation ir

openvinotoolkit · Apr 18, 2023 · 681ead6 · 681ead6
1 parent 6cbac46
commit 681ead6
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 2 deletions.
diff --git a/otx/algorithms/segmentation/adapters/mmseg/task.py b/otx/algorithms/segmentation/adapters/mmseg/task.py
@@ -57,7 +57,6 @@
 from otx.algorithms.segmentation.adapters.mmseg.utils.exporter import SegmentationExporter
 from otx.algorithms.segmentation.task import OTXSegmentationTask
 
-# from otx.algorithms.segmentation.utils import get_det_model_api_configuration
 from otx.api.configuration import cfg_helper
 from otx.api.configuration.helper.utils import ids_to_strings
 from otx.api.entities.datasets import DatasetEntity

diff --git a/otx/algorithms/segmentation/configs/base/configuration.py b/otx/algorithms/segmentation/configs/base/configuration.py
@@ -44,7 +44,7 @@
 class SegmentationConfig(BaseConfig):
     """Configurations of OTX Segmentation."""
 
-    header = string_attribute("Configuration for an object detection task of MPA")
+    header = string_attribute("Configuration for an object semantic segmentation task of OTX")
     description = header
 
     @attrs

diff --git a/otx/algorithms/segmentation/task.py b/otx/algorithms/segmentation/task.py
@@ -29,11 +29,13 @@
     InferenceProgressCallback,
     TrainingProgressCallback,
 )
+from otx.algorithms.common.utils.ir import embed_ir_model_data
 from otx.algorithms.common.utils.logger import get_logger
 from otx.algorithms.segmentation.adapters.openvino.model_wrappers.blur import (
     get_activation_map,
 )
 from otx.algorithms.segmentation.configs.base import SegmentationConfig
+from otx.algorithms.segmentation.utils.metadata import get_seg_model_api_configuration
 from otx.api.configuration import cfg_helper
 from otx.api.configuration.helper.utils import ids_to_strings
 from otx.api.entities.datasets import DatasetEntity
@@ -228,6 +230,11 @@ def export(
         xml_file = outputs.get("xml")
         onnx_file = outputs.get("onnx")
 
+        ir_extra_data = get_seg_model_api_configuration(
+            self._task_environment.label_schema, self._task_type, self._hyperparams
+        )
+        embed_ir_model_data(xml_file, ir_extra_data)
+
         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:

diff --git a/otx/algorithms/segmentation/utils/__init__.py b/otx/algorithms/segmentation/utils/__init__.py
@@ -13,3 +13,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions
 # and limitations under the License.
+
+
+from .metadata import get_seg_model_api_configuration