🚀 feat(model): add automatic download of DTD dataset to DRAEM model (#2866)

Added automatic download of DTD dataset to DRAEM model

Author: abc-125

docs/source/examples

@@ -1 +1 @@
-../../examples
+../../examples

examples/configs/model/draem.yaml

@@ -4,7 +4,7 @@ model:
     beta: [0.1, 1.0]
     enable_sspcab: false
     sspcab_lambda: 0.1
-    anomaly_source_path: null
+    dtd_dir: ./datasets/dtd
 
 trainer:
   max_epochs: 700

src/anomalib/data/utils/generators/perlin.py

@@ -200,7 +200,7 @@ class PerlinAnomalyGenerator(v2.Transform):
 
     def __init__(
         self,
-        anomaly_source_path: str | None = None,
+        anomaly_source_path: Path | str | None = None,
         probability: float = 0.5,
         blend_factor: float | tuple[float, float] = (0.2, 1.0),
         rotation_range: tuple[float, float] = (-90, 90),

src/anomalib/models/image/draem/README.md

@@ -8,8 +8,6 @@ Model Type: Segmentation
 
 DRAEM is a reconstruction based algorithm that consists of a reconstructive subnetwork and a discriminative subnetwork. DRAEM is trained on simulated anomaly images, generated by augmenting normal input images from the training set with a random Perlin noise mask extracted from an unrelated source of image data. The reconstructive subnetwork is an autoencoder architecture that is trained to reconstruct the original input images from the augmented images. The reconstructive submodel is trained using a combination of L2 loss and Structural Similarity loss. The input of the discriminative subnetwork consists of the channel-wise concatenation of the (augmented) input image and the output of the reconstructive subnetwork. The output of the discriminative subnetwork is an anomaly map that contains the predicted anomaly scores for each pixel location. The discriminative subnetwork is trained using Focal Loss.
 
-For optimal results, DRAEM requires specifying the path to a folder of image data that will be used as the source of the anomalous pixel regions in the simulated anomaly images. The path can be specified by editing the value of the `model.anomaly_source_path` parameter in the `config.yaml` file. The authors of the original paper recommend using the [DTD](https://www.robots.ox.ac.uk/~vgg/data/dtd/) dataset as anomaly source.
-
 ## Architecture
 
 ![DRAEM Architecture](/docs/source/images/draem/architecture.png "DRAEM Architecture")

src/anomalib/models/image/draem/lightning_model.py

@@ -14,6 +14,7 @@
 """
 
 from collections.abc import Callable
+from pathlib import Path
 from typing import Any
 
 import torch
@@ -23,6 +24,7 @@
 
 from anomalib import LearningType
 from anomalib.data import Batch
+from anomalib.data.utils import DownloadInfo, download_and_extract
 from anomalib.data.utils.generators.perlin import PerlinAnomalyGenerator
 from anomalib.metrics import Evaluator
 from anomalib.models.components import AnomalibModule
@@ -35,6 +37,12 @@
 
 __all__ = ["Draem"]
 
+DTD_DOWNLOAD_INFO = DownloadInfo(
+    name="dtd-r1.0.1.tar.gz",
+    url="https://www.robots.ox.ac.uk/~vgg/data/dtd/download/dtd-r1.0.1.tar.gz",
+    hashsum="e42855a52a4950a3b59612834602aa253914755c95b0cff9ead6d07395f8e205",
+)
+
 
 class Draem(AnomalibModule):
     """DRÆM.
@@ -47,6 +55,8 @@ class Draem(AnomalibModule):
     2. A discriminative network that learns to identify anomalous regions
 
     Args:
+        dtd_dir (Path | str): Directory path for the DTD dataset for anomaly deneration.
+            Defaults to ``./datasets/dtd``.
         enable_sspcab (bool, optional): Enable SSPCAB training.
             Defaults to ``False``.
         sspcab_lambda (float, optional): Weight factor for SSPCAB loss.
@@ -73,9 +83,9 @@ class Draem(AnomalibModule):
 
     def __init__(
         self,
+        dtd_dir: Path | str = "./datasets/dtd",
         enable_sspcab: bool = False,
         sspcab_lambda: float = 0.1,
-        anomaly_source_path: str | None = None,
         beta: float | tuple[float, float] = (0.1, 1.0),
         pre_processor: PreProcessor | bool = True,
         post_processor: PostProcessor | bool = True,
@@ -88,8 +98,10 @@ def __init__(
             evaluator=evaluator,
             visualizer=visualizer,
         )
-
-        self.augmenter = PerlinAnomalyGenerator(anomaly_source_path=anomaly_source_path, blend_factor=beta)
+        dtd_dir = Path(dtd_dir)
+        if not dtd_dir.is_dir():
+            download_and_extract(dtd_dir, DTD_DOWNLOAD_INFO)
+        self.augmenter = PerlinAnomalyGenerator(anomaly_source_path=dtd_dir, blend_factor=beta)
         self.model = DraemModel(sspcab=enable_sspcab)
         self.loss = DraemLoss()
         self.sspcab = enable_sspcab