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(