Added segmentation configuration for export to align new interface

src/otx/algo/classification/deit_tiny.py

@@ -18,6 +18,7 @@ def __init__(self, num_classes: int, num_multiclass_heads: int, num_multilabel_c
         config = read_mmconfig(model_name="deit_tiny", subdir_name="hlabel_classification")
         config.head.num_multiclass_heads = num_multiclass_heads
         config.head.num_multilabel_classes = num_multilabel_classes
+        self.image_size = config["data_preprocessor"].get("size", (224, 224))
         super().__init__(num_classes=num_classes, config=config)
 
 
@@ -26,6 +27,7 @@ class DeitTinyForMulticlassCls(MMPretrainMulticlassClsModel):
 
     def __init__(self, num_classes: int) -> None:
         config = read_mmconfig("deit_tiny", subdir_name="multiclass_classification")
+        self.image_size = config["data_preprocessor"].get("size", (224, 224))
         super().__init__(num_classes=num_classes, config=config)
 
 
@@ -34,4 +36,5 @@ class DeitTinyForMultilabelCls(MMPretrainMultilabelClsModel):
 
     def __init__(self, num_classes: int) -> None:
         config = read_mmconfig("deit_tiny", subdir_name="multilabel_classification")
+        self.image_size = config["data_preprocessor"].get("size", (224, 224))
         super().__init__(num_classes=num_classes, config=config)

src/otx/algo/classification/efficientnet_b0.py

@@ -18,6 +18,7 @@ def __init__(self, num_classes: int, num_multiclass_heads: int, num_multilabel_c
         config = read_mmconfig(model_name="efficientnet_b0_light", subdir_name="hlabel_classification")
         config.head.num_multiclass_heads = num_multiclass_heads
         config.head.num_multilabel_classes = num_multilabel_classes
+        self.image_size = config["data_preprocessor"].get("size", (224, 224))
         super().__init__(num_classes=num_classes, config=config)
 
 
@@ -27,6 +28,7 @@ class EfficientNetB0ForMulticlassCls(MMPretrainMulticlassClsModel):
     def __init__(self, num_classes: int, light: bool = False) -> None:
         model_name = "efficientnet_b0_light" if light else "efficientnet_b0"
         config = read_mmconfig(model_name=model_name, subdir_name="multiclass_classification")
+        self.image_size = config["data_preprocessor"].get("size", (224, 224))
         super().__init__(num_classes=num_classes, config=config)
 
 
@@ -35,4 +37,5 @@ class EfficientNetB0ForMultilabelCls(MMPretrainMultilabelClsModel):
 
     def __init__(self, num_classes: int) -> None:
         config = read_mmconfig(model_name="efficientnet_b0_light", subdir_name="multilabel_classification")
+        self.image_size = config["data_preprocessor"].get("size", (224, 224))
         super().__init__(num_classes=num_classes, config=config)

src/otx/algo/classification/efficientnet_v2.py

@@ -18,6 +18,7 @@ def __init__(self, num_classes: int, num_multiclass_heads: int, num_multilabel_c
         config = read_mmconfig("efficientnet_v2_light", subdir_name="hlabel_classification")
         config.head.num_multiclass_heads = num_multiclass_heads
         config.head.num_multilabel_classes = num_multilabel_classes
+        self.image_size = config["data_preprocessor"].get("size", (224, 224))
         super().__init__(num_classes=num_classes, config=config)
 
 
@@ -27,6 +28,7 @@ class EfficientNetV2ForMulticlassCls(MMPretrainMulticlassClsModel):
     def __init__(self, num_classes: int, light: bool = False) -> None:
         model_name = "efficientnet_v2_light" if light else "efficientnet_v2"
         config = read_mmconfig(model_name=model_name, subdir_name="multiclass_classification")
+        self.image_size = config["data_preprocessor"].get("size", (224, 224))
         super().__init__(num_classes=num_classes, config=config)
 
 
@@ -35,4 +37,5 @@ class EfficientNetV2ForMultilabelCls(MMPretrainMultilabelClsModel):
 
     def __init__(self, num_classes: int) -> None:
         config = read_mmconfig("efficientnet_v2_light", subdir_name="multilabel_classification")
+        self.image_size = config["data_preprocessor"].get("size", (224, 224))
         super().__init__(num_classes=num_classes, config=config)

src/otx/algo/classification/mobilenet_v3_large.py

@@ -18,6 +18,7 @@ def __init__(self, num_classes: int, num_multiclass_heads: int, num_multilabel_c
         config = read_mmconfig(model_name="mobilenet_v3_large_light", subdir_name="hlabel_classification")
         config.head.num_multiclass_heads = num_multiclass_heads
         config.head.num_multilabel_classes = num_multilabel_classes
+        self.image_size = config["data_preprocessor"].get("size", (224, 224))
         super().__init__(num_classes=num_classes, config=config)
 
     def _configure_export_parameters(self) -> None:
@@ -31,6 +32,7 @@ class MobileNetV3ForMulticlassCls(MMPretrainMulticlassClsModel):
     def __init__(self, num_classes: int, light: bool = False) -> None:
         model_name = "mobilenet_v3_large_light" if light else "mobilenet_v3_large"
         config = read_mmconfig(model_name=model_name, subdir_name="multiclass_classification")
+        self.image_size = config["data_preprocessor"].get("size", (224, 224))
         super().__init__(num_classes=num_classes, config=config)
 
     def _configure_export_parameters(self) -> None:
@@ -43,6 +45,7 @@ class MobileNetV3ForMultilabelCls(MMPretrainMultilabelClsModel):
 
     def __init__(self, num_classes: int) -> None:
         config = read_mmconfig("mobilenet_v3_large_light", subdir_name="multilabel_classification")
+        self.image_size = config["data_preprocessor"].get("size", (224, 224))
         super().__init__(num_classes=num_classes, config=config)
 
     def _configure_export_parameters(self) -> None:

src/otx/algo/segmentation/dino_v2_seg.py

@@ -14,3 +14,12 @@ def __init__(self, num_classes: int) -> None:
         model_name = "dino_v2_seg"
         config = read_mmconfig(model_name=model_name)
         super().__init__(num_classes=num_classes, config=config)
+
+    def export(self, *args) -> None:
+        """Export method for DinoV2Seg.
+
+        Model doesn't support export for now due to unsupported operations from xformers.
+        This method will raise an error.
+        """
+        msg = "{model_name} cannot be exported. It is not supported."
+        raise RuntimeError(msg)

src/otx/algo/segmentation/litehrnet.py

@@ -5,6 +5,8 @@
 
 from typing import Literal
 
+from torch.onnx import OperatorExportTypes
+
 from otx.algo.utils.mmconfig import read_mmconfig
 from otx.core.model.entity.segmentation import MMSegCompatibleModel
 
@@ -15,4 +17,12 @@ class LiteHRNet(MMSegCompatibleModel):
     def __init__(self, num_classes: int, variant: Literal["18", "s", "x"]) -> None:
         model_name = f"litehrnet_{variant}"
         config = read_mmconfig(model_name=model_name)
+        self.image_size = config["data_preprocessor"]["size"]
         super().__init__(num_classes=num_classes, config=config)
+
+    def _configure_export_parameters(self) -> None:
+        super()._configure_export_parameters()
+        self.export_params["via_onnx"] = True
+        self.export_params["onnx_export_configuration"] = {
+            "operator_export_type": OperatorExportTypes.ONNX_ATEN_FALLBACK,
+        }

src/otx/algo/segmentation/mmconfigs/segnext_b.yaml

@@ -88,7 +88,7 @@ decode_head:
     requires_grad: true
     type: GN
   num_classes: 150
-  type: LightHamHead
+  type: CustomLightHamHead
 pretrained: null
 test_cfg:
   mode: whole

src/otx/algo/segmentation/mmconfigs/segnext_s.yaml

@@ -88,7 +88,7 @@ decode_head:
     requires_grad: true
     type: GN
   num_classes: 4
-  type: LightHamHead
+  type: CustomLightHamHead
 pretrained: null
 test_cfg:
   mode: whole

src/otx/algo/segmentation/mmconfigs/segnext_t.yaml

@@ -88,7 +88,7 @@ decode_head:
     requires_grad: true
     type: GN
   num_classes: 150
-  type: LightHamHead
+  type: CustomLightHamHead
 pretrained: null
 test_cfg:
   mode: whole

src/otx/algo/segmentation/segnext.py

@@ -15,4 +15,9 @@ class SegNext(MMSegCompatibleModel):
     def __init__(self, num_classes: int, variant: Literal["b", "s", "t"]) -> None:
         model_name = f"segnext_{variant}"
         config = read_mmconfig(model_name=model_name)
+        self.image_size = config["data_preprocessor"].get("size", (512, 512))
         super().__init__(num_classes=num_classes, config=config)
+
+    def _configure_export_parameters(self) -> None:
+        super()._configure_export_parameters()
+        self.export_params["via_onnx"] = True

src/otx/core/model/entity/classification.py

@@ -166,11 +166,12 @@ def _customize_outputs(
         )
 
     def _configure_export_parameters(self) -> None:
+        image_size = getattr(self, "image_size", (224, 224))
         self.export_params["resize_mode"] = "standard"
         self.export_params["pad_value"] = 0
         self.export_params["swap_rgb"] = False
         self.export_params["via_onnx"] = False
-        self.export_params["input_size"] = (1, 3, 224, 224)
+        self.export_params["input_size"] = (1, 3, *image_size)
         self.export_params["onnx_export_configuration"] = None
 
     def export(
@@ -300,11 +301,12 @@ def _customize_outputs(
         )
 
     def _configure_export_parameters(self) -> None:
+        image_size = getattr(self, "image_size", (224, 224))
         self.export_params["resize_mode"] = "standard"
         self.export_params["pad_value"] = 0
         self.export_params["swap_rgb"] = False
         self.export_params["via_onnx"] = False
-        self.export_params["input_size"] = (1, 3, 224, 224)
+        self.export_params["input_size"] = (1, 3, *image_size)
         self.export_params["onnx_export_configuration"] = None
 
     def export(
@@ -434,11 +436,12 @@ def _customize_outputs(
         )
 
     def _configure_export_parameters(self) -> None:
+        image_size = getattr(self, "image_size", (224, 224))
         self.export_params["resize_mode"] = "standard"
         self.export_params["pad_value"] = 0
         self.export_params["swap_rgb"] = False
         self.export_params["via_onnx"] = False
-        self.export_params["input_size"] = (1, 3, 224, 224)
+        self.export_params["input_size"] = (1, 3, *image_size)
         self.export_params["onnx_export_configuration"] = None
 
     def export(
@@ -494,7 +497,6 @@ def _customize_outputs(
         outputs: list[ClassificationResult],
         inputs: MultilabelClsBatchDataEntity,
     ) -> MultilabelClsBatchPredEntity:
-
         predicted_labels = []
         pred_scores = []
         for out in outputs:

src/otx/core/model/entity/segmentation.py

@@ -41,7 +41,7 @@ def _generate_model_metadata(
         metadata[("model_info", "task_type")] = "segmentation"
         metadata[("model_info", "return_soft_prediction")] = str(True)
         metadata[("model_info", "soft_threshold")] = str(0.5)
-        metadata[("model_info", "blur_strength")] = str(1)
+        metadata[("model_info", "blur_strength")] = str(-1)
 
         return metadata
 
@@ -153,11 +153,12 @@ def _customize_outputs(
         )
 
     def _configure_export_parameters(self) -> None:
+        image_size = getattr(self, "image_size", (512, 512))
         self.export_params["resize_mode"] = "standard"
         self.export_params["pad_value"] = 0
         self.export_params["swap_rgb"] = False
         self.export_params["via_onnx"] = False
-        self.export_params["input_size"] = (1, 3, 512, 512)
+        self.export_params["input_size"] = (1, 3, *image_size)
         self.export_params["onnx_export_configuration"] = None
 
     def export(

src/otx/engine/engine.py

@@ -329,10 +329,11 @@ def export(self, output_dir: Path, cfg: ExportConfig) -> None:
             output_dir (Path): Directory path to save exported binary files.
         """
         if self.checkpoint is not None:
+            self.model.eval()
             lit_module = self._build_lightning_module(
-            model=self.model,
-            optimizer=self.optimizer,
-            scheduler=self.scheduler,
+                model=self.model,
+                optimizer=self.optimizer,
+                scheduler=self.scheduler,
             )
             lit_module.meta_info = self.datamodule.meta_info