From 18ecbe2864e278540cb929367fdd3190075c3207 Mon Sep 17 00:00:00 2001 From: Ashwin Vaidya Date: Fri, 8 Apr 2022 08:19:20 +0200 Subject: [PATCH 1/2] Initial callback update --- anomalib/utils/callbacks/cdf_normalization.py | 13 +++++++------ anomalib/utils/callbacks/min_max_normalization.py | 9 +++++---- anomalib/utils/callbacks/model_loader.py | 6 ++++-- anomalib/utils/callbacks/openvino.py | 7 +++---- tests/nightly/models/test_model_nightly.py | 10 ++++++++-- 5 files changed, 27 insertions(+), 18 deletions(-) diff --git a/anomalib/utils/callbacks/cdf_normalization.py b/anomalib/utils/callbacks/cdf_normalization.py index be072be236..d1ae80fae9 100644 --- a/anomalib/utils/callbacks/cdf_normalization.py +++ b/anomalib/utils/callbacks/cdf_normalization.py @@ -22,6 +22,7 @@ from torch.distributions import LogNormal from anomalib.models import get_model +from anomalib.models.components import AnomalyModule from anomalib.post_processing.normalization.cdf import normalize, standardize @@ -32,12 +33,12 @@ def __init__(self): self.image_dist: Optional[LogNormal] = None self.pixel_dist: Optional[LogNormal] = None - def on_test_start(self, _trainer: pl.Trainer, pl_module: pl.LightningModule) -> None: + def on_test_start(self, _trainer: pl.Trainer, pl_module: AnomalyModule) -> None: """Called when the test begins.""" pl_module.image_metrics.F1.threshold = 0.5 pl_module.pixel_metrics.F1.threshold = 0.5 - def on_validation_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None: + def on_validation_epoch_start(self, trainer: "pl.Trainer", pl_module: AnomalyModule) -> None: """Called when the validation starts after training. Use the current model to compute the anomaly score distributions @@ -49,7 +50,7 @@ def on_validation_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.Lightn def on_validation_batch_end( self, _trainer: pl.Trainer, - pl_module: pl.LightningModule, + pl_module: AnomalyModule, outputs: Optional[STEP_OUTPUT], _batch: Any, _batch_idx: int, @@ -61,7 +62,7 @@ def on_validation_batch_end( def on_test_batch_end( self, _trainer: pl.Trainer, - pl_module: pl.LightningModule, + pl_module: AnomalyModule, outputs: Optional[STEP_OUTPUT], _batch: Any, _batch_idx: int, @@ -74,7 +75,7 @@ def on_test_batch_end( def on_predict_batch_end( self, _trainer: pl.Trainer, - pl_module: pl.LightningModule, + pl_module: AnomalyModule, outputs: Dict, _batch: Any, _batch_idx: int, @@ -120,7 +121,7 @@ def _standardize_batch(outputs: STEP_OUTPUT, pl_module) -> None: ) @staticmethod - def _normalize_batch(outputs: STEP_OUTPUT, pl_module: pl.LightningModule) -> None: + def _normalize_batch(outputs: STEP_OUTPUT, pl_module: AnomalyModule) -> None: outputs["pred_scores"] = normalize(outputs["pred_scores"], pl_module.image_threshold.value) if "anomaly_maps" in outputs.keys(): outputs["anomaly_maps"] = normalize(outputs["anomaly_maps"], pl_module.pixel_threshold.value) diff --git a/anomalib/utils/callbacks/min_max_normalization.py b/anomalib/utils/callbacks/min_max_normalization.py index acdbb03cb7..9fd5dde1d2 100644 --- a/anomalib/utils/callbacks/min_max_normalization.py +++ b/anomalib/utils/callbacks/min_max_normalization.py @@ -20,13 +20,14 @@ from pytorch_lightning import Callback from pytorch_lightning.utilities.types import STEP_OUTPUT +from anomalib.models.components import AnomalyModule from anomalib.post_processing.normalization.min_max import normalize class MinMaxNormalizationCallback(Callback): """Callback that normalizes the image-level and pixel-level anomaly scores using min-max normalization.""" - def on_test_start(self, _trainer: pl.Trainer, pl_module: pl.LightningModule) -> None: + def on_test_start(self, _trainer: pl.Trainer, pl_module: AnomalyModule) -> None: """Called when the test begins.""" pl_module.image_metrics.F1.threshold = 0.5 pl_module.pixel_metrics.F1.threshold = 0.5 @@ -34,7 +35,7 @@ def on_test_start(self, _trainer: pl.Trainer, pl_module: pl.LightningModule) -> def on_validation_batch_end( self, _trainer: pl.Trainer, - pl_module: pl.LightningModule, + pl_module: AnomalyModule, outputs: STEP_OUTPUT, _batch: Any, _batch_idx: int, @@ -49,7 +50,7 @@ def on_validation_batch_end( def on_test_batch_end( self, _trainer: pl.Trainer, - pl_module: pl.LightningModule, + pl_module: AnomalyModule, outputs: STEP_OUTPUT, _batch: Any, _batch_idx: int, @@ -61,7 +62,7 @@ def on_test_batch_end( def on_predict_batch_end( self, _trainer: pl.Trainer, - pl_module: pl.LightningModule, + pl_module: AnomalyModule, outputs: Dict, _batch: Any, _batch_idx: int, diff --git a/anomalib/utils/callbacks/model_loader.py b/anomalib/utils/callbacks/model_loader.py index 6c4fa0f278..1c156be0bc 100644 --- a/anomalib/utils/callbacks/model_loader.py +++ b/anomalib/utils/callbacks/model_loader.py @@ -15,7 +15,9 @@ # and limitations under the License. import torch -from pytorch_lightning import Callback, LightningModule +from pytorch_lightning import Callback + +from anomalib.models.components import AnomalyModule class LoadModelCallback(Callback): @@ -24,7 +26,7 @@ class LoadModelCallback(Callback): def __init__(self, weights_path): self.weights_path = weights_path - def on_test_start(self, trainer, pl_module: LightningModule) -> None: # pylint: disable=W0613 + def on_test_start(self, trainer, pl_module: AnomalyModule) -> None: # pylint: disable=W0613 """Call when the test begins. Loads the model weights from ``weights_path`` into the PyTorch module. diff --git a/anomalib/utils/callbacks/openvino.py b/anomalib/utils/callbacks/openvino.py index b07cd7de03..4cff1f814a 100644 --- a/anomalib/utils/callbacks/openvino.py +++ b/anomalib/utils/callbacks/openvino.py @@ -15,9 +15,9 @@ # and limitations under the License. import os -from typing import Tuple, cast +from typing import Tuple -from pytorch_lightning import Callback, LightningModule +from pytorch_lightning import Callback from anomalib.deploy import export_convert from anomalib.models.components import AnomalyModule @@ -39,7 +39,7 @@ def __init__(self, input_size: Tuple[int, int], dirpath: str, filename: str): self.dirpath = dirpath self.filename = filename - def on_train_end(self, trainer, pl_module: LightningModule) -> None: # pylint: disable=W0613 + def on_train_end(self, trainer, pl_module: AnomalyModule) -> None: # pylint: disable=W0613 """Call when the train ends. Converts the model to ``onnx`` format and then calls OpenVINO's model optimizer to get the @@ -47,7 +47,6 @@ def on_train_end(self, trainer, pl_module: LightningModule) -> None: # pylint: """ os.makedirs(self.dirpath, exist_ok=True) onnx_path = os.path.join(self.dirpath, self.filename + ".onnx") - pl_module = cast(AnomalyModule, pl_module) export_convert( model=pl_module, input_size=self.input_size, diff --git a/tests/nightly/models/test_model_nightly.py b/tests/nightly/models/test_model_nightly.py index 8a32da9533..30dfda6da1 100644 --- a/tests/nightly/models/test_model_nightly.py +++ b/tests/nightly/models/test_model_nightly.py @@ -90,13 +90,19 @@ def _test_metrics(self, trainer, config, model, datamodule): threshold = thresholds[config.model.name][config.dataset.category] if "optimization" in config.keys() and config.optimization.nncf.apply: threshold = threshold.nncf - if not (np.isclose(results["image_AUROC"], threshold["image_AUROC"], rtol=0.02) or (results["image_AUROC"] >= threshold["image_AUROC"])): + if not ( + np.isclose(results["image_AUROC"], threshold["image_AUROC"], rtol=0.02) + or (results["image_AUROC"] >= threshold["image_AUROC"]) + ): raise AssertionError( f"results['image_AUROC']:{results['image_AUROC']} >= threshold['image_AUROC']:{threshold['image_AUROC']}" ) if config.dataset.task == "segmentation": - if not (np.isclose(results["pixel_AUROC"] ,threshold["pixel_AUROC"], rtol=0.02) or (results["pixel_AUROC"] >= threshold["pixel_AUROC"])): + if not ( + np.isclose(results["pixel_AUROC"], threshold["pixel_AUROC"], rtol=0.02) + or (results["pixel_AUROC"] >= threshold["pixel_AUROC"]) + ): raise AssertionError( f"results['pixel_AUROC']:{results['pixel_AUROC']} >= threshold['pixel_AUROC']:{threshold['pixel_AUROC']}" ) From 0b6d5e5a5dee8a412473579a90166c2898b6c7b4 Mon Sep 17 00:00:00 2001 From: Ashwin Vaidya Date: Fri, 8 Apr 2022 08:59:37 +0200 Subject: [PATCH 2/2] Modify visualizer callback --- anomalib/utils/callbacks/visualizer_callback.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/anomalib/utils/callbacks/visualizer_callback.py b/anomalib/utils/callbacks/visualizer_callback.py index 7879306624..5e6c9e82ca 100644 --- a/anomalib/utils/callbacks/visualizer_callback.py +++ b/anomalib/utils/callbacks/visualizer_callback.py @@ -87,7 +87,7 @@ def _add_images( def on_test_batch_end( self, _trainer: pl.Trainer, - pl_module: pl.LightningModule, + pl_module: AnomalyModule, outputs: Optional[STEP_OUTPUT], _batch: Any, _batch_idx: int, @@ -149,7 +149,7 @@ def on_test_batch_end( self._add_images(visualizer, pl_module, Path(filename)) visualizer.close() - def on_test_end(self, _trainer: pl.Trainer, pl_module: pl.LightningModule) -> None: + def on_test_end(self, _trainer: pl.Trainer, pl_module: AnomalyModule) -> None: """Sync logs. Currently only ``AnomalibWandbLogger`` is called from this method. This is because logging as a single batch @@ -157,7 +157,7 @@ def on_test_end(self, _trainer: pl.Trainer, pl_module: pl.LightningModule) -> No Args: _trainer (pl.Trainer): Pytorch Lightning trainer (unused) - pl_module (pl.LightningModule): Anomaly module + pl_module (AnomalyModule): Anomaly module """ if pl_module.logger is not None and isinstance(pl_module.logger, AnomalibWandbLogger): pl_module.logger.save()