From ef2deb08fb790ce3c4740579174a5e7f0c8e84d6 Mon Sep 17 00:00:00 2001
From: Samet Akcay <samet.akcay@intel.com>
Date: Mon, 9 Dec 2024 13:02:05 +0000
Subject: [PATCH 1/9] Refactor post-processor to match the pre-processor and
 evaluator pattern

Signed-off-by: Samet Akcay <samet.akcay@intel.com>
---
 .../models/components/base/anomaly_module.py  | 56 ++++++++++++++++---
 .../models/image/cfa/lightning_model.py       |  2 +-
 .../models/image/cflow/lightning_model.py     |  2 +-
 .../models/image/csflow/lightning_model.py    |  2 +-
 .../models/image/dfkde/lightning_model.py     |  2 +-
 .../models/image/dfm/lightning_model.py       |  2 +-
 .../models/image/draem/lightning_model.py     |  2 +-
 .../models/image/dsr/lightning_model.py       |  2 +-
 .../image/efficient_ad/lightning_model.py     |  2 +-
 .../models/image/fastflow/lightning_model.py  |  2 +-
 .../models/image/fre/lightning_model.py       |  2 +-
 .../models/image/ganomaly/lightning_model.py  |  2 +-
 .../models/image/padim/lightning_model.py     |  4 +-
 .../models/image/patchcore/lightning_model.py |  4 +-
 .../reverse_distillation/lightning_model.py   |  2 +-
 .../models/image/stfpm/lightning_model.py     |  2 +-
 .../models/image/uflow/lightning_model.py     |  2 +-
 .../models/image/vlm_ad/lightning_model.py    |  3 +-
 .../models/image/winclip/lightning_model.py   |  4 +-
 .../models/video/ai_vad/lightning_model.py    |  5 +-
 .../dummy_lightning_model.py                  |  2 +-
 21 files changed, 73 insertions(+), 33 deletions(-)

diff --git a/src/anomalib/models/components/base/anomaly_module.py b/src/anomalib/models/components/base/anomaly_module.py
index 408200231a..c7b0d0d420 100644
--- a/src/anomalib/models/components/base/anomaly_module.py
+++ b/src/anomalib/models/components/base/anomaly_module.py
@@ -40,7 +40,7 @@ class AnomalibModule(ExportMixin, pl.LightningModule, ABC):
     def __init__(
         self,
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__()
@@ -52,11 +52,11 @@ def __init__(
         self.callbacks: list[Callback]
 
         self.pre_processor = self._resolve_pre_processor(pre_processor)
-        self.post_processor = post_processor or self.default_post_processor()
+        self.post_processor = self._resolve_post_processor(post_processor)
         self.evaluator = self._resolve_evaluator(evaluator)
 
         self._input_size: tuple[int, int] | None = None
-        self._is_setup = False  # flag to track if setup has been called from the trainer
+        self._is_setup = False
 
     @property
     def name(self) -> str:
@@ -214,15 +214,53 @@ def configure_pre_processor(cls, image_size: tuple[int, int] | None = None) -> P
             ]),
         )
 
-    def default_post_processor(self) -> PostProcessor | None:
-        """Default post processor.
+    def _resolve_post_processor(self, post_processor: PostProcessor | bool) -> PostProcessor | None:
+        """Resolve and validate which post-processor to use.
 
-        Override in subclass for model-specific post-processing behaviour.
+        Args:
+            post_processor: Post-processor configuration
+                - True -> use default post-processor
+                - False -> no post-processor
+                - PostProcessor -> use the provided post-processor
+
+        Returns:
+            Configured post-processor
+        """
+        if isinstance(post_processor, PostProcessor):
+            return post_processor
+        if isinstance(post_processor, bool):
+            return self.configure_post_processor() if post_processor else None
+        msg = f"Invalid post-processor type: {type(post_processor)}"
+        raise TypeError(msg)
+
+    @classmethod
+    def configure_post_processor(cls) -> PostProcessor | None:
+        """Configure the default post-processor based on the learning type.
+
+        Returns:
+            PostProcessor: Configured post-processor instance.
+
+        Raises:
+            NotImplementedError: If no default post-processor is available for the model's learning type.
+
+        Examples:
+            Get default post-processor:
+
+            >>> post_processor = AnomalibModule.configure_post_processor()
+
+            Create model with custom post-processor:
+
+            >>> custom_post_processor = CustomPostProcessor()
+            >>> model = PatchCore(post_processor=custom_post_processor)
+
+            Disable post-processing:
+
+            >>> model = PatchCore(post_processor=False)
         """
-        if self.learning_type == LearningType.ONE_CLASS:
+        if cls.learning_type == LearningType.ONE_CLASS:
             return OneClassPostProcessor()
-        msg = f"No default post-processor available for model {self.__name__} with learning type {self.learning_type}. \
-              Please override the default_post_processor method in the model implementation."
+        msg = f"No default post-processor available for model with learning type {cls.learning_type}. \
+              Please override the configure_post_processor method in the model implementation."
         raise NotImplementedError(msg)
 
     def _resolve_evaluator(self, evaluator: Evaluator | bool) -> Evaluator | None:
diff --git a/src/anomalib/models/image/cfa/lightning_model.py b/src/anomalib/models/image/cfa/lightning_model.py
index ea4bf3a2bd..4c2a9cffe2 100644
--- a/src/anomalib/models/image/cfa/lightning_model.py
+++ b/src/anomalib/models/image/cfa/lightning_model.py
@@ -59,7 +59,7 @@ def __init__(
         num_hard_negative_features: int = 3,
         radius: float = 1e-5,
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
diff --git a/src/anomalib/models/image/cflow/lightning_model.py b/src/anomalib/models/image/cflow/lightning_model.py
index d6b39751d0..f82872cb84 100644
--- a/src/anomalib/models/image/cflow/lightning_model.py
+++ b/src/anomalib/models/image/cflow/lightning_model.py
@@ -71,7 +71,7 @@ def __init__(
         permute_soft: bool = False,
         lr: float = 0.0001,
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
diff --git a/src/anomalib/models/image/csflow/lightning_model.py b/src/anomalib/models/image/csflow/lightning_model.py
index 3be936cc90..3324972422 100644
--- a/src/anomalib/models/image/csflow/lightning_model.py
+++ b/src/anomalib/models/image/csflow/lightning_model.py
@@ -48,7 +48,7 @@ def __init__(
         clamp: int = 3,
         num_channels: int = 3,
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
diff --git a/src/anomalib/models/image/dfkde/lightning_model.py b/src/anomalib/models/image/dfkde/lightning_model.py
index d1b1fba497..be141c45ef 100644
--- a/src/anomalib/models/image/dfkde/lightning_model.py
+++ b/src/anomalib/models/image/dfkde/lightning_model.py
@@ -50,7 +50,7 @@ def __init__(
         feature_scaling_method: FeatureScalingMethod = FeatureScalingMethod.SCALE,
         max_training_points: int = 40000,
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
diff --git a/src/anomalib/models/image/dfm/lightning_model.py b/src/anomalib/models/image/dfm/lightning_model.py
index cc39b4e398..61595434c9 100644
--- a/src/anomalib/models/image/dfm/lightning_model.py
+++ b/src/anomalib/models/image/dfm/lightning_model.py
@@ -54,7 +54,7 @@ def __init__(
         pca_level: float = 0.97,
         score_type: str = "fre",
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
diff --git a/src/anomalib/models/image/draem/lightning_model.py b/src/anomalib/models/image/draem/lightning_model.py
index 66e87a904b..13fa1346e7 100644
--- a/src/anomalib/models/image/draem/lightning_model.py
+++ b/src/anomalib/models/image/draem/lightning_model.py
@@ -51,7 +51,7 @@ def __init__(
         anomaly_source_path: str | None = None,
         beta: float | tuple[float, float] = (0.1, 1.0),
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
diff --git a/src/anomalib/models/image/dsr/lightning_model.py b/src/anomalib/models/image/dsr/lightning_model.py
index 8aa3de08e2..5f9c1c625d 100644
--- a/src/anomalib/models/image/dsr/lightning_model.py
+++ b/src/anomalib/models/image/dsr/lightning_model.py
@@ -53,7 +53,7 @@ def __init__(
         latent_anomaly_strength: float = 0.2,
         upsampling_train_ratio: float = 0.7,
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
diff --git a/src/anomalib/models/image/efficient_ad/lightning_model.py b/src/anomalib/models/image/efficient_ad/lightning_model.py
index 47ace2a073..af90173c52 100644
--- a/src/anomalib/models/image/efficient_ad/lightning_model.py
+++ b/src/anomalib/models/image/efficient_ad/lightning_model.py
@@ -76,7 +76,7 @@ def __init__(
         padding: bool = False,
         pad_maps: bool = True,
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
diff --git a/src/anomalib/models/image/fastflow/lightning_model.py b/src/anomalib/models/image/fastflow/lightning_model.py
index 35a5f8dddb..9d51f99489 100644
--- a/src/anomalib/models/image/fastflow/lightning_model.py
+++ b/src/anomalib/models/image/fastflow/lightning_model.py
@@ -50,7 +50,7 @@ def __init__(
         conv3x3_only: bool = False,
         hidden_ratio: float = 1.0,
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
diff --git a/src/anomalib/models/image/fre/lightning_model.py b/src/anomalib/models/image/fre/lightning_model.py
index 505beb7f1b..c748705c3a 100755
--- a/src/anomalib/models/image/fre/lightning_model.py
+++ b/src/anomalib/models/image/fre/lightning_model.py
@@ -56,7 +56,7 @@ def __init__(
         input_dim: int = 65536,
         latent_dim: int = 220,
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
diff --git a/src/anomalib/models/image/ganomaly/lightning_model.py b/src/anomalib/models/image/ganomaly/lightning_model.py
index 982bd93d33..cf7d7525d0 100644
--- a/src/anomalib/models/image/ganomaly/lightning_model.py
+++ b/src/anomalib/models/image/ganomaly/lightning_model.py
@@ -71,7 +71,7 @@ def __init__(
         beta1: float = 0.5,
         beta2: float = 0.999,
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
diff --git a/src/anomalib/models/image/padim/lightning_model.py b/src/anomalib/models/image/padim/lightning_model.py
index 28d59fc9eb..d658b1cfa1 100644
--- a/src/anomalib/models/image/padim/lightning_model.py
+++ b/src/anomalib/models/image/padim/lightning_model.py
@@ -50,7 +50,7 @@ def __init__(
         pre_trained: bool = True,
         n_features: int | None = None,
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
@@ -132,6 +132,6 @@ def learning_type(self) -> LearningType:
         return LearningType.ONE_CLASS
 
     @staticmethod
-    def default_post_processor() -> OneClassPostProcessor:
+    def configure_post_processor() -> OneClassPostProcessor:
         """Return the default post-processor for PADIM."""
         return OneClassPostProcessor()
diff --git a/src/anomalib/models/image/patchcore/lightning_model.py b/src/anomalib/models/image/patchcore/lightning_model.py
index d22a97d891..d2fb922da3 100644
--- a/src/anomalib/models/image/patchcore/lightning_model.py
+++ b/src/anomalib/models/image/patchcore/lightning_model.py
@@ -53,7 +53,7 @@ def __init__(
         coreset_sampling_ratio: float = 0.1,
         num_neighbors: int = 9,
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
@@ -154,7 +154,7 @@ def learning_type(self) -> LearningType:
         return LearningType.ONE_CLASS
 
     @staticmethod
-    def default_post_processor() -> OneClassPostProcessor:
+    def configure_post_processor() -> OneClassPostProcessor:
         """Return the default post-processor for the model.
 
         Returns:
diff --git a/src/anomalib/models/image/reverse_distillation/lightning_model.py b/src/anomalib/models/image/reverse_distillation/lightning_model.py
index e052c864cb..4fb6e06b2f 100644
--- a/src/anomalib/models/image/reverse_distillation/lightning_model.py
+++ b/src/anomalib/models/image/reverse_distillation/lightning_model.py
@@ -48,7 +48,7 @@ def __init__(
         anomaly_map_mode: AnomalyMapGenerationMode = AnomalyMapGenerationMode.ADD,
         pre_trained: bool = True,
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
diff --git a/src/anomalib/models/image/stfpm/lightning_model.py b/src/anomalib/models/image/stfpm/lightning_model.py
index 9f37e46239..b94d5b6639 100644
--- a/src/anomalib/models/image/stfpm/lightning_model.py
+++ b/src/anomalib/models/image/stfpm/lightning_model.py
@@ -44,7 +44,7 @@ def __init__(
         backbone: str = "resnet18",
         layers: Sequence[str] = ("layer1", "layer2", "layer3"),
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
diff --git a/src/anomalib/models/image/uflow/lightning_model.py b/src/anomalib/models/image/uflow/lightning_model.py
index 60e7a26752..445cfcd6ee 100644
--- a/src/anomalib/models/image/uflow/lightning_model.py
+++ b/src/anomalib/models/image/uflow/lightning_model.py
@@ -49,7 +49,7 @@ def __init__(
         affine_subnet_channels_ratio: float = 1.0,
         permute_soft: bool = False,
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         """Uflow model.
diff --git a/src/anomalib/models/image/vlm_ad/lightning_model.py b/src/anomalib/models/image/vlm_ad/lightning_model.py
index 6e47e58027..7340474f29 100644
--- a/src/anomalib/models/image/vlm_ad/lightning_model.py
+++ b/src/anomalib/models/image/vlm_ad/lightning_model.py
@@ -102,7 +102,8 @@ def configure_transforms(image_size: tuple[int, int] | None = None) -> None:
         if image_size is not None:
             logger.warning("Ignoring image_size argument as each backend has its own transforms.")
 
-    def default_post_processor(self) -> PostProcessor | None:  # noqa: PLR6301
+    @classmethod
+    def configure_post_processor(cls) -> PostProcessor | None:
         """Post processing is not required for this model."""
         return None
 
diff --git a/src/anomalib/models/image/winclip/lightning_model.py b/src/anomalib/models/image/winclip/lightning_model.py
index d3f7481df7..9f16558619 100644
--- a/src/anomalib/models/image/winclip/lightning_model.py
+++ b/src/anomalib/models/image/winclip/lightning_model.py
@@ -56,7 +56,7 @@ def __init__(
         scales: tuple = (2, 3),
         few_shot_source: Path | str | None = None,
         pre_processor: PreProcessor | bool = True,
-        post_processor: PostProcessor | None = None,
+        post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
     ) -> None:
         super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
@@ -195,6 +195,6 @@ def configure_pre_processor(cls, image_size: tuple[int, int] | None = None) -> P
         return PreProcessor(val_transform=transform, test_transform=transform)
 
     @staticmethod
-    def default_post_processor() -> OneClassPostProcessor:
+    def configure_post_processor() -> OneClassPostProcessor:
         """Return the default post-processor for WinCLIP."""
         return OneClassPostProcessor()
diff --git a/src/anomalib/models/video/ai_vad/lightning_model.py b/src/anomalib/models/video/ai_vad/lightning_model.py
index 750746f259..bbebbe5edf 100644
--- a/src/anomalib/models/video/ai_vad/lightning_model.py
+++ b/src/anomalib/models/video/ai_vad/lightning_model.py
@@ -81,9 +81,10 @@ def __init__(
         n_neighbors_pose: int = 1,
         n_neighbors_deep: int = 1,
         pre_processor: PreProcessor | bool = True,
+        post_processor: PostProcessor | bool = True,
         **kwargs,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, **kwargs)
+        super().__init__(pre_processor=pre_processor, post_processor=post_processor, **kwargs)
         self.model = AiVadModel(
             box_score_thresh=box_score_thresh,
             persons_only=persons_only,
@@ -179,6 +180,6 @@ def configure_pre_processor(cls, image_size: tuple[int, int] | None = None) -> P
         return PreProcessor()  # A pre-processor with no transforms.
 
     @staticmethod
-    def default_post_processor() -> PostProcessor:
+    def configure_post_processor() -> PostProcessor:
         """Return the default post-processor for AI-VAD."""
         return OneClassPostProcessor()
diff --git a/tests/unit/utils/callbacks/visualizer_callback/dummy_lightning_model.py b/tests/unit/utils/callbacks/visualizer_callback/dummy_lightning_model.py
index dc9c31e8db..45389c949f 100644
--- a/tests/unit/utils/callbacks/visualizer_callback/dummy_lightning_model.py
+++ b/tests/unit/utils/callbacks/visualizer_callback/dummy_lightning_model.py
@@ -74,6 +74,6 @@ def learning_type(self) -> LearningType:
         return LearningType.ZERO_SHOT
 
     @staticmethod
-    def default_post_processor() -> PostProcessor:
+    def configure_post_processor() -> PostProcessor:
         """Returns a dummy post-processor."""
         return DummyPostProcessor()

From e11c663f05c24711b1dc3f88c7f11e52d80b6697 Mon Sep 17 00:00:00 2001
From: Samet Akcay <samet.akcay@intel.com>
Date: Mon, 9 Dec 2024 13:38:36 +0000
Subject: [PATCH 2/9] Add visualizer to models

Signed-off-by: Samet Akcay <samet.akcay@intel.com>
---
 .../models/components/base/anomaly_module.py  | 98 +++++++++++++++----
 .../models/image/cfa/lightning_model.py       | 10 +-
 .../models/image/cflow/lightning_model.py     |  9 +-
 .../models/image/csflow/lightning_model.py    |  9 +-
 .../models/image/dfkde/lightning_model.py     |  9 +-
 .../models/image/dfm/lightning_model.py       |  9 +-
 .../models/image/draem/lightning_model.py     |  9 +-
 .../models/image/dsr/lightning_model.py       |  9 +-
 .../image/efficient_ad/lightning_model.py     |  9 +-
 .../models/image/fastflow/lightning_model.py  |  9 +-
 .../models/image/fre/lightning_model.py       |  9 +-
 .../models/image/ganomaly/lightning_model.py  |  9 +-
 .../models/image/padim/lightning_model.py     |  9 +-
 .../models/image/patchcore/lightning_model.py |  9 +-
 .../reverse_distillation/lightning_model.py   |  9 +-
 .../models/image/stfpm/lightning_model.py     |  9 +-
 .../models/image/uflow/lightning_model.py     | 12 ++-
 .../models/image/winclip/lightning_model.py   |  9 +-
 src/anomalib/visualization/__init__.py        |  3 +
 src/anomalib/visualization/base.py            | 14 +++
 .../visualization/image/visualizer.py         |  6 +-
 21 files changed, 240 insertions(+), 38 deletions(-)
 create mode 100644 src/anomalib/visualization/base.py

diff --git a/src/anomalib/models/components/base/anomaly_module.py b/src/anomalib/models/components/base/anomaly_module.py
index c7b0d0d420..dace0a7082 100644
--- a/src/anomalib/models/components/base/anomaly_module.py
+++ b/src/anomalib/models/components/base/anomaly_module.py
@@ -25,6 +25,7 @@
 from anomalib.metrics.threshold import Threshold
 from anomalib.post_processing import OneClassPostProcessor, PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import ImageVisualizer, Visualizer
 
 from .export_mixin import ExportMixin
 
@@ -42,6 +43,7 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
         super().__init__()
         logger.info("Initializing %s model.", self.__class__.__name__)
@@ -54,6 +56,7 @@ def __init__(
         self.pre_processor = self._resolve_pre_processor(pre_processor)
         self.post_processor = self._resolve_post_processor(post_processor)
         self.evaluator = self._resolve_evaluator(evaluator)
+        self.visualizer = self._resolve_visualizer(visualizer)
 
         self._input_size: tuple[int, int] | None = None
         self._is_setup = False
@@ -79,25 +82,6 @@ def _setup(self) -> None:
         initialization.
         """
 
-    def _resolve_pre_processor(self, pre_processor: PreProcessor | bool) -> PreProcessor | None:
-        """Resolve and validate which pre-processor to use..
-
-        Args:
-            pre_processor: Pre-processor configuration
-                - True -> use default pre-processor
-                - False -> no pre-processor
-                - PreProcessor -> use the provided pre-processor
-
-        Returns:
-            Configured pre-processor
-        """
-        if isinstance(pre_processor, PreProcessor):
-            return pre_processor
-        if isinstance(pre_processor, bool):
-            return self.configure_pre_processor() if pre_processor else None
-        msg = f"Invalid pre-processor type: {type(pre_processor)}"
-        raise TypeError(msg)
-
     def configure_callbacks(self) -> Sequence[Callback] | Callback:
         """Configure default callbacks for AnomalibModule."""
         return [self.pre_processor] if self.pre_processor else []
@@ -170,6 +154,25 @@ def learning_type(self) -> LearningType:
         """Learning type of the model."""
         raise NotImplementedError
 
+    def _resolve_pre_processor(self, pre_processor: PreProcessor | bool) -> PreProcessor | None:
+        """Resolve and validate which pre-processor to use..
+
+        Args:
+            pre_processor: Pre-processor configuration
+                - True -> use default pre-processor
+                - False -> no pre-processor
+                - PreProcessor -> use the provided pre-processor
+
+        Returns:
+            Configured pre-processor
+        """
+        if isinstance(pre_processor, PreProcessor):
+            return pre_processor
+        if isinstance(pre_processor, bool):
+            return self.configure_pre_processor() if pre_processor else None
+        msg = f"Invalid pre-processor type: {type(pre_processor)}"
+        raise TypeError(msg)
+
     @classmethod
     def configure_pre_processor(cls, image_size: tuple[int, int] | None = None) -> PreProcessor:
         """Configure the pre-processor.
@@ -289,6 +292,63 @@ def configure_evaluator() -> Evaluator:
         test_metrics = [image_auroc, image_f1score, pixel_auroc, pixel_f1score]
         return Evaluator(test_metrics=test_metrics)
 
+    def _resolve_visualizer(self, visualizer: Visualizer | bool) -> Visualizer | None:
+        """Resolve and validate which visualizer to use.
+
+        Args:
+            visualizer: Visualizer configuration
+                - True -> use default visualizer
+                - False -> no visualizer
+                - Visualizer -> use the provided visualizer
+
+        Returns:
+            Configured visualizer
+        """
+        if isinstance(visualizer, Visualizer):
+            return visualizer
+        if isinstance(visualizer, bool):
+            return self.configure_visualizer() if visualizer else None
+        msg = f"Visualizer must be of type Visualizer or bool, got {type(visualizer)}"
+        raise TypeError(msg)
+
+    @classmethod
+    def configure_visualizer(cls) -> ImageVisualizer:
+        """Configure the default visualizer.
+
+        By default, this method returns an ImageVisualizer instance, which is suitable for
+        visualizing image-based anomaly detection results. However, the visualizer can be
+        customized based on your needs - for example, using VideoVisualizer for video data
+        or implementing a custom visualizer for specific visualization requirements.
+
+        Returns:
+            Visualizer: Configured visualizer instance (ImageVisualizer by default).
+
+        Examples:
+            Get default ImageVisualizer:
+
+            >>> visualizer = AnomalibModule.configure_visualizer()
+
+            Create model with VideoVisualizer:
+
+            >>> from custom_module import VideoVisualizer
+            >>> video_visualizer = VideoVisualizer()
+            >>> model = PatchCore(visualizer=video_visualizer)
+
+            Create model with custom visualizer:
+
+            >>> class CustomVisualizer(Visualizer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         # Custom visualization logic
+            >>> custom_visualizer = CustomVisualizer()
+            >>> model = PatchCore(visualizer=custom_visualizer)
+
+            Disable visualization:
+
+            >>> model = PatchCore(visualizer=False)
+        """
+        return ImageVisualizer()
+
     @property
     def input_size(self) -> tuple[int, int] | None:
         """Return the effective input size of the model.
diff --git a/src/anomalib/models/image/cfa/lightning_model.py b/src/anomalib/models/image/cfa/lightning_model.py
index 4c2a9cffe2..9eed15b6a7 100644
--- a/src/anomalib/models/image/cfa/lightning_model.py
+++ b/src/anomalib/models/image/cfa/lightning_model.py
@@ -20,6 +20,7 @@
 from anomalib.models.components import AnomalibModule
 from anomalib.post_processing import PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .loss import CfaLoss
 from .torch_model import CfaModel
@@ -58,11 +59,18 @@ def __init__(
         num_nearest_neighbors: int = 3,
         num_hard_negative_features: int = 3,
         radius: float = 1e-5,
+        # Anomalib's Auxiliary Components
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
         self.model: CfaModel = CfaModel(
             backbone=backbone,
             gamma_c=gamma_c,
diff --git a/src/anomalib/models/image/cflow/lightning_model.py b/src/anomalib/models/image/cflow/lightning_model.py
index f82872cb84..4dd9c25850 100644
--- a/src/anomalib/models/image/cflow/lightning_model.py
+++ b/src/anomalib/models/image/cflow/lightning_model.py
@@ -27,6 +27,7 @@
 from anomalib.models.components import AnomalibModule
 from anomalib.post_processing import PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .torch_model import CflowModel
 from .utils import get_logp, positional_encoding_2d
@@ -73,8 +74,14 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
 
         self.model: CflowModel = CflowModel(
             backbone=backbone,
diff --git a/src/anomalib/models/image/csflow/lightning_model.py b/src/anomalib/models/image/csflow/lightning_model.py
index 3324972422..8e9994631a 100644
--- a/src/anomalib/models/image/csflow/lightning_model.py
+++ b/src/anomalib/models/image/csflow/lightning_model.py
@@ -18,6 +18,7 @@
 from anomalib.models.components import AnomalibModule
 from anomalib.post_processing import PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .loss import CsFlowLoss
 from .torch_model import CsFlowModel
@@ -50,8 +51,14 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
         if self.input_size is None:
             msg = "CsFlow needs input size to build torch model."
             raise ValueError(msg)
diff --git a/src/anomalib/models/image/dfkde/lightning_model.py b/src/anomalib/models/image/dfkde/lightning_model.py
index be141c45ef..16ccac6403 100644
--- a/src/anomalib/models/image/dfkde/lightning_model.py
+++ b/src/anomalib/models/image/dfkde/lightning_model.py
@@ -17,6 +17,7 @@
 from anomalib.models.components.classification import FeatureScalingMethod
 from anomalib.post_processing import PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .torch_model import DfkdeModel
 
@@ -52,8 +53,14 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
 
         self.model = DfkdeModel(
             layers=layers,
diff --git a/src/anomalib/models/image/dfm/lightning_model.py b/src/anomalib/models/image/dfm/lightning_model.py
index 61595434c9..96a4388835 100644
--- a/src/anomalib/models/image/dfm/lightning_model.py
+++ b/src/anomalib/models/image/dfm/lightning_model.py
@@ -18,6 +18,7 @@
 from anomalib.models.components import AnomalibModule, MemoryBankMixin
 from anomalib.post_processing import PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .torch_model import DFMModel
 
@@ -56,8 +57,14 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
 
         self.model: DFMModel = DFMModel(
             backbone=backbone,
diff --git a/src/anomalib/models/image/draem/lightning_model.py b/src/anomalib/models/image/draem/lightning_model.py
index 13fa1346e7..84b143f3f5 100644
--- a/src/anomalib/models/image/draem/lightning_model.py
+++ b/src/anomalib/models/image/draem/lightning_model.py
@@ -21,6 +21,7 @@
 from anomalib.models.components import AnomalibModule
 from anomalib.post_processing import PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .loss import DraemLoss
 from .torch_model import DraemModel
@@ -53,8 +54,14 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
 
         self.augmenter = PerlinAnomalyGenerator(anomaly_source_path=anomaly_source_path, blend_factor=beta)
         self.model = DraemModel(sspcab=enable_sspcab)
diff --git a/src/anomalib/models/image/dsr/lightning_model.py b/src/anomalib/models/image/dsr/lightning_model.py
index 5f9c1c625d..dd80e88ba7 100644
--- a/src/anomalib/models/image/dsr/lightning_model.py
+++ b/src/anomalib/models/image/dsr/lightning_model.py
@@ -25,6 +25,7 @@
 from anomalib.models.image.dsr.torch_model import DsrModel
 from anomalib.post_processing import PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 __all__ = ["Dsr"]
 
@@ -55,8 +56,14 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
 
         self.automatic_optimization = False
         self.upsampling_train_ratio = upsampling_train_ratio
diff --git a/src/anomalib/models/image/efficient_ad/lightning_model.py b/src/anomalib/models/image/efficient_ad/lightning_model.py
index af90173c52..aa99d6a439 100644
--- a/src/anomalib/models/image/efficient_ad/lightning_model.py
+++ b/src/anomalib/models/image/efficient_ad/lightning_model.py
@@ -24,6 +24,7 @@
 from anomalib.models.components import AnomalibModule
 from anomalib.post_processing import PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .torch_model import EfficientAdModel, EfficientAdModelSize, reduce_tensor_elems
 
@@ -78,8 +79,14 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
 
         self.imagenet_dir = Path(imagenet_dir)
         if not isinstance(model_size, EfficientAdModelSize):
diff --git a/src/anomalib/models/image/fastflow/lightning_model.py b/src/anomalib/models/image/fastflow/lightning_model.py
index 9d51f99489..8a98ea9e7a 100644
--- a/src/anomalib/models/image/fastflow/lightning_model.py
+++ b/src/anomalib/models/image/fastflow/lightning_model.py
@@ -18,6 +18,7 @@
 from anomalib.models.components import AnomalibModule
 from anomalib.post_processing import PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .loss import FastflowLoss
 from .torch_model import FastflowModel
@@ -52,8 +53,14 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
         if self.input_size is None:
             msg = "Fastflow needs input size to build torch model."
             raise ValueError(msg)
diff --git a/src/anomalib/models/image/fre/lightning_model.py b/src/anomalib/models/image/fre/lightning_model.py
index c748705c3a..953fcd4322 100755
--- a/src/anomalib/models/image/fre/lightning_model.py
+++ b/src/anomalib/models/image/fre/lightning_model.py
@@ -19,6 +19,7 @@
 from anomalib.models.components import AnomalibModule
 from anomalib.post_processing import PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .torch_model import FREModel
 
@@ -58,8 +59,14 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
 
         self.model: FREModel = FREModel(
             backbone=backbone,
diff --git a/src/anomalib/models/image/ganomaly/lightning_model.py b/src/anomalib/models/image/ganomaly/lightning_model.py
index cf7d7525d0..4b48b0b633 100644
--- a/src/anomalib/models/image/ganomaly/lightning_model.py
+++ b/src/anomalib/models/image/ganomaly/lightning_model.py
@@ -19,6 +19,7 @@
 from anomalib.models.components import AnomalibModule
 from anomalib.post_processing import PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .loss import DiscriminatorLoss, GeneratorLoss
 from .torch_model import GanomalyModel
@@ -73,8 +74,14 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
         if self.input_size is None:
             msg = "GANomaly needs input size to build torch model."
             raise ValueError(msg)
diff --git a/src/anomalib/models/image/padim/lightning_model.py b/src/anomalib/models/image/padim/lightning_model.py
index d658b1cfa1..4a223c9e62 100644
--- a/src/anomalib/models/image/padim/lightning_model.py
+++ b/src/anomalib/models/image/padim/lightning_model.py
@@ -17,6 +17,7 @@
 from anomalib.models.components import AnomalibModule, MemoryBankMixin
 from anomalib.post_processing import OneClassPostProcessor, PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .torch_model import PadimModel
 
@@ -52,8 +53,14 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
 
         self.model: PadimModel = PadimModel(
             backbone=backbone,
diff --git a/src/anomalib/models/image/patchcore/lightning_model.py b/src/anomalib/models/image/patchcore/lightning_model.py
index d2fb922da3..689b7ac81f 100644
--- a/src/anomalib/models/image/patchcore/lightning_model.py
+++ b/src/anomalib/models/image/patchcore/lightning_model.py
@@ -20,6 +20,7 @@
 from anomalib.models.components import AnomalibModule, MemoryBankMixin
 from anomalib.post_processing import OneClassPostProcessor, PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .torch_model import PatchcoreModel
 
@@ -55,8 +56,14 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
 
         self.model: PatchcoreModel = PatchcoreModel(
             backbone=backbone,
diff --git a/src/anomalib/models/image/reverse_distillation/lightning_model.py b/src/anomalib/models/image/reverse_distillation/lightning_model.py
index 4fb6e06b2f..3eb3bf903c 100644
--- a/src/anomalib/models/image/reverse_distillation/lightning_model.py
+++ b/src/anomalib/models/image/reverse_distillation/lightning_model.py
@@ -18,6 +18,7 @@
 from anomalib.models.components import AnomalibModule
 from anomalib.post_processing import PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .anomaly_map import AnomalyMapGenerationMode
 from .loss import ReverseDistillationLoss
@@ -50,8 +51,14 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
         if self.input_size is None:
             msg = "Input size is required for Reverse Distillation model."
             raise ValueError(msg)
diff --git a/src/anomalib/models/image/stfpm/lightning_model.py b/src/anomalib/models/image/stfpm/lightning_model.py
index b94d5b6639..f3daafe407 100644
--- a/src/anomalib/models/image/stfpm/lightning_model.py
+++ b/src/anomalib/models/image/stfpm/lightning_model.py
@@ -19,6 +19,7 @@
 from anomalib.models.components import AnomalibModule
 from anomalib.post_processing import PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .loss import STFPMLoss
 from .torch_model import STFPMModel
@@ -46,8 +47,14 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
 
         self.model = STFPMModel(backbone=backbone, layers=layers)
         self.loss = STFPMLoss()
diff --git a/src/anomalib/models/image/uflow/lightning_model.py b/src/anomalib/models/image/uflow/lightning_model.py
index 445cfcd6ee..bfd51195ca 100644
--- a/src/anomalib/models/image/uflow/lightning_model.py
+++ b/src/anomalib/models/image/uflow/lightning_model.py
@@ -21,6 +21,7 @@
 from anomalib.models.components import AnomalibModule
 from anomalib.post_processing import PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .loss import UFlowLoss
 from .torch_model import UflowModel
@@ -51,6 +52,7 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
         """Uflow model.
 
@@ -69,8 +71,16 @@ def __init__(
             evaluator (Evaluator, optional): Evaluator for the model.
                 This is used to evaluate the model.
                 Defaults to ``True``.
+            visualizer (Visualizer, optional): Visualizer for the model.
+                This is used to visualize the model.
+                Defaults to ``True``.
         """
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
         if self.input_size is None:
             msg = "Input size is required for UFlow model."
             raise ValueError(msg)
diff --git a/src/anomalib/models/image/winclip/lightning_model.py b/src/anomalib/models/image/winclip/lightning_model.py
index 9f16558619..23a7cf23a1 100644
--- a/src/anomalib/models/image/winclip/lightning_model.py
+++ b/src/anomalib/models/image/winclip/lightning_model.py
@@ -22,6 +22,7 @@
 from anomalib.models.components import AnomalibModule
 from anomalib.post_processing import OneClassPostProcessor, PostProcessor
 from anomalib.pre_processing import PreProcessor
+from anomalib.visualization import Visualizer
 
 from .torch_model import WinClipModel
 
@@ -58,8 +59,14 @@ def __init__(
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
         evaluator: Evaluator | bool = True,
+        visualizer: Visualizer | bool = True,
     ) -> None:
-        super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator)
+        super().__init__(
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            evaluator=evaluator,
+            visualizer=visualizer,
+        )
 
         self.model = WinClipModel(scales=scales, apply_transform=False)
         self.class_name = class_name
diff --git a/src/anomalib/visualization/__init__.py b/src/anomalib/visualization/__init__.py
index ca0b7bc138..989f4cc34c 100644
--- a/src/anomalib/visualization/__init__.py
+++ b/src/anomalib/visualization/__init__.py
@@ -3,10 +3,13 @@
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
+from .base import Visualizer
 from .image import ImageVisualizer, visualize_anomaly_map, visualize_mask
 from .image.item_visualizer import visualize_image_item
 
 __all__ = [
+    # Base visualizer class
+    "Visualizer",
     # Image visualizer class
     "ImageVisualizer",
     # Image visualization functions
diff --git a/src/anomalib/visualization/base.py b/src/anomalib/visualization/base.py
new file mode 100644
index 0000000000..dc49a85401
--- /dev/null
+++ b/src/anomalib/visualization/base.py
@@ -0,0 +1,14 @@
+"""Base Visualizer."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from lightning.pytorch import Callback
+
+
+class Visualizer(Callback):
+    """Base class for all visualizers.
+
+    In Anomalib, the visualizer is used to visualize the results of the model
+    during the testing and prediction phases.
+    """
diff --git a/src/anomalib/visualization/image/visualizer.py b/src/anomalib/visualization/image/visualizer.py
index c8eafc8ae8..97d11546ea 100644
--- a/src/anomalib/visualization/image/visualizer.py
+++ b/src/anomalib/visualization/image/visualizer.py
@@ -6,11 +6,12 @@
 from pathlib import Path
 from typing import Any
 
-from lightning.pytorch import Callback, Trainer
+from lightning.pytorch import Trainer
 
 from anomalib.data import ImageBatch
 from anomalib.models import AnomalibModule
 from anomalib.utils.path import generate_output_filename
+from anomalib.visualization.base import Visualizer
 
 from .item_visualizer import (
     DEFAULT_FIELDS_CONFIG,
@@ -20,7 +21,7 @@
 )
 
 
-class ImageVisualizer(Callback):
+class ImageVisualizer(Visualizer):
     """Image Visualizer.
 
     This class is responsible for visualizing images and their corresponding anomaly maps
@@ -127,6 +128,7 @@ def __init__(
         text_config: dict[str, Any] | None = None,
         output_dir: str | Path | None = None,
     ) -> None:
+        super().__init__()
         self.fields = fields or ["image", "gt_mask"]
         self.overlay_fields = overlay_fields or [("image", ["anomaly_map"]), ("image", ["pred_mask"])]
         self.field_size = field_size

From fab2bb64d6f06d18b1a63cdc93f15d0a327aef32 Mon Sep 17 00:00:00 2001
From: Samet Akcay <samet.akcay@intel.com>
Date: Mon, 9 Dec 2024 13:42:46 +0000
Subject: [PATCH 3/9] Refactor `configure_callbacks` method in `AnomalibModule`
 to dynamically include available callbacks (pre-processor, post-processor,
 evaluator, visualizer) based on their existence. Update docstring for clarity
 on return values.

---
 .../models/components/base/anomaly_module.py      | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/anomalib/models/components/base/anomaly_module.py b/src/anomalib/models/components/base/anomaly_module.py
index dace0a7082..6f42023e8b 100644
--- a/src/anomalib/models/components/base/anomaly_module.py
+++ b/src/anomalib/models/components/base/anomaly_module.py
@@ -83,8 +83,19 @@ def _setup(self) -> None:
         """
 
     def configure_callbacks(self) -> Sequence[Callback] | Callback:
-        """Configure default callbacks for AnomalibModule."""
-        return [self.pre_processor] if self.pre_processor else []
+        """Configure default callbacks for AnomalibModule.
+
+        Returns:
+            List of callbacks that includes the pre-processor, post-processor, evaluator,
+            and visualizer if they are available and inherit from Callback.
+        """
+        callbacks: list[Callback] = []
+        callbacks.extend(
+            component
+            for component in (self.pre_processor, self.post_processor, self.evaluator, self.visualizer)
+            if isinstance(component, Callback)
+        )
+        return callbacks
 
     def forward(self, batch: torch.Tensor, *args, **kwargs) -> InferenceBatch:
         """Perform the forward-pass by passing input tensor to the module.

From b3c656d4cce748c9db8ddcb18ddeba04250a0715 Mon Sep 17 00:00:00 2001
From: Samet Akcay <samet.akcay@intel.com>
Date: Mon, 9 Dec 2024 13:43:41 +0000
Subject: [PATCH 4/9] Refactor callback handling in `Engine` class by removing
 post-processor and metrics callbacks. Simplify callback configuration to
 streamline the process and enhance maintainability.

---
 src/anomalib/engine/engine.py | 17 ++---------------
 1 file changed, 2 insertions(+), 15 deletions(-)

diff --git a/src/anomalib/engine/engine.py b/src/anomalib/engine/engine.py
index fe823f3729..ecd0a4f062 100644
--- a/src/anomalib/engine/engine.py
+++ b/src/anomalib/engine/engine.py
@@ -22,7 +22,6 @@
 from anomalib.deploy import CompressionType, ExportType
 from anomalib.models import AnomalibModule
 from anomalib.utils.path import create_versioned_dir
-from anomalib.visualization import ImageVisualizer
 
 logger = logging.getLogger(__name__)
 
@@ -258,7 +257,7 @@ def _setup_trainer(self, model: AnomalibModule) -> None:
             self._cache.update(model)
 
         # Setup anomalib callbacks to be used with the trainer
-        self._setup_anomalib_callbacks(model)
+        self._setup_anomalib_callbacks()
 
         # Temporarily set devices to 1 to avoid issues with multiple processes
         self._cache.args["devices"] = 1
@@ -267,7 +266,7 @@ def _setup_trainer(self, model: AnomalibModule) -> None:
         if self._trainer is None:
             self._trainer = Trainer(**self._cache.args)
 
-    def _setup_anomalib_callbacks(self, model: AnomalibModule) -> None:
+    def _setup_anomalib_callbacks(self) -> None:
         """Set up callbacks for the trainer."""
         _callbacks: list[Callback] = []
 
@@ -282,18 +281,6 @@ def _setup_anomalib_callbacks(self, model: AnomalibModule) -> None:
                 ),
             )
 
-        # Add the post-processor callback.
-        if isinstance(model.post_processor, Callback):
-            _callbacks.append(model.post_processor)
-
-        # Add the metrics callback.
-        if isinstance(model.evaluator, Callback):
-            _callbacks.append(model.evaluator)
-
-        # Add the image visualizer callback if it is passed by the user.
-        if not any(isinstance(callback, ImageVisualizer) for callback in self._cache.args["callbacks"]):
-            _callbacks.append(ImageVisualizer())
-
         _callbacks.append(TimerCallback())
 
         # Combine the callbacks, and update the trainer callbacks.

From 28e951ffd380edd8c1df2be69c91398820417ac1 Mon Sep 17 00:00:00 2001
From: Samet Akcay <samet.akcay@intel.com>
Date: Mon, 9 Dec 2024 13:53:34 +0000
Subject: [PATCH 5/9] Fix circular import issue

Signed-off-by: Samet Akcay <samet.akcay@intel.com>
---
 src/anomalib/visualization/image/visualizer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/anomalib/visualization/image/visualizer.py b/src/anomalib/visualization/image/visualizer.py
index 97d11546ea..69bc088724 100644
--- a/src/anomalib/visualization/image/visualizer.py
+++ b/src/anomalib/visualization/image/visualizer.py
@@ -9,7 +9,7 @@
 from lightning.pytorch import Trainer
 
 from anomalib.data import ImageBatch
-from anomalib.models import AnomalibModule
+from anomalib.models.components.base import AnomalibModule
 from anomalib.utils.path import generate_output_filename
 from anomalib.visualization.base import Visualizer
 

From b86eb33af17d96d19a2e5a1c089c60bfcd4b47f6 Mon Sep 17 00:00:00 2001
From: Samet Akcay <samet.akcay@intel.com>
Date: Mon, 9 Dec 2024 14:39:07 +0000
Subject: [PATCH 6/9] Fix circular import issue

Signed-off-by: Samet Akcay <samet.akcay@intel.com>
---
 .../visualization/image/visualizer.py         | 27 ++++++++++---------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/src/anomalib/visualization/image/visualizer.py b/src/anomalib/visualization/image/visualizer.py
index 69bc088724..c230bdde03 100644
--- a/src/anomalib/visualization/image/visualizer.py
+++ b/src/anomalib/visualization/image/visualizer.py
@@ -4,12 +4,15 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
-from lightning.pytorch import Trainer
+# Only import types during type checking to avoid circular imports
+if TYPE_CHECKING:
+    from lightning.pytorch import Trainer
+
+    from anomalib.data import ImageBatch
+    from anomalib.models import AnomalibModule
 
-from anomalib.data import ImageBatch
-from anomalib.models.components.base import AnomalibModule
 from anomalib.utils.path import generate_output_filename
 from anomalib.visualization.base import Visualizer
 
@@ -139,10 +142,10 @@ def __init__(
 
     def on_test_batch_end(
         self,
-        trainer: Trainer,
-        pl_module: AnomalibModule,
-        outputs: ImageBatch,
-        batch: ImageBatch,
+        trainer: "Trainer",
+        pl_module: "AnomalibModule",
+        outputs: "ImageBatch",
+        batch: "ImageBatch",
         batch_idx: int,
         dataloader_idx: int = 0,
     ) -> None:
@@ -177,10 +180,10 @@ def on_test_batch_end(
 
     def on_predict_batch_end(
         self,
-        trainer: Trainer,
-        pl_module: AnomalibModule,
-        outputs: ImageBatch,
-        batch: ImageBatch,
+        trainer: "Trainer",
+        pl_module: "AnomalibModule",
+        outputs: "ImageBatch",
+        batch: "ImageBatch",
         batch_idx: int,
         dataloader_idx: int = 0,
     ) -> None:

From bca8b324689a75d31760c34e5d5ca9a8d7d6f0c2 Mon Sep 17 00:00:00 2001
From: Samet Akcay <samet.akcay@intel.com>
Date: Mon, 9 Dec 2024 15:33:04 +0000
Subject: [PATCH 7/9] Convert `configure_post_processor` to a method

Signed-off-by: Samet Akcay <samet.akcay@intel.com>
---
 src/anomalib/models/components/base/anomaly_module.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/anomalib/models/components/base/anomaly_module.py b/src/anomalib/models/components/base/anomaly_module.py
index 6f42023e8b..3fd5557032 100644
--- a/src/anomalib/models/components/base/anomaly_module.py
+++ b/src/anomalib/models/components/base/anomaly_module.py
@@ -247,8 +247,7 @@ def _resolve_post_processor(self, post_processor: PostProcessor | bool) -> PostP
         msg = f"Invalid post-processor type: {type(post_processor)}"
         raise TypeError(msg)
 
-    @classmethod
-    def configure_post_processor(cls) -> PostProcessor | None:
+    def configure_post_processor(self) -> PostProcessor | None:
         """Configure the default post-processor based on the learning type.
 
         Returns:
@@ -271,10 +270,12 @@ def configure_post_processor(cls) -> PostProcessor | None:
 
             >>> model = PatchCore(post_processor=False)
         """
-        if cls.learning_type == LearningType.ONE_CLASS:
+        if self.learning_type == LearningType.ONE_CLASS:
             return OneClassPostProcessor()
-        msg = f"No default post-processor available for model with learning type {cls.learning_type}. \
-              Please override the configure_post_processor method in the model implementation."
+        msg = (
+            f"No default post-processor available for model with learning type {self.learning_type}. "
+            "Please override the configure_post_processor method in the model implementation."
+        )
         raise NotImplementedError(msg)
 
     def _resolve_evaluator(self, evaluator: Evaluator | bool) -> Evaluator | None:

From e2472dbbd06643c0c3361dd59b0a8d71735c2f57 Mon Sep 17 00:00:00 2001
From: Samet Akcay <samet.akcay@intel.com>
Date: Mon, 9 Dec 2024 17:59:01 +0000
Subject: [PATCH 8/9] Fix upgrade tests

Signed-off-by: Samet Akcay <samet.akcay@intel.com>
---
 tests/integration/tools/upgrade/expected_draem_v1.yaml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/integration/tools/upgrade/expected_draem_v1.yaml b/tests/integration/tools/upgrade/expected_draem_v1.yaml
index 438c49fd73..feb059214d 100644
--- a/tests/integration/tools/upgrade/expected_draem_v1.yaml
+++ b/tests/integration/tools/upgrade/expected_draem_v1.yaml
@@ -21,8 +21,9 @@ model:
       - 0.1
       - 1.0
     pre_processor: true
-    post_processor: null
+    post_processor: true
     evaluator: true
+    visualizer: true
 normalization:
   normalization_method: min_max
 metrics:

From ca22800c109298b5aa9ef3fcf8132307f6d8bb40 Mon Sep 17 00:00:00 2001
From: Samet Akcay <samet.akcay@intel.com>
Date: Tue, 10 Dec 2024 10:59:38 +0000
Subject: [PATCH 9/9] Rename anomaly_module to anomalib_module

Signed-off-by: Samet Akcay <samet.akcay@intel.com>
---
 src/anomalib/models/components/base/__init__.py                 | 2 +-
 .../components/base/{anomaly_module.py => anomalib_module.py}   | 0
 2 files changed, 1 insertion(+), 1 deletion(-)
 rename src/anomalib/models/components/base/{anomaly_module.py => anomalib_module.py} (100%)

diff --git a/src/anomalib/models/components/base/__init__.py b/src/anomalib/models/components/base/__init__.py
index 5214f966dc..250eec5045 100644
--- a/src/anomalib/models/components/base/__init__.py
+++ b/src/anomalib/models/components/base/__init__.py
@@ -3,7 +3,7 @@
 # Copyright (C) 2022-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-from .anomaly_module import AnomalibModule
+from .anomalib_module import AnomalibModule
 from .buffer_list import BufferListMixin
 from .dynamic_buffer import DynamicBufferMixin
 from .memory_bank_module import MemoryBankMixin
diff --git a/src/anomalib/models/components/base/anomaly_module.py b/src/anomalib/models/components/base/anomalib_module.py
similarity index 100%
rename from src/anomalib/models/components/base/anomaly_module.py
rename to src/anomalib/models/components/base/anomalib_module.py