Add mmdeploy export

pyproject.toml

@@ -83,8 +83,8 @@ mmlab = [
     "mmpretrain==1.1.1",
     "mmsegmentation==1.2.1",
     "mmaction2==1.2.0",
+    "mmdeploy==1.3.*",
     # Without the pip cache, oss2 will sometimes install to a lower version. This is related to the installation of the mmlab library.
-    # This causes an error when training the mm model, so fix the version first.
     "oss2==2.17.0",
 ]
 
@@ -249,6 +249,7 @@ exclude = [
     "node_modules",
     "venv",
     "tests/assets",
+    "src/otx/algo/**/mmdeploy/",
 
     # it will be cleaned up later
     "src/otx/algo/classification/backbones/*",

src/otx/algo/detection/atss.py

@@ -3,7 +3,9 @@
 #
 """ATSS model implementations."""
 
-from typing import Literal
+from __future__ import annotations
+
+from typing import Any, Literal
 
 from otx.algo.utils.mmconfig import read_mmconfig
 from otx.algo.utils.support_otx_v1 import OTXv1Helper
@@ -17,7 +19,40 @@ def __init__(self, num_classes: int, variant: Literal["mobilenetv2", "r50_fpn",
         model_name = f"atss_{variant}"
         config = read_mmconfig(model_name=model_name)
         super().__init__(num_classes=num_classes, config=config)
+        self.image_size = (1, 3, 736, 992)
+
+    @property
+    def _export_parameters(self) -> dict[str, Any]:
+        """Parameters for an exporter."""
+        export_params = super()._export_parameters
+        export_params["deploy_cfg"] = "otx.algo.detection.mmdeploy.atss"
+        export_params["input_size"] = self.image_size
+        export_params["resize_mode"] = "standard"
+        export_params["pad_value"] = 0
+        export_params["swap_rgb"] = False
+
+        return export_params
 
     def load_from_otx_v1_ckpt(self, state_dict: dict, add_prefix: str = "model.model.") -> dict:
         """Load the previous OTX ckpt according to OTX2.0."""
         return OTXv1Helper.load_det_ckpt(state_dict, add_prefix)
+
+
+class ATSSR50FPN(ATSS):
+    """ATSSR50FPN Model."""
+
+    def __init__(self, num_classes: int) -> None:
+        super().__init__(num_classes=num_classes, variant="r50_fpn")
+        self.image_size = (1, 3, 800, 1333)
+
+    @property
+    def _export_parameters(self) -> dict[str, Any]:
+        """Parameters for an exporter."""
+        export_params = super()._export_parameters
+        export_params["deploy_cfg"] = "otx.algo.detection.mmdeploy.atss_r50_fpn"
+        export_params["input_size"] = self.image_size
+        export_params["resize_mode"] = "standard"  # [TODO](@Eunwoo): need to revert it to fit_to_window after resolving
+        export_params["pad_value"] = 0
+        export_params["swap_rgb"] = False
+
+        return export_params

src/otx/algo/detection/mmdeploy/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+#
+"""MMDeploy config for detection models."""

src/otx/algo/detection/mmdeploy/atss.py

@@ -0,0 +1,11 @@
+"""MMDeploy config of ATSS model for Detection Task."""
+
+_base_ = ["./base_detection.py"]
+
+ir_config = dict(
+    output_names=["boxes", "labels"],
+)
+
+backend_config = dict(
+    model_inputs=[dict(opt_shapes=dict(input=[-1, 3, 736, 992]))],
+)

src/otx/algo/detection/mmdeploy/atss_r50_fpn.py

@@ -0,0 +1,11 @@
+"""MMDeploy config of ATSS model for Detection Task."""
+
+_base_ = ["./base_detection.py"]
+
+ir_config = dict(
+    output_names=["boxes", "labels"],
+)
+
+backend_config = dict(
+    model_inputs=[dict(opt_shapes=dict(input=[-1, 3, 800, 1333]))],
+)

src/otx/algo/detection/mmdeploy/base_detection.py

@@ -0,0 +1,51 @@
+"""Detection models base deploy config."""
+
+ir_config = dict(
+    type="onnx",
+    export_params=True,
+    keep_initializers_as_inputs=False,
+    opset_version=11,
+    save_file="end2end.onnx",
+    input_names=["image"],
+    output_names=["boxes", "labels"],
+    input_shape=None,
+    # TODO
+    # optimizing onnx graph mess up NNCF graph at some point
+    # where we need to look into
+    optimize=False,
+    dynamic_axes={
+        "image": {
+            0: "batch",
+            2: "height",
+            3: "width",
+        },
+        "boxes": {
+            0: "batch",
+            1: "num_dets",
+        },
+        "labels": {
+            0: "batch",
+            1: "num_dets",
+        },
+    },
+)
+
+codebase_config = dict(
+    type="mmdet",
+    task="ObjectDetection",
+    model_type="end2end",
+    post_processing=dict(
+        score_threshold=0.05,
+        confidence_threshold=0.005,  # for YOLOv3
+        iou_threshold=0.5,
+        max_output_boxes_per_class=200,
+        pre_top_k=5000,
+        keep_top_k=100,
+        background_label_id=-1,
+    ),
+)
+
+backend_config = dict(
+    type="openvino",
+    mo_options=None,
+)

src/otx/algo/detection/mmdeploy/rtmdet.py

@@ -0,0 +1,11 @@
+"""MMDeploy config of RTMdet model for Detection Task."""
+
+_base_ = ["./base_detection.py"]
+
+ir_config = dict(
+    output_names=["boxes", "labels"],
+)
+
+backend_config = dict(
+    model_inputs=[dict(opt_shapes=dict(input=[-1, 3, 640, 640]))],
+)

src/otx/algo/detection/mmdeploy/ssd_mobilenetv2.py

@@ -0,0 +1,11 @@
+"""MMDeploy config of SSD model for Detection Task."""
+
+_base_ = ["./base_detection.py"]
+
+ir_config = dict(
+    output_names=["boxes", "labels"],
+)
+
+backend_config = dict(
+    model_inputs=[dict(opt_shapes=dict(input=[-1, 3, 864, 864]))],
+)

src/otx/algo/detection/mmdeploy/yolox.py

@@ -0,0 +1,11 @@
+"""MMDeploy config of YOLOX models except YOLOX_tiny for Detection Task."""
+
+_base_ = ["./base_detection.py"]
+
+ir_config = dict(
+    output_names=["boxes", "labels"],
+)
+
+backend_config = dict(
+    model_inputs=[dict(opt_shapes=dict(input=[-1, 3, 640, 640]))],
+)

src/otx/algo/detection/mmdeploy/yolox_tiny.py

@@ -0,0 +1,11 @@
+"""MMDeploy config of YOLOX Tiny model for Detection Task."""
+
+_base_ = ["./base_detection.py"]
+
+ir_config = dict(
+    output_names=["boxes", "labels"],
+)
+
+backend_config = dict(
+    model_inputs=[dict(opt_shapes=dict(input=[-1, 3, 416, 416]))],
+)

src/otx/algo/detection/rtmdet.py

@@ -3,7 +3,9 @@
 #
 """RTMDetTiny model implementations."""
 
-from typing import Literal
+from __future__ import annotations
+
+from typing import Any, Literal
 
 from otx.algo.utils.mmconfig import read_mmconfig
 from otx.algo.utils.support_otx_v1 import OTXv1Helper
@@ -17,6 +19,19 @@ def __init__(self, num_classes: int, variant: Literal["tiny"]) -> None:
         model_name = f"rtmdet_{variant}"
         config = read_mmconfig(model_name=model_name)
         super().__init__(num_classes=num_classes, config=config)
+        self.image_size = (1, 3, 640, 640)
+
+    @property
+    def _export_parameters(self) -> dict[str, Any]:
+        """Parameters for an exporter."""
+        export_params = super()._export_parameters
+        export_params["deploy_cfg"] = "otx.algo.detection.mmdeploy.rtmdet"
+        export_params["input_size"] = self.image_size
+        export_params["resize_mode"] = "fit_to_window_letterbox"
+        export_params["pad_value"] = 114
+        export_params["swap_rgb"] = False
+
+        return export_params
 
     def load_from_otx_v1_ckpt(self, state_dict: dict, add_prefix: str = "model.model.") -> dict:
         """Load the previous OTX ckpt according to OTX2.0."""

src/otx/algo/detection/ssd.py

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 from copy import deepcopy
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 from otx.algo.utils.mmconfig import read_mmconfig
 from otx.algo.utils.support_otx_v1 import OTXv1Helper
@@ -27,6 +27,7 @@ def __init__(self, num_classes: int, variant: Literal["mobilenetv2"]) -> None:
         model_name = f"ssd_{variant}"
         config = read_mmconfig(model_name=model_name)
         super().__init__(num_classes=num_classes, config=config)
+        self.image_size = (1, 3, 864, 864)
 
     def _create_model(self) -> nn.Module:
         from mmdet.models.data_preprocessors import (
@@ -125,6 +126,18 @@ def load_state_dict_pre_hook(self, state_dict: dict[str, torch.Tensor], prefix:
             # Replace checkpoint weight by mixed weights
             state_dict[prefix + param_name] = model_param
 
+    @property
+    def _export_parameters(self) -> dict[str, Any]:
+        """Parameters for an exporter."""
+        export_params = super()._export_parameters
+        export_params["deploy_cfg"] = "otx.algo.detection.mmdeploy.ssd_mobilenetv2"
+        export_params["input_size"] = self.image_size
+        export_params["resize_mode"] = "standard"
+        export_params["pad_value"] = 0
+        export_params["swap_rgb"] = False
+
+        return export_params
+
     def load_from_otx_v1_ckpt(self, state_dict: dict, add_prefix: str = "model.model.") -> dict:
         """Load the previous OTX ckpt according to OTX2.0."""
         return OTXv1Helper.load_det_ckpt(state_dict, add_prefix)

src/otx/algo/detection/yolox.py

@@ -1,9 +1,11 @@
-# Copyright (C) 2023 Intel Corporation
+# Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 #
 """YOLOX model implementations."""
 
-from typing import Literal
+from __future__ import annotations
+
+from typing import Any, Literal
 
 from otx.algo.utils.mmconfig import read_mmconfig
 from otx.algo.utils.support_otx_v1 import OTXv1Helper
@@ -17,7 +19,40 @@ def __init__(self, num_classes: int, variant: Literal["l", "s", "tiny", "x"]) ->
         model_name = f"yolox_{variant}"
         config = read_mmconfig(model_name=model_name)
         super().__init__(num_classes=num_classes, config=config)
+        self.image_size = (1, 3, 640, 640)
+
+    @property
+    def _export_parameters(self) -> dict[str, Any]:
+        """Parameters for an exporter."""
+        export_params = super()._export_parameters
+        export_params["resize_mode"] = "fit_to_window_letterbox"
+        export_params["pad_value"] = 114
+        export_params["swap_rgb"] = True
+        export_params["input_size"] = self.image_size
+        export_params["deploy_cfg"] = "otx.algo.detection.mmdeploy.yolox"
+
+        return export_params
 
     def load_from_otx_v1_ckpt(self, state_dict: dict, add_prefix: str = "model.model.") -> dict:
         """Load the previous OTX ckpt according to OTX2.0."""
         return OTXv1Helper.load_det_ckpt(state_dict, add_prefix)
+
+
+class YoloXTiny(YoloX):
+    """YoloX tiny Model."""
+
+    def __init__(self, num_classes: int) -> None:
+        super().__init__(num_classes=num_classes, variant="tiny")
+        self.image_size = (1, 3, 416, 416)
+
+    @property
+    def _export_parameters(self) -> dict[str, Any]:
+        """Parameters for an exporter."""
+        export_params = super()._export_parameters
+        export_params["resize_mode"] = "fit_to_window_letterbox"
+        export_params["pad_value"] = 114
+        export_params["swap_rgb"] = False
+        export_params["input_size"] = self.image_size
+        export_params["deploy_cfg"] = "otx.algo.detection.mmdeploy.yolox_tiny"
+
+        return export_params

src/otx/algo/instance_segmentation/maskrcnn.py

@@ -1,9 +1,11 @@
-# Copyright (C) 2023 Intel Corporation
+# Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 #
-"""ATSS model implementations."""
+"""MaskRCNN model implementations."""
 
-from typing import Literal
+from __future__ import annotations
+
+from typing import Any, Literal
 
 from otx.algo.utils.mmconfig import read_mmconfig
 from otx.algo.utils.support_otx_v1 import OTXv1Helper
@@ -17,7 +19,40 @@ def __init__(self, num_classes: int, variant: Literal["efficientnetb2b", "r50",
         model_name = f"maskrcnn_{variant}"
         config = read_mmconfig(model_name=model_name)
         super().__init__(num_classes=num_classes, config=config)
+        self.image_size = (1, 3, 1024, 1024)
+
+    @property
+    def _export_parameters(self) -> dict[str, Any]:
+        """Parameters for an exporter."""
+        export_params = super()._export_parameters
+        export_params["deploy_cfg"] = "otx.algo.instance_segmentation.mmdeploy.maskrcnn"
+        export_params["input_size"] = self.image_size
+        export_params["resize_mode"] = "standard"  # [TODO](@Eunwoo): need to revert it to fit_to_window after resolving
+        export_params["pad_value"] = 0
+        export_params["swap_rgb"] = False
+
+        return export_params
 
     def load_from_otx_v1_ckpt(self, state_dict: dict, add_prefix: str = "model.model.") -> dict:
         """Load the previous OTX ckpt according to OTX2.0."""
         return OTXv1Helper.load_iseg_ckpt(state_dict, add_prefix)
+
+
+class MaskRCNNSwinT(MaskRCNN):
+    """MaskRCNNSwinT Model."""
+
+    def __init__(self, num_classes: int) -> None:
+        super().__init__(num_classes=num_classes, variant="swint")
+        self.image_size = (1, 3, 1344, 1344)
+
+    @property
+    def _export_parameters(self) -> dict[str, Any]:
+        """Parameters for an exporter."""
+        export_params = super()._export_parameters
+        export_params["deploy_cfg"] = "otx.algo.instance_segmentation.mmdeploy.maskrcnn_swint"
+        export_params["input_size"] = self.image_size
+        export_params["resize_mode"] = "standard"  # [TODO](@Eunwoo): need to revert it to fit_to_window after resolving
+        export_params["pad_value"] = 0
+        export_params["swap_rgb"] = False
+
+        return export_params

src/otx/algo/instance_segmentation/mmdeploy/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+#
+"""MMDeploy config for instance segmentation models."""