[OTX/Anomaly] Add changes from external to otx (#1452)

* Add changes from external to otx

* Address PR comments

* Update config files + remove backbone from base

* Fix pre-merge checks

Co-authored-by: Ashwin Vaidya <ashwinitinvaidya@gmail.com>

otx/algorithms/anomaly/adapters/anomalib/callbacks/__init__.py

@@ -16,6 +16,5 @@
 
 from .inference import AnomalyInferenceCallback
 from .progress import ProgressCallback
-from .score_report import ScoreReportingCallback
 
-__all__ = ["AnomalyInferenceCallback", "ProgressCallback", "ScoreReportingCallback"]
+__all__ = ["AnomalyInferenceCallback", "ProgressCallback"]

otx/algorithms/anomaly/adapters/anomalib/callbacks/progress.py

@@ -1,4 +1,7 @@
-"""Progressbar Callback for OTX task."""
+"""Progressbar  and Score Reporting callback Callback for OTX task.
+
+TODO Since only one progressbar callback is supported HPO is combined into one callback. Remove this after the refactor
+"""
 
 # Copyright (C) 2021 Intel Corporation
 #
@@ -38,9 +41,9 @@ def __init__(
         self._progress: float = 0
 
         if parameters is not None:
-            self.update_progress_callback = parameters.update_progress
+            self.progress_and_hpo_callback = parameters.update_progress
         else:
-            self.update_progress_callback = default_progress_callback
+            self.progress_and_hpo_callback = default_progress_callback
 
     def on_train_start(self, trainer, pl_module):
         """Store max epochs and current epoch from trainer."""
@@ -75,6 +78,22 @@ def on_test_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, datal
         super().on_test_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx)
         self._update_progress(stage="test")
 
+    def on_validation_epoch_end(self, trainer, pl_module):  # pylint: disable=unused-argument
+        """If score exists in trainer.logged_metrics, report the score."""
+        if self.progress_and_hpo_callback is not None:
+            score = None
+            metric = getattr(self.progress_and_hpo_callback, "metric", None)
+            print(f"[DEBUG-HPO] logged_metrics = {trainer.logged_metrics}")
+            if metric in trainer.logged_metrics:
+                score = float(trainer.logged_metrics[metric])
+                if score < 1.0:
+                    score = score + int(trainer.global_step)
+                else:
+                    score = -(score + int(trainer.global_step))
+
+            # Always assumes that hpo validation step is called during training.
+            self.progress_and_hpo_callback(int(self._get_progress("train")), score)  # pylint: disable=not-callable
+
     def _reset_progress(self):
         self._progress = 0.0
 
@@ -104,4 +123,4 @@ def _get_progress(self, stage: str = "train") -> float:
 
     def _update_progress(self, stage: str):
         progress = self._get_progress(stage)
-        self.update_progress_callback(int(progress), None)
+        self.progress_and_hpo_callback(int(progress), None)

otx/algorithms/anomaly/configs/classification/draem/configuration.yaml

@@ -190,7 +190,8 @@ nncf_optimization:
       operator: AND
       rules: []
       type: UI_RULES
-    visible_in_ui: true
+    value: false
+    visible_in_ui: false
     warning: null
   type: PARAMETER_GROUP
   visible_in_ui: true

otx/algorithms/anomaly/configs/classification/padim/configuration.yaml

@@ -130,7 +130,7 @@ nncf_optimization:
       rules: []
       type: UI_RULES
     value: false
-    visible_in_ui: true
+    visible_in_ui: false
     warning: null
   type: PARAMETER_GROUP
   visible_in_ui: true

otx/algorithms/anomaly/configs/classification/stfpm/configuration.yaml

@@ -259,7 +259,7 @@ nncf_optimization:
       rules: []
       type: UI_RULES
     value: false
-    visible_in_ui: true
+    visible_in_ui: false
     warning: null
   type: PARAMETER_GROUP
   visible_in_ui: true

otx/algorithms/anomaly/configs/detection/draem/configuration.yaml

@@ -190,7 +190,8 @@ nncf_optimization:
       operator: AND
       rules: []
       type: UI_RULES
-    visible_in_ui: true
+    value: false
+    visible_in_ui: false
     warning: null
   type: PARAMETER_GROUP
   visible_in_ui: true

otx/algorithms/anomaly/configs/detection/padim/configuration.yaml

@@ -130,7 +130,7 @@ nncf_optimization:
       rules: []
       type: UI_RULES
     value: false
-    visible_in_ui: true
+    visible_in_ui: false
     warning: null
   type: PARAMETER_GROUP
   visible_in_ui: true

otx/algorithms/anomaly/configs/detection/stfpm/configuration.yaml

@@ -259,7 +259,7 @@ nncf_optimization:
       rules: []
       type: UI_RULES
     value: false
-    visible_in_ui: true
+    visible_in_ui: false
     warning: null
   type: PARAMETER_GROUP
   visible_in_ui: true

otx/algorithms/anomaly/configs/segmentation/draem/configuration.yaml

@@ -190,7 +190,8 @@ nncf_optimization:
       operator: AND
       rules: []
       type: UI_RULES
-    visible_in_ui: true
+    value: false
+    visible_in_ui: false
     warning: null
   type: PARAMETER_GROUP
   visible_in_ui: true

otx/algorithms/anomaly/configs/segmentation/padim/configuration.yaml

@@ -130,7 +130,7 @@ nncf_optimization:
       rules: []
       type: UI_RULES
     value: false
-    visible_in_ui: true
+    visible_in_ui: false
     warning: null
   type: PARAMETER_GROUP
   visible_in_ui: true

otx/algorithms/anomaly/configs/segmentation/stfpm/configuration.yaml

@@ -259,7 +259,7 @@ nncf_optimization:
       rules: []
       type: UI_RULES
     value: false
-    visible_in_ui: true
+    visible_in_ui: false
     warning: null
   type: PARAMETER_GROUP
   visible_in_ui: true

otx/algorithms/anomaly/tasks/inference.py

@@ -128,8 +128,8 @@ def load_model(self, otx_model: Optional[ModelEntity]) -> AnomalyModule:
             AnomalyModule: Anomalib
                 classification or segmentation model with/without weights.
         """
-        model = get_model(config=self.config)
         if otx_model is None:
+            model = get_model(config=self.config)
             logger.info(
                 "No trained model in project yet. Created new model with '%s'",
                 self.model_name,
@@ -138,10 +138,16 @@ def load_model(self, otx_model: Optional[ModelEntity]) -> AnomalyModule:
             buffer = io.BytesIO(otx_model.get_data("weights.pth"))
             model_data = torch.load(buffer, map_location=torch.device("cpu"))
 
+            if model_data["config"]["model"]["backbone"] != self.config["model"]["backbone"]:
+                logger.warning(
+                    "Backbone of the model in the Task Environment is different from the one in the template. "
+                    f"creating model with backbone={model_data['config']['model']['backbone']}"
+                )
+                self.config["model"]["backbone"] = model_data["config"]["model"]["backbone"]
             try:
+                model = get_model(config=self.config)
                 model.load_state_dict(model_data["model"])
                 logger.info("Loaded model weights from Task Environment")
-
             except BaseException as exception:
                 raise ValueError("Could not load the saved model. The model file structure is invalid.") from exception
 
@@ -251,8 +257,8 @@ def export(self, export_type: ExportType, output_model: ModelEntity) -> None:
         logger.info("Exporting the OpenVINO model.")
         onnx_path = os.path.join(self.config.project.path, "onnx_model.onnx")
         self._export_to_onnx(onnx_path)
-        optimize_command = "mo --input_model " + onnx_path + " --output_dir " + self.config.project.path
-        subprocess.call(optimize_command, shell=True)
+        optimize_command = ["mo", "--input_model", onnx_path, "--output_dir", self.config.project.path]
+        subprocess.run(optimize_command, check=True)
         bin_file = glob(os.path.join(self.config.project.path, "*.bin"))[0]
         xml_file = glob(os.path.join(self.config.project.path, "*.xml"))[0]
         with open(bin_file, "rb") as file:
@@ -266,7 +272,7 @@ def export(self, export_type: ExportType, output_model: ModelEntity) -> None:
         output_model.set_data("label_schema.json", label_schema_to_bytes(self.task_environment.label_schema))
         self._set_metadata(output_model)
 
-    def _model_info(self) -> Dict:
+    def model_info(self) -> Dict:
         """Return model info to save the model weights.
 
         Returns:
@@ -285,7 +291,7 @@ def save_model(self, output_model: ModelEntity) -> None:
             output_model (ModelEntity): Output model onto which the weights are saved.
         """
         logger.info("Saving the model weights.")
-        model_info = self._model_info()
+        model_info = self.model_info()
         buffer = io.BytesIO()
         torch.save(model_info, buffer)
         output_model.set_data("weights.pth", buffer.getvalue())
@@ -301,8 +307,10 @@ def save_model(self, output_model: ModelEntity) -> None:
         output_model.optimization_methods = self.optimization_methods
 
     def _set_metadata(self, output_model: ModelEntity):
-        output_model.set_data("image_threshold", self.model.image_threshold.value.cpu().numpy().tobytes())
-        output_model.set_data("pixel_threshold", self.model.pixel_threshold.value.cpu().numpy().tobytes())
+        if hasattr(self.model, "image_threshold"):
+            output_model.set_data("image_threshold", self.model.image_threshold.value.cpu().numpy().tobytes())
+        if hasattr(self.model, "pixel_threshold"):
+            output_model.set_data("pixel_threshold", self.model.pixel_threshold.value.cpu().numpy().tobytes())
         if hasattr(self.model, "normalization_metrics"):
             output_model.set_data("min", self.model.normalization_metrics.state_dict()["min"].cpu().numpy().tobytes())
             output_model.set_data("max", self.model.normalization_metrics.state_dict()["max"].cpu().numpy().tobytes())

otx/algorithms/anomaly/tasks/nncf.py

@@ -202,7 +202,7 @@ def optimize(
 
         logger.info("Training completed.")
 
-    def _model_info(self) -> Dict:
+    def model_info(self) -> Dict:
         """Return model info to save the model weights.
 
         Returns:

otx/algorithms/anomaly/tasks/train.py

@@ -14,18 +14,18 @@
 # See the License for the specific language governing permissions
 # and limitations under the License.
 
+import io
 from typing import Optional
 
+import torch
+from anomalib.models import AnomalyModule, get_model
 from anomalib.utils.callbacks import (
     MetricsConfigurationCallback,
     MinMaxNormalizationCallback,
 )
 from pytorch_lightning import Trainer, seed_everything
 
-from otx.algorithms.anomaly.adapters.anomalib.callbacks import (
-    ProgressCallback,
-    ScoreReportingCallback,
-)
+from otx.algorithms.anomaly.adapters.anomalib.callbacks import ProgressCallback
 from otx.algorithms.anomaly.adapters.anomalib.data import OTXAnomalyDataModule
 from otx.algorithms.anomaly.adapters.anomalib.logger import get_logger
 from otx.api.entities.datasets import DatasetEntity
@@ -72,7 +72,6 @@ def train(
         callbacks = [
             ProgressCallback(parameters=train_parameters),
             MinMaxNormalizationCallback(),
-            ScoreReportingCallback(parameters=train_parameters),
             MetricsConfigurationCallback(
                 adaptive_threshold=config.metrics.threshold.adaptive,
                 default_image_threshold=config.metrics.threshold.image_default,
@@ -88,3 +87,42 @@ def train(
         self.save_model(output_model)
 
         logger.info("Training completed.")
+
+    def load_model(self, otx_model: Optional[ModelEntity]) -> AnomalyModule:
+        """Create and Load Anomalib Module from OTE Model.
+
+        This method checks if the task environment has a saved OTE Model,
+        and creates one. If the OTE model already exists, it returns the
+        the model with the saved weights.
+
+        Args:
+            otx_model (Optional[ModelEntity]): OTE Model from the
+                task environment.
+
+        Returns:
+            AnomalyModule: Anomalib
+                classification or segmentation model with/without weights.
+        """
+        model = get_model(config=self.config)
+        if otx_model is None:
+            logger.info(
+                "No trained model in project yet. Created new model with '%s'",
+                self.model_name,
+            )
+        else:
+            buffer = io.BytesIO(otx_model.get_data("weights.pth"))
+            model_data = torch.load(buffer, map_location=torch.device("cpu"))
+
+            try:
+                if model_data["config"]["model"]["backbone"] == self.config["model"]["backbone"]:
+                    model.load_state_dict(model_data["model"])
+                    logger.info("Loaded model weights from Task Environment")
+                else:
+                    logger.info(
+                        "Model backbone does not match. Created new model with '%s'",
+                        self.model_name,
+                    )
+            except BaseException as exception:
+                raise ValueError("Could not load the saved model. The model file structure is invalid.") from exception
+
+        return model

otx/api/utils/vis_utils.py

@@ -3,7 +3,7 @@
 # Copyright (C) 2021-2022 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Union
+from typing import Iterable, Union
 
 import cv2
 import numpy as np