From d55512eb7843bdb913f94b1580d43506d0c72624 Mon Sep 17 00:00:00 2001 From: Harim Kang Date: Thu, 27 Jun 2024 10:37:20 +0900 Subject: [PATCH] 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 | 9 ++++++++- src/otx/algo/anomaly/padim.py | 14 +++++++++++--- src/otx/algo/anomaly/stfpm.py | 14 +++++++++++--- src/otx/engine/engine.py | 13 ------------- .../anomaly_classification/openvino_model.yaml | 1 + src/otx/recipe/anomaly_classification/padim.yaml | 1 + src/otx/recipe/anomaly_classification/stfpm.yaml | 1 + .../recipe/anomaly_detection/openvino_model.yaml | 1 + src/otx/recipe/anomaly_detection/padim.yaml | 1 + src/otx/recipe/anomaly_detection/stfpm.yaml | 1 + .../anomaly_segmentation/openvino_model.yaml | 1 + src/otx/recipe/anomaly_segmentation/padim.yaml | 1 + src/otx/recipe/anomaly_segmentation/stfpm.yaml | 1 + 13 files changed, 39 insertions(+), 20 deletions(-) diff --git a/src/otx/algo/anomaly/openvino_model.py b/src/otx/algo/anomaly/openvino_model.py index 4b47a99a05d..ab74d8b15ff 100644 --- a/src/otx/algo/anomaly/openvino_model.py +++ b/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 diff --git a/src/otx/algo/anomaly/padim.py b/src/otx/algo/anomaly/padim.py index 57d731b7cb1..201b0230a02 100644 --- a/src/otx/algo/anomaly/padim.py +++ b/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.""" diff --git a/src/otx/algo/anomaly/stfpm.py b/src/otx/algo/anomaly/stfpm.py index cf84d3c84cf..72dd30e8aa3 100644 --- a/src/otx/algo/anomaly/stfpm.py +++ b/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: diff --git a/src/otx/engine/engine.py b/src/otx/engine/engine.py index 9ffc0c2fe34..fcd3b7ecce5 100644 --- a/src/otx/engine/engine.py +++ b/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 diff --git a/src/otx/recipe/anomaly_classification/openvino_model.yaml b/src/otx/recipe/anomaly_classification/openvino_model.yaml index feac4c139d7..a677af762f1 100644 --- a/src/otx/recipe/anomaly_classification/openvino_model.yaml +++ b/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 diff --git a/src/otx/recipe/anomaly_classification/padim.yaml b/src/otx/recipe/anomaly_classification/padim.yaml index 6faff2ab27b..70cd9fa0961 100644 --- a/src/otx/recipe/anomaly_classification/padim.yaml +++ b/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 diff --git a/src/otx/recipe/anomaly_classification/stfpm.yaml b/src/otx/recipe/anomaly_classification/stfpm.yaml index 3364156dcb0..879365d9f63 100644 --- a/src/otx/recipe/anomaly_classification/stfpm.yaml +++ b/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 diff --git a/src/otx/recipe/anomaly_detection/openvino_model.yaml b/src/otx/recipe/anomaly_detection/openvino_model.yaml index 1e86094e832..b26e47afe77 100644 --- a/src/otx/recipe/anomaly_detection/openvino_model.yaml +++ b/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 diff --git a/src/otx/recipe/anomaly_detection/padim.yaml b/src/otx/recipe/anomaly_detection/padim.yaml index 99a3e86615c..41383ed9026 100644 --- a/src/otx/recipe/anomaly_detection/padim.yaml +++ b/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 diff --git a/src/otx/recipe/anomaly_detection/stfpm.yaml b/src/otx/recipe/anomaly_detection/stfpm.yaml index e16d75c8f00..5b17fc66915 100644 --- a/src/otx/recipe/anomaly_detection/stfpm.yaml +++ b/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 diff --git a/src/otx/recipe/anomaly_segmentation/openvino_model.yaml b/src/otx/recipe/anomaly_segmentation/openvino_model.yaml index 3d478c1d15f..8c27b0cb274 100644 --- a/src/otx/recipe/anomaly_segmentation/openvino_model.yaml +++ b/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 diff --git a/src/otx/recipe/anomaly_segmentation/padim.yaml b/src/otx/recipe/anomaly_segmentation/padim.yaml index eaee80185a4..f067d1781d9 100644 --- a/src/otx/recipe/anomaly_segmentation/padim.yaml +++ b/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 diff --git a/src/otx/recipe/anomaly_segmentation/stfpm.yaml b/src/otx/recipe/anomaly_segmentation/stfpm.yaml index a9c7cdfb3b4..84e880e47f3 100644 --- a/src/otx/recipe/anomaly_segmentation/stfpm.yaml +++ b/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