From c871af400d7c2a09941676b654fd1fefd80c9d82 Mon Sep 17 00:00:00 2001 From: Ashwin Vaidya Date: Thu, 9 May 2024 03:09:56 +0200 Subject: [PATCH] Add metrics to Anomaly OV Task (#3471) add metrics to openvino model Signed-off-by: Ashwin Vaidya --- src/otx/algo/anomaly/openvino_model.py | 86 ++++++++++++++++++++++++-- src/otx/core/model/anomaly.py | 44 ++++++------- 2 files changed, 99 insertions(+), 31 deletions(-) diff --git a/src/otx/algo/anomaly/openvino_model.py b/src/otx/algo/anomaly/openvino_model.py index 80120523691..04357807e4b 100644 --- a/src/otx/algo/anomaly/openvino_model.py +++ b/src/otx/algo/anomaly/openvino_model.py @@ -11,23 +11,85 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Sequence +import numpy as np +import torch +from anomalib.metrics import create_metric_collection +from lightning import Callback, Trainer +from torchvision.transforms.functional import resize + +from otx.core.data.entity.anomaly import AnomalyClassificationDataBatch from otx.core.metrics.types import MetricCallable, NullMetricCallable from otx.core.model.anomaly import AnomalyModelInputs from otx.core.model.base import OVModel if TYPE_CHECKING: + from anomalib.metrics import AnomalibMetricCollection from model_api.models import Model from model_api.models.anomaly import AnomalyResult +class _OVMetricCallback(Callback): + def __init__(self) -> None: + super().__init__() + + def on_test_epoch_start(self, trainer: Trainer, pl_module: AnomalyOpenVINO) -> None: + pl_module.image_metrics.reset() + pl_module.pixel_metrics.reset() + + def on_test_batch_end( + self, + trainer: Trainer, + pl_module: AnomalyOpenVINO, + outputs: list[AnomalyResult], + batch: AnomalyModelInputs, + batch_idx: int, + dataloader_idx: int = 0, + ) -> None: + # Convert modelAPI scores to anomaly scores. i.e flip scores with Normal label. + score_dict = { + "pred_scores": torch.tensor( + [output.pred_score if output.pred_label == "Anomaly" else 1 - output.pred_score for output in outputs], + ), + "labels": torch.tensor(batch.labels) if batch.batch_size == 1 else torch.vstack(batch.labels), + } + if not isinstance(batch, AnomalyClassificationDataBatch): + score_dict["anomaly_maps"] = torch.tensor(np.array([output.anomaly_map for output in outputs])) / 255.0 + score_dict["masks"] = batch.masks if batch.batch_size == 1 else torch.vstack(batch.masks) + # resize masks and anomaly maps to 256,256 as this is the size used in Anomalib + score_dict["masks"] = resize(score_dict["masks"], (256, 256)) + score_dict["anomaly_maps"] = resize(score_dict["anomaly_maps"], (256, 256)) + + self._update_metrics(pl_module.image_metrics, pl_module.pixel_metrics, score_dict) + + def on_test_epoch_end(self, trainer: Trainer, pl_module: AnomalyOpenVINO) -> None: + self._log_metrics(pl_module) + + def _update_metrics( + self, + image_metric: AnomalibMetricCollection, + pixel_metric: AnomalibMetricCollection, + outputs: dict[str, torch.Tensor], + ) -> None: + """Update performance metrics.""" + image_metric.update(outputs["pred_scores"], outputs["labels"].int()) + if "masks" in outputs and "anomaly_maps" in outputs: + pixel_metric.update(outputs["anomaly_maps"], outputs["masks"].int()) + + @staticmethod + def _log_metrics(pl_module: AnomalyOpenVINO) -> None: + """Log computed performance metrics.""" + if pl_module.pixel_metrics._update_called: # noqa: SLF001 + pl_module.log_dict(pl_module.pixel_metrics, prog_bar=True) + pl_module.log_dict(pl_module.image_metrics, prog_bar=False) + else: + pl_module.log_dict(pl_module.image_metrics, prog_bar=True) + + class AnomalyOpenVINO(OVModel): """Anomaly OpenVINO model.""" - # [TODO](ashwinvaidya17): Remove LightningModule once OTXModel is updated to use LightningModule. - # NOTE: Ideally OVModel should not be a LightningModule - def __init__( self, model_name: str, @@ -35,7 +97,7 @@ def __init__( max_num_requests: int | None = None, use_throughput_mode: bool = True, model_api_configuration: dict[str, Any] | None = None, - metric: MetricCallable = NullMetricCallable, + metric: MetricCallable = NullMetricCallable, # Metrics is computed using Anomalib's metric **kwargs, ) -> None: super().__init__( @@ -47,6 +109,9 @@ def __init__( model_api_configuration=model_api_configuration, metric=metric, ) + 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_") def _create_model(self) -> Model: from model_api.adapters import OpenvinoAdapter, create_core, get_user_config @@ -68,6 +133,10 @@ def _create_model(self) -> Model: configuration=self.model_api_configuration, ) + def configure_callbacks(self) -> Sequence[Callback] | Callback: + """Return the metric callback.""" + return _OVMetricCallback() + def test_step(self, inputs: AnomalyModelInputs, batch_idx: int) -> list[AnomalyResult]: """Return outputs from the OpenVINO model.""" return self.forward(inputs) # type: ignore[return-value] @@ -79,3 +148,10 @@ def predict_step(self, inputs: AnomalyModelInputs, batch_idx: int) -> list[Anoma def _customize_outputs(self, outputs: list[AnomalyResult], inputs: AnomalyModelInputs) -> list[AnomalyResult]: """Return outputs from the OpenVINO model as is.""" return outputs + + def _customize_inputs(self, inputs: AnomalyModelInputs) -> dict[str, np.ndarray]: + """Return inputs as is.""" + inputs = super()._customize_inputs(inputs) + # model needs inputs in range 0-1 + inputs["inputs"] = [value / 255.0 for value in inputs["inputs"]] + return inputs diff --git a/src/otx/core/model/anomaly.py b/src/otx/core/model/anomaly.py index d9e70afe4fb..40765da41ce 100644 --- a/src/otx/core/model/anomaly.py +++ b/src/otx/core/model/anomaly.py @@ -39,7 +39,6 @@ from lightning.pytorch.callbacks.callback import Callback from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable from torchmetrics import Metric - from torchvision.transforms.v2 import Transform AnomalyModelInputs: TypeAlias = ( @@ -78,7 +77,7 @@ def __init__( input_size=(1, 3, *image_shape), mean=mean_values, std=scale_values, - swap_rgb=False, # default value. Ideally, modelAPI should pass RGB inputs after the pre-processing step + swap_rgb=True, # BGR -> RGB ) @property @@ -145,8 +144,6 @@ def __init__(self) -> None: self.optimizer: list[OptimizerCallable] | OptimizerCallable = None self.scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = None self._input_size: tuple[int, int] = (256, 256) - self.mean_values: tuple[float, float, float] = (0.0, 0.0, 0.0) - self.scale_values: tuple[float, float, float] = (1.0, 1.0, 1.0) self.trainer: Trainer self.model: nn.Module self.image_threshold: BaseThreshold @@ -161,14 +158,12 @@ def on_save_checkpoint(self, checkpoint: dict[str, Any]) -> None: """Callback on saving checkpoint.""" super().on_save_checkpoint(checkpoint) # type: ignore[misc] - attrs = ["_task_type", "_input_size", "mean_values", "scale_values", "image_threshold", "pixel_threshold"] - + attrs = ["_task_type", "_input_size", "image_threshold", "pixel_threshold"] checkpoint["anomaly"] = {key: getattr(self, key, None) for key in attrs} def on_load_checkpoint(self, checkpoint: dict[str, Any]) -> None: """Callback on loading checkpoint.""" super().on_load_checkpoint(checkpoint) # type: ignore[misc] - if anomaly_attrs := checkpoint.get("anomaly"): for key, value in anomaly_attrs.items(): setattr(self, key, value) @@ -206,15 +201,21 @@ def task(self, value: OTXTaskType) -> None: msg = f"Unexpected task type: {value}" raise ValueError(msg) - def _extract_mean_scale_from_transforms(self, transforms: list[Transform]) -> None: - """Extract mean and scale values from transforms.""" - for transform in transforms: + def _get_values_from_transforms( + self, + key_name: str, + ) -> tuple: + """Get the value requested value from default transforms.""" + for transform in self.configure_transforms().transforms: # type: ignore[attr-defined] name = transform.__class__.__name__ - if "Resize" in name: - self.input_size = transform.size * 2 # transform.size has value [size], so *2 gives (size, size) + if "Resize" in name and key_name == "input_size": + image_size = transform.size elif "Normalize" in name: - self.mean_values = transform.mean - self.scale_values = transform.std + if key_name == "mean": + mean_value = transform.mean + elif key_name == "scale": + std_value = transform.std + return image_size, mean_value, std_value @property def trainable_model(self) -> str | None: @@ -228,15 +229,6 @@ def trainable_model(self) -> str | None: """ return None - def setup(self, stage: str | None = None) -> None: - """Setup the model.""" - super().setup(stage) # type: ignore[misc] - if stage == "fit" and hasattr(self.trainer, "datamodule") and hasattr(self.trainer.datamodule, "config"): - if hasattr(self.trainer.datamodule.config, "test_subset"): - self._extract_mean_scale_from_transforms(self.trainer.datamodule.config.test_subset.transforms) - elif hasattr(self.trainer.datamodule.config, "val_subset"): - self._extract_mean_scale_from_transforms(self.trainer.datamodule.config.val_subset.transforms) - def configure_callbacks(self) -> list[Callback]: """Get all necessary callbacks required for training and post-processing on Anomalib models.""" image_metrics = ["AUROC", "F1Score"] @@ -414,14 +406,14 @@ def export( """ min_val = self.normalization_metrics.state_dict()["min"].cpu().numpy().tolist() max_val = self.normalization_metrics.state_dict()["max"].cpu().numpy().tolist() - image_shape = (256, 256) if self.input_size is None else self.input_size + image_shape, mean_values, scale_values = self._get_values_from_transforms("input_size") exporter = _AnomalyModelExporter( image_shape=image_shape, image_threshold=self.image_threshold.value.cpu().numpy().tolist(), pixel_threshold=self.pixel_threshold.value.cpu().numpy().tolist(), task=self.task, - mean_values=self.mean_values, - scale_values=self.scale_values, + mean_values=mean_values, + scale_values=scale_values, normalization_scale=max_val - min_val, ) return exporter.export(