🔀 Merge main to v1 (#1652)

* Address tiler issues (#1411) * fix tiler * deprecate random tile locations * restore random tiling in tile method * check tiling section in config * disable tiling for ganomalu * pad -> padding * Refactor Reverse Distillation to match official code (#1389) * Non-mandatory early stopping * Added conv4 and bn4 to OCBE * Loss as in the official code (flattened arrays) * Added comment on how to use torchvision model as an encoder to reproduce results in the paper * Remove early stop from config, change default anomaly_map_mode to add * pre-commit fix * Updated results * Update src/anomalib/models/reverse_distillation/README.md Co-authored-by: Samet Akcay <[email protected]> * Update src/anomalib/models/reverse_distillation/README.md Co-authored-by: Samet Akcay <[email protected]> * Update src/anomalib/models/reverse_distillation/README.md Co-authored-by: Samet Akcay <[email protected]> * Remove early_stopping * Update src/anomalib/models/reverse_distillation/lightning_model.py Co-authored-by: Samet Akcay <[email protected]> * Easier to read code --------- Co-authored-by: Samet Akcay <[email protected]> * Patch for the WinError183 on the OpenVino export mode (#1386) * Fix WinError183 (Windows Error) * Add commentary of the change --------- Co-authored-by: Youho99 <[email protected]> * Add DSR model (#1142) * added license, init.py and draft readme * added draft DSR files * minor comment update * Implemented dsr model + comments * added dsr discrete model * added defect generation in torch model + dsr to list of existing methods in init.py * fixed torch model, started implementing lightning model, implemented anomaly generator * added loss file for DSR * Added loss, improved lightning module * Finished up global implementation of DSR second phase * minor fixes * Bugfixes * Fixed DSR loss calculation * on_training_start -> on_train_start * pre-commit run * updated DSR documentation * reset config file * added automatic pretraining weight download * testing pretrained weights. fixed embedding size in upsampling module and image recon module, to be fixed in original branch * successful testing on pretrained dsr weights * checked test quality with pretrained weights, fixed anomaly score calculation * training is functional * Fixed training procedure * test still working * working upsampling module training and testing * fixed minor bugs * updated documentation * added tests and doc * adapted learning schedule to steps * Update src/anomalib/models/dsr/anomaly_generator.py Co-authored-by: Ashwin Vaidya <[email protected]> * Apply suggestions from code review Co-authored-by: Samet Akcay <[email protected]> Co-authored-by: Ashwin Vaidya <[email protected]> * refactored outputs into dicts * remove super() args * changed downloading weights from anomalib releases + minor fixes * pre commit hooks + minor fixes * removed configurable ckpt path refs + default iteration nb from paper * cleaned up dsr.rst and turned exceptions into RuntimeErrors * Added upsampling ratio parameter to set third training phase epochs * Added batched evalaution + minor code simplification * pre commit hooks * squeeze output image score tensor * readded new path check in efficient ad * fixed double step count with manual optimization * fixed trailing whitespace * Fix black issues * Apply suggestions from code review Co-authored-by: Samet Akcay <[email protected]> * review suggestions * updated architecture image links * Address mypy * changed output types for dsr model * readded dict outputs, adapted to TorchInferencer * fixed error in output dict * removed default imagenet norm --------- Co-authored-by: Samet Akcay <[email protected]> Co-authored-by: Ashwin Vaidya <[email protected]> * Fix unexpected key pixel_metrics.AUPRO.fpr_limit (#1055) * fix unexpected key pixel_metrics.AUPRO.fpr_limit Signed-off-by: FanJiangIntel <[email protected]> * load AUPRO before create_metric_collection Signed-off-by: FanJiangIntel <[email protected]> * code refine Signed-off-by: FanJiangIntel <[email protected]> * fix comment Signed-off-by: FanJiangIntel <[email protected]> * fix Signed-off-by: FanJiangIntel <[email protected]> * Support test Signed-off-by: Kang Wenjing <[email protected]> * Update test Signed-off-by: Kang Wenjing <[email protected]> * Update test Signed-off-by: Kang Wenjing <[email protected]> --------- Signed-off-by: FanJiangIntel <[email protected]> Signed-off-by: Kang Wenjing <[email protected]> Co-authored-by: FanJiangIntel <[email protected]> Co-authored-by: Samet Akcay <[email protected]> * Improved speed and memory usage of mean+std calculation (#1457) * preexisting OpenCV version check added to `setup.py`, ran formatting pre-commit hooks on previous contribution. (#1424) * testing upstream switch * picked up on stale OpenCV `setup.py` issue #1041 * 🐞 Hotfix: Limit Gradio Version (#1458) * Fix metadata path * Ignore hidden directories in folder dataset * Add check for mask_dir for segmentation tasks in Folder dataset * Limit the gradio version to <4 * Fix/efficient ad normalize before every validation (#1441) * Normalize anomaly maps before every validation * Remove print statement --------- Co-authored-by: Samet Akcay <[email protected]> * Fix DRAEM (#1431) * Fix beta in augmenter * Add scheduler * Change normalization to none * Replace two lr schedulers with MultiStepLR * Revert change to beta * Disable early stopping default * Format config * Add opacity parameter beta to config * Adding U-Flow method (#1415) * Added uflow model * Added documentation (README) for uflow model * Added uflow to the list of available models, and main README updated * Added missing images for the documentation * Update src/anomalib/models/uflow/anomaly_map.py Co-authored-by: Samet Akcay <[email protected]> * Update src/anomalib/models/uflow/anomaly_map.py Co-authored-by: Samet Akcay <[email protected]> * Update src/anomalib/models/uflow/feature_extraction.py Co-authored-by: Samet Akcay <[email protected]> * Update src/anomalib/models/uflow/torch_model.py Co-authored-by: Samet Akcay <[email protected]> * Added uflow to the reference guide in docs * Added uflow to the pre-merge tests * removed the _step function, and merged the code with training_step * added as a comment the values used in the paper * re-factorized feature extractors to use the TimmFeatureExtractor class * added annotations for some functions, where the flow graph is created * updated readme to fix images loading * Added link in the README to the original code for reproducing the results * Removed unused kwargs * Added docstrigs with args explanations to UFlow classes * Added models in a github release, and linked here * Passing all pre-commit checks * Changed freia's AllInOneBlock by Anomalib's version, and converted the subnet contructor to a Class, in order to be pickable, that is needed to export the model to torch --------- Co-authored-by: Samet Akcay <[email protected]> * Update README.md * 📘 Announce anomalib v1 on the main `README.md` (#1542) * Fix metadata path * Ignore hidden directories in folder dataset * Add check for mask_dir for segmentation tasks in Folder dataset * Limit the gradio version to <4 * Announce anomalib v1 on readme * Add the installation instructions and update the documentation link * Fixed DSR (#1486) * fixed DSR squeeze bug * added comment * Refactor/extensions custom dataset (#1562) * Explanation how to use extension names in the config file * Added information about extensions to the error message and control of the user input * Easier to read code * Replacing assert with raise * 📚 Modify the PR template (#1611) Update pull_request_template.md * Fix result image URLs (#1510) * Fix tests * refactor path + fix issues + fix linting issues * Migrate docs * fix typing * fix failing model tests * Fix tests * Address PR comments * Fixed shape error, allowing arbitary image sizes for EfficientAD (#1537) * Fixed shape error, allowing arbitrary image sizes. Replaced integer parsing by floor operation * Replaced calculation by ceil operation. Solution of shape error is to round up and not down for the last upsample layer * Add comment for ceil oepration * Formatting with pre-commit hook * Clean up badge --------- Signed-off-by: FanJiangIntel <[email protected]> Signed-off-by: Kang Wenjing <[email protected]> Co-authored-by: Dick Ameln <[email protected]> Co-authored-by: abc-125 <[email protected]> Co-authored-by: Samet Akcay <[email protected]> Co-authored-by: ggiret-thinkdeep <[email protected]> Co-authored-by: Youho99 <[email protected]> Co-authored-by: Philippe Carvalho <[email protected]> Co-authored-by: Wenjing Kang <[email protected]> Co-authored-by: FanJiangIntel <[email protected]> Co-authored-by: belfner <[email protected]> Co-authored-by: Abdulla Al Blooshi <[email protected]> Co-authored-by: Blaž Rolih <[email protected]> Co-authored-by: Matías Tailanian <[email protected]> Co-authored-by: Jan Schlüter <[email protected]> Co-authored-by: Ashwin Vaidya <[email protected]> Co-authored-by: Christopher <[email protected]>
openvinotoolkit · Jan 23, 2024 · 7c1e3dc · 7c1e3dc
1 parent 926bae2
commit 7c1e3dc
Show file tree

Hide file tree

Showing 54 changed files with 3,159 additions and 56 deletions.
diff --git a/README.md b/README.md
@@ -14,7 +14,7 @@
 [![python](https://img.shields.io/badge/python-3.7%2B-green)]()
 [![pytorch](https://img.shields.io/badge/pytorch-1.8.1%2B-orange)]()
 [![openvino](https://img.shields.io/badge/openvino-2022.3.0-purple)]()
-<br/>
+
 [![Pre-Merge Checks](https://github.com/openvinotoolkit/anomalib/actions/workflows/pre_merge.yml/badge.svg)](https://github.com/openvinotoolkit/anomalib/actions/workflows/pre_merge.yml)
 [![Documentation Status](https://readthedocs.org/projects/anomalib/badge/?version=latest)](https://anomalib.readthedocs.io/en/latest/?badge=latest)
 [![codecov](https://codecov.io/gh/openvinotoolkit/anomalib/branch/main/graph/badge.svg?token=Z6A07N1BZK)](https://codecov.io/gh/openvinotoolkit/anomalib)

diff --git a/docs/source/images/dsr/architecture.png b/docs/source/images/dsr/architecture.png
diff --git a/docs/source/images/uflow/diagram.png b/docs/source/images/uflow/diagram.png
diff --git a/docs/source/images/uflow/iou.png b/docs/source/images/uflow/iou.png
diff --git a/docs/source/images/uflow/more-results.png b/docs/source/images/uflow/more-results.png
diff --git a/docs/source/images/uflow/pixel-aupro.png b/docs/source/images/uflow/pixel-aupro.png
diff --git a/docs/source/images/uflow/pixel-auroc.png b/docs/source/images/uflow/pixel-auroc.png
diff --git a/docs/source/images/uflow/results-mvtec-anomalies.jpg b/docs/source/images/uflow/results-mvtec-anomalies.jpg
diff --git a/docs/source/images/uflow/results-mvtec-good.jpg b/docs/source/images/uflow/results-mvtec-good.jpg
diff --git a/docs/source/images/uflow/results-others-anomalies.jpg b/docs/source/images/uflow/results-others-anomalies.jpg
diff --git a/docs/source/images/uflow/results-others-good.jpg b/docs/source/images/uflow/results-others-good.jpg
diff --git a/docs/source/images/uflow/teaser.jpg b/docs/source/images/uflow/teaser.jpg
diff --git a/docs/source/markdown/guides/reference/models/image/dsr.md b/docs/source/markdown/guides/reference/models/image/dsr.md
@@ -0,0 +1,44 @@
+# DSR
+
+This is the implementation of the [DSR](https://link.springer.com/chapter/10.1007/978-3-031-19821-2_31) paper.
+
+Model Type: Segmentation
+
+## Description
+
+DSR is a quantized-feature based algorithm that consists of an autoencoder with one encoder and two decoders, coupled with an anomaly detection module. DSR learns a codebook of quantized representations on ImageNet, which are then used to encode input images. These quantized representations also serve to sample near-in-distribution anomalies, since they do not rely on external datasets. Training takes place in three phases. The encoder and "general object decoder", as well as the codebook, are pretrained on ImageNet. Defects are then generated at the feature level using the codebook on the quantized representations, and are used to train the object-specific decoder as well as the anomaly detection module. In the final phase of training, the upsampling module is trained on simulated image-level smudges in order to output more robust anomaly maps.
+
+## Architecture
+
+```{eval-rst}
+.. image:: https://raw.githubusercontent.com/openvinotoolkit/anomalib/main/docs/source/images/dsr/architecture.png
+    :alt: DSR Architecture
+```
+
+```{eval-rst}
+.. automodule:: anomalib.models.image.dsr.torch_model
+   :members:
+   :undoc-members:
+   :show-inheritance:
+```
+
+```{eval-rst}
+.. automodule:: anomalib.models.image.dsr.lightning_model
+   :members:
+   :undoc-members:
+   :show-inheritance:
+```
+
+```{eval-rst}
+.. automodule:: anomalib.models.image.dsr.anomaly_generator
+   :members:
+   :undoc-members:
+   :show-inheritance:
+```
+
+```{eval-rst}
+.. automodule:: anomalib.models.image.dsr.loss
+   :members:
+   :undoc-members:
+   :show-inheritance:
+```
diff --git a/docs/source/markdown/guides/reference/models/image/index.md b/docs/source/markdown/guides/reference/models/image/index.md
@@ -46,6 +46,13 @@ Probabilistic Modeling of Deep Features for Out-of-Distribution and Adversarial
 DRÆM – A discriminatively trained reconstruction embedding for surface anomaly detection
 :::
 
+:::{grid-item-card} {material-regular}`model_training;1.5em` DSR
+:link: ./dsr
+:link-type: doc
+
+DSR – A Dual Subspace Re-Projection Network for Surface Anomaly Detection
+:::
+
 :::{grid-item-card} {material-regular}`model_training;1.5em` Efficient AD
 :link: ./efficient_ad
 :link-type: doc
@@ -102,6 +109,11 @@ Region-Based Kernel Density Estimation (RKDE)
 Student-Teacher Feature Pyramid Matching for Unsupervised Anomaly Detection
 :::
 
+:::{grid-item-card} {material-regular}`model_training;1.5em` U-Flow
+:link: ./uflow
+:link-type: doc
+
+U-Flow: A U-shaped Normalizing Flow for Anomaly Detection with Unsupervised Threshold
 :::{grid-item-card} {material-regular}`model_training;1.5em` WinCLIP
 :link: ./winclip
 :link-type: doc

diff --git a/docs/source/markdown/guides/reference/models/image/uflow.md b/docs/source/markdown/guides/reference/models/image/uflow.md
@@ -0,0 +1,41 @@
+# U-Flow
+
+This is the implementation of the [U-Flow](https://arxiv.org/abs/2211.12353) paper.
+
+Model Type: Segmentation
+
+## Description
+
+U-Flow is a U-Shaped normalizing flow-based probability distribution estimator.
+The method consists of three phases.
+(1) Multi-scale feature extraction: a rich multi-scale representation is obtained with MSCaiT, by combining pre-trained image Transformers acting at different image scales. It can also be used any other feature extractor, such as ResNet.
+(2) U-shaped Normalizing Flow: by adapting the widely used U-like architecture to NFs, a fully invertible architecture is designed. This architecture is capable of merging the information from different scales while ensuring independence both intra- and inter-scales. To make it fully invertible, split and invertible up-sampling operations are used.
+(3) Anomaly score and segmentation computation: besides generating the anomaly map based on the likelihood of test data, we also propose to adapt the a contrario framework to obtain an automatic threshold by controlling the allowed number of false alarms.
+
+## Architecture
+
+```{eval-rst}
+.. image:: ../../images/uflow/diagram.png
+    :alt: U-Flow Architecture
+```
+
+```{eval-rst}
+.. automodule:: anomalib.models.uflow.torch_model
+   :members:
+   :undoc-members:
+   :show-inheritance:
+```
+
+```{eval-rst}
+.. automodule:: anomalib.models.uflow.lightning_model
+   :members:
+   :undoc-members:
+   :show-inheritance:
+```
+
+```{eval-rst}
+.. automodule:: anomalib.models.uflow.anomaly_map
+   :members:
+   :undoc-members:
+   :show-inheritance:
+```
diff --git a/setup.py b/setup.py
@@ -86,6 +86,22 @@ def get_required_packages(requirement_files: list[str]) -> list[str]:
     "full": get_required_packages(requirement_files=["loggers", "notebooks", "openvino"]),
 }
 
+try:
+    # OpenCV installed via conda.
+    import cv2
+
+    major, minor, *rest = cv2.__version__.split(".")
+    opencv_base = next((req for req in INSTALL_REQUIRES if req.startswith("opencv-python")), None)
+    if opencv_base is not None:
+        opencv_base_version = opencv_base.split(">=")[-1]
+        req_major, req_minor, *req_rest = opencv_base_version.split(".")
+        if int(major) < int(req_major) and int(minor) < int(req_minor):
+            msg = f"OpenCV >={req_major}.{req_minor} is required but {cv2.__version__} is installed"
+            raise RuntimeError(msg)
+        print("Removing OpenCV requirement since it was found")
+        INSTALL_REQUIRES.remove(opencv_base)
+except ImportError:
+    print("Installing OpenCV since no installation was found")
 
 setup(
     name="anomalib",

diff --git a/src/anomalib/callbacks/metrics.py b/src/anomalib/callbacks/metrics.py
@@ -35,21 +35,21 @@ class _MetricsCallback(Callback):
     these to the lightning module.
 
     Args:
-        task (TaskType): Task type of the current run.
+        task (TaskType | str): Task type of the current run.
         image_metrics (list[str] | str | None): List of image-level metrics.
         pixel_metrics (list[str] | str | None): List of pixel-level metrics.
         device (str): Whether to compute metrics on cpu or gpu. Defaults to cpu.
     """
 
     def __init__(
         self,
-        task: TaskType = TaskType.SEGMENTATION,
+        task: TaskType | str = TaskType.SEGMENTATION,
         image_metrics: list[str] | str | None = None,
         pixel_metrics: list[str] | str | None = None,
         device: Device = Device.CPU,
     ) -> None:
         super().__init__()
-        self.task = task
+        self.task = TaskType(task)
         self.image_metric_names = image_metrics
         self.pixel_metric_names = pixel_metrics
         self.device = device
@@ -90,7 +90,13 @@ def setup(
 
         if isinstance(pl_module, AnomalyModule):
             pl_module.image_metrics = create_metric_collection(image_metric_names, "image_")
-            pl_module.pixel_metrics = create_metric_collection(pixel_metric_names, "pixel_")
+            if hasattr(pl_module, "pixel_metrics"):  # incase metrics are loaded from model checkpoint
+                new_metrics = create_metric_collection(pixel_metric_names, "pixel_")
+                for name in new_metrics:
+                    if name not in pl_module.pixel_metrics:
+                        pl_module.pixel_metrics.add_metrics(new_metrics[name])
+            else:
+                pl_module.pixel_metrics = create_metric_collection(pixel_metric_names, "pixel_")
             self._set_threshold(pl_module)
 
     def on_validation_epoch_start(

diff --git a/src/anomalib/models/__init__.py b/src/anomalib/models/__init__.py
@@ -28,6 +28,7 @@
     ReverseDistillation,
     Rkde,
     Stfpm,
+    Uflow,
     WinClip,
 )
 from .video import AiVad
@@ -44,6 +45,7 @@ class UnknownModelError(ModuleNotFoundError):
     "Dfkde",
     "Dfm",
     "Draem",
+    "Dsr",
     "EfficientAd",
     "Fastflow",
     "Ganomaly",
@@ -52,6 +54,7 @@ class UnknownModelError(ModuleNotFoundError):
     "ReverseDistillation",
     "Rkde",
     "Stfpm",
+    "Uflow",
     "AiVad",
     "WinClip",
 ]

diff --git a/src/anomalib/models/components/base/anomaly_module.py b/src/anomalib/models/components/base/anomaly_module.py
@@ -15,13 +15,13 @@
 from torch import nn
 
 from anomalib import LearningType
+from anomalib.metrics import AnomalibMetricCollection
 from anomalib.metrics.threshold import BaseThreshold
 
 if TYPE_CHECKING:
     from lightning.pytorch.callbacks import Callback
     from torchmetrics import Metric
 
-    from anomalib.metrics import AnomalibMetricCollection
 
 logger = logging.getLogger(__name__)
 
@@ -136,9 +136,39 @@ def load_state_dict(self, state_dict: OrderedDict[str, Any], strict: bool = True
             self.pixel_threshold = self._get_instance(state_dict, "pixel_threshold_class")
         if "normalization_class" in state_dict:
             self.normalization_metrics = self._get_instance(state_dict, "normalization_class")
+        # Used to load metrics if there is any related data in state_dict
+        self._load_metrics(state_dict)
 
         return super().load_state_dict(state_dict, strict)
 
+    def _load_metrics(self, state_dict: OrderedDict[str, torch.Tensor]) -> None:
+        """Load metrics from saved checkpoint."""
+        self._add_metrics("pixel", state_dict)
+        self._add_metrics("image", state_dict)
+
+    def _add_metrics(self, name: str, state_dict: OrderedDict[str, torch.Tensor]) -> None:
+        """Sets the pixel/image metrics.
+
+        Args:
+            name (str): is it pixel or image.
+            state_dict (OrderedDict[str, Tensor]): state dict of the model.
+        """
+        metric_keys = [key for key in state_dict if key.startswith(f"{name}_metrics")]
+        if any(metric_keys):
+            if not hasattr(self, f"{name}_metrics"):
+                setattr(self, f"{name}_metrics", AnomalibMetricCollection([], prefix=name))
+            metrics = getattr(self, f"{name}_metrics")
+            for key in metric_keys:
+                class_name = key.split(".")[1]
+                try:
+                    metrics_module = importlib.import_module("anomalib.metrics")
+                    metrics_cls = getattr(metrics_module, class_name)
+                except (ImportError, AttributeError) as exception:
+                    msg = f"Class {class_name} not found in module anomalib.metrics"
+                    raise ImportError(msg) from exception
+                logger.info("Loading %s metrics from state dict", class_name)
+                metrics.add_metrics(metrics_cls())
+
     def _get_instance(self, state_dict: OrderedDict[str, Any], dict_key: str) -> BaseThreshold:
         """Get the threshold class from the ``state_dict``."""
         class_path = state_dict.pop(dict_key)

diff --git a/src/anomalib/models/image/__init__.py b/src/anomalib/models/image/__init__.py
@@ -9,6 +9,7 @@
 from .dfkde import Dfkde
 from .dfm import Dfm
 from .draem import Draem
+from .dsr import Dsr
 from .efficient_ad import EfficientAd
 from .fastflow import Fastflow
 from .ganomaly import Ganomaly
@@ -17,6 +18,7 @@
 from .reverse_distillation import ReverseDistillation
 from .rkde import Rkde
 from .stfpm import Stfpm
+from .uflow import Uflow
 from .winclip import WinClip
 
 __all__ = [
@@ -26,6 +28,7 @@
     "Dfkde",
     "Dfm",
     "Draem",
+    "Dsr",
     "EfficientAd",
     "Fastflow",
     "Ganomaly",
@@ -34,5 +37,6 @@
     "ReverseDistillation",
     "Rkde",
     "Stfpm",
+    "Uflow",
     "WinClip",
 ]
diff --git a/src/anomalib/models/image/cfa/README.md b/src/anomalib/models/image/cfa/README.md
@@ -139,11 +139,11 @@ When the numbers are produced, early stopping callback (patience: 5) is used. It
 
 ### Sample Results
 
-![Sample Result 1](../../../docs/source/images/cfa/results/0.png "Sample Result 1")
+![Sample Result 1](../../../../docs/source/images/cfa/results/0.png "Sample Result 1")
 
-![Sample Result 2](../../../docs/source/images/cfa/results/1.png "Sample Result 2")
+![Sample Result 2](../../../../docs/source/images/cfa/results/1.png "Sample Result 2")
 
-![Sample Result 3](../../../docs/source/images/cfa/results/2.png "Sample Result 3")
+![Sample Result 3](../../../../docs/source/images/cfa/results/2.png "Sample Result 3")
 
 ## Reference
 

diff --git a/src/anomalib/models/image/cflow/README.md b/src/anomalib/models/image/cflow/README.md
@@ -42,8 +42,8 @@ All results gathered with seed `42`.
 
 ### Sample Results
 
-![Sample Result 1](../../../docs/source/images/cflow/results/0.png "Sample Result 1")
+![Sample Result 1](../../../../docs/source/images/cflow/results/0.png "Sample Result 1")
 
-![Sample Result 2](../../../docs/source/images/cflow/results/1.png "Sample Result 2")
+![Sample Result 2](../../../../docs/source/images/cflow/results/1.png "Sample Result 2")
 
-![Sample Result 3](../../../docs/source/images/cflow/results/2.png "Sample Result 3")
+![Sample Result 3](../../../../docs/source/images/cflow/results/2.png "Sample Result 3")
diff --git a/src/anomalib/models/image/csflow/README.md b/src/anomalib/models/image/csflow/README.md
@@ -93,8 +93,8 @@ All results gathered with seed `42`.
 
 ### Sample Results
 
-![Sample Result 1](../../../docs/source/images/csflow/results/0.png "Sample Result 1")
+![Sample Result 1](../../../../docs/source/images/csflow/results/0.png "Sample Result 1")
 
-![Sample Result 2](../../../docs/source/images/csflow/results/1.png "Sample Result 2")
+![Sample Result 2](../../../../docs/source/images/csflow/results/1.png "Sample Result 2")
 
-![Sample Result 3](../../../docs/source/images/csflow/results/2.png "Sample Result 3")
+![Sample Result 3](../../../../docs/source/images/csflow/results/2.png "Sample Result 3")
diff --git a/src/anomalib/models/image/draem/lightning_model.py b/src/anomalib/models/image/draem/lightning_model.py
@@ -42,10 +42,11 @@ def __init__(
         enable_sspcab: bool = False,
         sspcab_lambda: float = 0.1,
         anomaly_source_path: str | None = None,
+        beta: float | tuple[float, float] = (0.1, 1.0),
     ) -> None:
         super().__init__()
 
-        self.augmenter = Augmenter(anomaly_source_path)
+        self.augmenter = Augmenter(anomaly_source_path, beta=beta)
         self.model = DraemModel(sspcab=enable_sspcab)
         self.loss = DraemLoss()
         self.sspcab = enable_sspcab
@@ -138,7 +139,9 @@ def trainer_arguments(self) -> dict[str, Any]:
 
     def configure_optimizers(self) -> torch.optim.Optimizer:
         """Configure the Adam optimizer."""
-        return torch.optim.Adam(params=self.model.parameters(), lr=0.0001)
+        optimizer = torch.optim.Adam(params=self.model.parameters(), lr=0.0001)
+        scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[400, 600], gamma=0.1)
+        return [optimizer], [scheduler]
 
     @property
     def learning_type(self) -> LearningType: