diff --git a/src/otx/algo/segmentation/dino_v2_seg.py b/src/otx/algo/segmentation/dino_v2_seg.py index b0e946cf826..7bc4ee8ed6a 100644 --- a/src/otx/algo/segmentation/dino_v2_seg.py +++ b/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) diff --git a/src/otx/algo/segmentation/litehrnet.py b/src/otx/algo/segmentation/litehrnet.py index 118b0b8d3ed..ae6258100d4 100644 --- a/src/otx/algo/segmentation/litehrnet.py +++ b/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 @@ -16,3 +18,10 @@ def __init__(self, num_classes: int, variant: Literal["18", "s", "x"]) -> None: model_name = f"litehrnet_{variant}" config = read_mmconfig(model_name=model_name) 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, + } diff --git a/src/otx/algo/segmentation/mmconfigs/segnext_b.yaml b/src/otx/algo/segmentation/mmconfigs/segnext_b.yaml index 30de9e8c8cf..8843448f2a0 100644 --- a/src/otx/algo/segmentation/mmconfigs/segnext_b.yaml +++ b/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 diff --git a/src/otx/algo/segmentation/mmconfigs/segnext_s.yaml b/src/otx/algo/segmentation/mmconfigs/segnext_s.yaml index 8e1980585e4..80dfe9a59f3 100644 --- a/src/otx/algo/segmentation/mmconfigs/segnext_s.yaml +++ b/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 diff --git a/src/otx/algo/segmentation/mmconfigs/segnext_t.yaml b/src/otx/algo/segmentation/mmconfigs/segnext_t.yaml index a666165ad14..d6467f4fb64 100644 --- a/src/otx/algo/segmentation/mmconfigs/segnext_t.yaml +++ b/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 diff --git a/src/otx/algo/segmentation/segnext.py b/src/otx/algo/segmentation/segnext.py index f5e06740c15..acbf79cd07f 100644 --- a/src/otx/algo/segmentation/segnext.py +++ b/src/otx/algo/segmentation/segnext.py @@ -16,3 +16,7 @@ def __init__(self, num_classes: int, variant: Literal["b", "s", "t"]) -> None: model_name = f"segnext_{variant}" config = read_mmconfig(model_name=model_name) super().__init__(num_classes=num_classes, config=config) + + def _configure_export_parameters(self) -> None: + super()._configure_export_parameters() + self.export_params["via_onnx"] = True diff --git a/src/otx/core/model/entity/classification.py b/src/otx/core/model/entity/classification.py index 23a59f0cf66..8ea4cee1461 100644 --- a/src/otx/core/model/entity/classification.py +++ b/src/otx/core/model/entity/classification.py @@ -89,6 +89,7 @@ def __init__(self, num_classes: int, config: DictConfig) -> None: self.config = config self.export_params = _get_export_params_from_cls_mmconfig(config) self.load_from = config.pop("load_from", None) + self.image_size = (224, 224) super().__init__(num_classes=num_classes) def _create_model(self) -> nn.Module: @@ -164,7 +165,7 @@ def _configure_export_parameters(self) -> None: 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, *self.image_size) self.export_params["onnx_export_configuration"] = None def _create_exporter( @@ -284,7 +285,7 @@ def _configure_export_parameters(self) -> None: 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, *self.image_size) self.export_params["onnx_export_configuration"] = None def _create_exporter( @@ -404,7 +405,7 @@ def _configure_export_parameters(self) -> None: 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, *self.image_size) self.export_params["onnx_export_configuration"] = None def _create_exporter( diff --git a/src/otx/core/model/entity/segmentation.py b/src/otx/core/model/entity/segmentation.py index c115b910f65..1d64d254bd4 100644 --- a/src/otx/core/model/entity/segmentation.py +++ b/src/otx/core/model/entity/segmentation.py @@ -35,7 +35,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 @@ -60,6 +60,7 @@ def __init__(self, num_classes: int, config: DictConfig) -> None: self.config = config self.export_params = _get_export_params_from_seg_mmconfig(config) self.load_from = self.config.pop("load_from", None) + self.image_size = (544, 544) super().__init__(num_classes=num_classes) def _create_model(self) -> nn.Module: @@ -151,7 +152,7 @@ def _configure_export_parameters(self) -> None: 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, *self.image_size) self.export_params["onnx_export_configuration"] = None def _create_exporter( diff --git a/src/otx/engine/engine.py b/src/otx/engine/engine.py index 610db145292..27a8a31dc9d 100644 --- a/src/otx/engine/engine.py +++ b/src/otx/engine/engine.py @@ -329,6 +329,7 @@ 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,