Move assigning tasks to Models from Engine to Anomaly Model Classes (#…

…3683)

* Move Task assign into Model with Anomaly Task

* Fix openvino model class

src/otx/algo/anomaly/openvino_model.py

@@ -11,7 +11,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any, Literal, Sequence
 
 import numpy as np
 import openvino
@@ -26,6 +26,7 @@
 from otx.core.model.anomaly import AnomalyModelInputs
 from otx.core.model.base import OVModel
 from otx.core.types.label import AnomalyLabelInfo
+from otx.core.types.task import OTXTaskType
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -103,6 +104,11 @@ def __init__(
         use_throughput_mode: bool = True,
         model_api_configuration: dict[str, Any] | None = None,
         metric: MetricCallable = NullMetricCallable,  # Metrics is computed using Anomalib's metric
+        task: Literal[
+            OTXTaskType.ANOMALY_CLASSIFICATION,
+            OTXTaskType.ANOMALY_DETECTION,
+            OTXTaskType.ANOMALY_SEGMENTATION,
+        ] = OTXTaskType.ANOMALY_CLASSIFICATION,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -117,6 +123,7 @@ def __init__(
         metric_names = ["AUROC", "F1Score"]
         self.image_metrics: AnomalibMetricCollection = create_metric_collection(metric_names, prefix="image_")
         self.pixel_metrics: AnomalibMetricCollection = create_metric_collection(metric_names, prefix="pixel_")
+        self.task = task
 
     def _create_model(self) -> Model:
         from model_api.adapters import OpenvinoAdapter, create_core, get_user_config

src/otx/algo/anomaly/padim.py

@@ -7,13 +7,14 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
 from anomalib.models.image import Padim as AnomalibPadim
 
 from otx.core.model.anomaly import OTXAnomaly
 from otx.core.model.base import OTXModel
 from otx.core.types.label import AnomalyLabelInfo
+from otx.core.types.task import OTXTaskType
 
 if TYPE_CHECKING:
     from lightning.pytorch.utilities.types import STEP_OUTPUT
@@ -30,8 +31,9 @@ class Padim(OTXAnomaly, OTXModel, AnomalibPadim):
         layers (list[str], optional): Feature extractor layers. Defaults to ["layer1", "layer2", "layer3"].
         pre_trained (bool, optional): Pretrained backbone. Defaults to True.
         n_features (int | None, optional): Number of features. Defaults to None.
-        num_classes (int, optional): Anoamly don't use num_classes ,
-            but OTXModel always receives num_classes, so need this.
+        task (Literal[
+                OTXTaskType.ANOMALY_CLASSIFICATION, OTXTaskType.ANOMALY_DETECTION, OTXTaskType.ANOMALY_SEGMENTATION
+            ], optional): Task type of Anomaly Task. Defaults to OTXTaskType.ANOMALY_CLASSIFICATION.
     """
 
     def __init__(
@@ -40,6 +42,11 @@ def __init__(
         layers: list[str] = ["layer1", "layer2", "layer3"],  # noqa: B006
         pre_trained: bool = True,
         n_features: int | None = None,
+        task: Literal[
+            OTXTaskType.ANOMALY_CLASSIFICATION,
+            OTXTaskType.ANOMALY_DETECTION,
+            OTXTaskType.ANOMALY_SEGMENTATION,
+        ] = OTXTaskType.ANOMALY_CLASSIFICATION,
     ) -> None:
         OTXAnomaly.__init__(self)
         OTXModel.__init__(self, label_info=AnomalyLabelInfo())
@@ -50,6 +57,7 @@ def __init__(
             pre_trained=pre_trained,
             n_features=n_features,
         )
+        self.task = task
 
     def configure_optimizers(self) -> tuple[list[Optimizer], list[Optimizer]] | None:
         """PADIM doesn't require optimization, therefore returns no optimizers."""

src/otx/algo/anomaly/stfpm.py

@@ -7,13 +7,14 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING, Literal, Sequence
 
 from anomalib.models.image.stfpm import Stfpm as AnomalibStfpm
 
 from otx.core.model.anomaly import OTXAnomaly
 from otx.core.model.base import OTXModel
 from otx.core.types.label import AnomalyLabelInfo
+from otx.core.types.task import OTXTaskType
 
 if TYPE_CHECKING:
     from lightning.pytorch.utilities.types import STEP_OUTPUT
@@ -28,14 +29,20 @@ class Stfpm(OTXAnomaly, OTXModel, AnomalibStfpm):
     Args:
         layers (Sequence[str]): Feature extractor layers.
         backbone (str, optional): Feature extractor backbone. Defaults to "resnet18".
-        num_classes (int, optional): Anoamly don't use num_classes ,
-            but OTXModel always receives num_classes, so need this.
+        task (Literal[
+                OTXTaskType.ANOMALY_CLASSIFICATION, OTXTaskType.ANOMALY_DETECTION, OTXTaskType.ANOMALY_SEGMENTATION
+            ], optional): Task type of Anomaly Task. Defaults to OTXTaskType.ANOMALY_CLASSIFICATION.
     """
 
     def __init__(
         self,
         layers: Sequence[str] = ["layer1", "layer2", "layer3"],
         backbone: str = "resnet18",
+        task: Literal[
+            OTXTaskType.ANOMALY_CLASSIFICATION,
+            OTXTaskType.ANOMALY_DETECTION,
+            OTXTaskType.ANOMALY_SEGMENTATION,
+        ] = OTXTaskType.ANOMALY_CLASSIFICATION,
         **kwargs,
     ) -> None:
         OTXAnomaly.__init__(self)
@@ -45,6 +52,7 @@ def __init__(
             backbone=backbone,
             layers=layers,
         )
+        self.task = task
 
     @property
     def trainable_model(self) -> str:

src/otx/engine/engine.py

@@ -157,14 +157,6 @@ def __init__(
             )
         )
 
-        # [TODO](ashwinvaidya17): Need to revisit how task, optimizer, and scheduler are assigned to the model
-        if self.task in (
-            OTXTaskType.ANOMALY_CLASSIFICATION,
-            OTXTaskType.ANOMALY_DETECTION,
-            OTXTaskType.ANOMALY_SEGMENTATION,
-        ):
-            self._model = self._get_anomaly_model(self._model)
-
     # ------------------------------------------------------------------------ #
     # General OTX Entry Points
     # ------------------------------------------------------------------------ #
@@ -1029,8 +1021,3 @@ def datamodule(self) -> OTXDataModule:
             msg = "Please include the `data_root` or `datamodule` when creating the Engine."
             raise RuntimeError(msg)
         return self._datamodule
-
-    def _get_anomaly_model(self, model: OTXModel) -> OTXModel:
-        # [TODO](ashwinvaidya17): Need to revisit how task, optimizer, and scheduler are assigned to the model
-        model.task = self.task
-        return model

src/otx/recipe/anomaly_classification/openvino_model.yaml

@@ -4,6 +4,7 @@ model:
     model_name: openvino.xml
     async_inference: True
     use_throughput_mode: False
+    task: ANOMALY_CLASSIFICATION
 
 engine:
   task: ANOMALY_CLASSIFICATION

src/otx/recipe/anomaly_classification/padim.yaml

@@ -5,6 +5,7 @@ model:
     backbone: "resnet18"
     pre_trained: True
     n_features: null
+    task: ANOMALY_CLASSIFICATION
 
 engine:
   task: ANOMALY_CLASSIFICATION

src/otx/recipe/anomaly_classification/stfpm.yaml

@@ -3,6 +3,7 @@ model:
   init_args:
     layers: ["layer1", "layer2", "layer3"]
     backbone: "resnet18"
+    task: ANOMALY_CLASSIFICATION
 
     optimizer:
       class_path: torch.optim.SGD

src/otx/recipe/anomaly_detection/openvino_model.yaml

@@ -4,6 +4,7 @@ model:
     model_name: openvino.xml
     async_inference: True
     use_throughput_mode: False
+    task: ANOMALY_DETECTION
 
 engine:
   task: ANOMALY_DETECTION

src/otx/recipe/anomaly_detection/padim.yaml

@@ -5,6 +5,7 @@ model:
     backbone: "resnet18"
     pre_trained: True
     n_features: null
+    task: ANOMALY_DETECTION
 
 engine:
   task: ANOMALY_DETECTION

src/otx/recipe/anomaly_detection/stfpm.yaml

@@ -3,6 +3,7 @@ model:
   init_args:
     layers: ["layer1", "layer2", "layer3"]
     backbone: "resnet18"
+    task: ANOMALY_DETECTION
 
     optimizer:
       class_path: torch.optim.SGD

src/otx/recipe/anomaly_segmentation/openvino_model.yaml

@@ -4,6 +4,7 @@ model:
     model_name: openvino.xml
     async_inference: True
     use_throughput_mode: False
+    task: ANOMALY_SEGMENTATION
 
 engine:
   task: ANOMALY_SEGMENTATION

src/otx/recipe/anomaly_segmentation/padim.yaml

@@ -5,6 +5,7 @@ model:
     backbone: "resnet18"
     pre_trained: True
     n_features: null
+    task: ANOMALY_SEGMENTATION
 
 engine:
   task: ANOMALY_SEGMENTATION

src/otx/recipe/anomaly_segmentation/stfpm.yaml

@@ -3,6 +3,7 @@ model:
   init_args:
     layers: ["layer1", "layer2", "layer3"]
     backbone: "resnet18"
+    task: ANOMALY_SEGMENTATION
 
     optimizer:
       class_path: torch.optim.SGD