Introduce Classification Factory and Simplify Model Imports (#4456)

* add factory for classficaiton

* add mising files

* minor

* fix imports

* fix imports in tests 2

* fix ruff

* fix unit test

* update factory. Reply comments

* add literal to other backbones

Author: kprokofi

src/otx/backend/native/models/__init__.py

@@ -5,21 +5,11 @@
 
 from .anomaly import Padim, Stfpm, Uflow
 from .classification import (
-    EfficientNetHLabelCls,
-    EfficientNetMulticlassCls,
-    EfficientNetMultilabelCls,
-    MobileNetV3HLabelCls,
-    MobileNetV3MulticlassCls,
-    MobileNetV3MultilabelCls,
-    TimmModelHLabelCls,
-    TimmModelMulticlassCls,
-    TimmModelMultilabelCls,
-    TVModelHLabelCls,
-    TVModelMulticlassCls,
-    TVModelMultilabelCls,
-    VisionTransformerHLabelCls,
-    VisionTransformerMulticlassCls,
-    VisionTransformerMultilabelCls,
+    EfficientNet,
+    MobileNetV3,
+    TimmModel,
+    TVModel,
+    VisionTransformer,
 )
 from .detection import ATSS, RTDETR, SSD, DFine, RTMDet
 from .instance_segmentation import MaskRCNN, MaskRCNNTV, RTMDetInst
@@ -30,21 +20,11 @@
     "Padim",
     "Stfpm",
     "Uflow",
-    "EfficientNetHLabelCls",
-    "EfficientNetMulticlassCls",
-    "EfficientNetMultilabelCls",
-    "MobileNetV3HLabelCls",
-    "MobileNetV3MulticlassCls",
-    "MobileNetV3MultilabelCls",
-    "TimmModelHLabelCls",
-    "TimmModelMulticlassCls",
-    "TimmModelMultilabelCls",
-    "TVModelHLabelCls",
-    "TVModelMulticlassCls",
-    "TVModelMultilabelCls",
-    "VisionTransformerHLabelCls",
-    "VisionTransformerMulticlassCls",
-    "VisionTransformerMultilabelCls",
+    "EfficientNet",
+    "TimmModel",
+    "MobileNetV3",
+    "TVModel",
+    "VisionTransformer",
     "ATSS",
     "DFine",
     "SSD",

src/otx/backend/native/models/classification/__init__.py

@@ -3,42 +3,18 @@
 
 """Module for OTX classification models."""
 
-from .hlabel_models import (
-    EfficientNetHLabelCls,
-    MobileNetV3HLabelCls,
-    TimmModelHLabelCls,
-    TVModelHLabelCls,
-    VisionTransformerHLabelCls,
-)
-from .multiclass_models import (
-    EfficientNetMulticlassCls,
-    MobileNetV3MulticlassCls,
-    TimmModelMulticlassCls,
-    TVModelMulticlassCls,
-    VisionTransformerMulticlassCls,
-)
-from .multilabel_models import (
-    EfficientNetMultilabelCls,
-    MobileNetV3MultilabelCls,
-    TimmModelMultilabelCls,
-    TVModelMultilabelCls,
-    VisionTransformerMultilabelCls,
+from .factory import (
+    EfficientNet,
+    MobileNetV3,
+    TimmModel,
+    TVModel,
+    VisionTransformer,
 )
 
 __all__ = [
-    "EfficientNetMulticlassCls",
-    "TimmModelMulticlassCls",
-    "MobileNetV3MulticlassCls",
-    "TVModelMulticlassCls",
-    "VisionTransformerMulticlassCls",
-    "EfficientNetHLabelCls",
-    "TimmModelHLabelCls",
-    "MobileNetV3HLabelCls",
-    "TVModelHLabelCls",
-    "VisionTransformerHLabelCls",
-    "EfficientNetMultilabelCls",
-    "TimmModelMultilabelCls",
-    "MobileNetV3MultilabelCls",
-    "TVModelMultilabelCls",
-    "VisionTransformerMultilabelCls",
+    "EfficientNet",
+    "TimmModel",
+    "MobileNetV3",
+    "TVModel",
+    "VisionTransformer",
 ]

src/otx/backend/native/models/classification/backbones/__init__.py

@@ -7,6 +7,12 @@
 from .mobilenet_v3 import MobileNetV3Backbone
 from .timm import TimmBackbone
 from .torchvision import TorchvisionBackbone
-from .vision_transformer import VisionTransformer
+from .vision_transformer import VisionTransformerBackbone
 
-__all__ = ["EfficientNetBackbone", "TimmBackbone", "MobileNetV3Backbone", "VisionTransformer", "TorchvisionBackbone"]
+__all__ = [
+    "EfficientNetBackbone",
+    "TimmBackbone",
+    "MobileNetV3Backbone",
+    "VisionTransformerBackbone",
+    "TorchvisionBackbone",
+]

src/otx/backend/native/models/classification/backbones/efficientnet.py

@@ -404,8 +404,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.conv(x)
 
 
-class EfficientNet(nn.Module):
-    """EfficientNet.
+class EfficientNetFeatureExtractor(nn.Module):
+    """Implementation of the EfficientNet Feature Extactor.
 
     Args:
         channels : list of list of int. Number of output channels for each unit.
@@ -611,7 +611,7 @@ def __new__(
         input_size: tuple[int, int] | None = None,
         pretrained: bool = True,
         **kwargs,
-    ) -> EfficientNet:
+    ) -> EfficientNetFeatureExtractor:
         """Create a new instance of the EfficientNet class.
 
         Args:
@@ -621,7 +621,7 @@ def __new__(
             **kwargs: Additional keyword arguments to be passed to the EfficientNet constructor.
 
         Returns:
-            EfficientNet: The created EfficientNet model instance.
+            EfficientNetFeatureExtractor: The created EfficientNetFeatureExtractor model instance.
         """
         origin_input_size, depth_factor, width_factor = cls.EFFICIENTNET_CFG[model_name].values()
         input_size = input_size or origin_input_size
@@ -657,7 +657,7 @@ def __new__(
         if width_factor > 1.0:
             final_block_channels = round_channels(final_block_channels * width_factor)
 
-        model = EfficientNet(
+        model = EfficientNetFeatureExtractor(
             channels=channels,
             init_block_channels=init_block_channels,
             final_block_channels=final_block_channels,

src/otx/backend/native/models/classification/backbones/mobilenet_v3.py

@@ -268,8 +268,8 @@ def forward(
         return self.extract_features(x)
 
 
-class MobileNetV3(MobileNetV3Base):
-    """MobileNetV3 constructor.
+class MobileNetV3FeatureExtractor(MobileNetV3Base):
+    """MobileNetV3FeatureExtractor constructor.
 
     Args:
         layer_cfgs (list): List of layer configurations.
@@ -396,7 +396,7 @@ def __new__(
         width_mult: float = 1.0,
         pretrained: bool = True,
         **kwargs,
-    ) -> MobileNetV3:
+    ) -> MobileNetV3FeatureExtractor:
         """Create a new instance of the MobileNetV3 class.
 
         Args:
@@ -412,7 +412,7 @@ def __new__(
             msg = f"Unknown MobileNetV3 model: {model_name}"
             raise ValueError(msg)
 
-        model = MobileNetV3(
+        model = MobileNetV3FeatureExtractor(
             layer_cfgs=cls.MV3_CFG[model_name]["layer_cfgs"],
             width_mult=width_mult,
             **kwargs,

src/otx/backend/native/models/classification/backbones/vision_transformer.py

@@ -35,7 +35,7 @@
     import numpy as np
 
 
-class VisionTransformer(BaseModule):
+class VisionTransformerBackbone(BaseModule):
     """Implementation of Vision Transformer from Timm.
 
     A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`
@@ -553,7 +553,7 @@ def forward(
     @torch.no_grad()
     def _load_npz_weights(  # noqa: C901
         self,
-        model: VisionTransformer,
+        model: VisionTransformerBackbone,
         checkpoint_path: str,
         prefix: str = "",
     ) -> None: