✨ Add Reverse Distillation

anomalib/models/__init__.py

@@ -24,14 +24,26 @@
 from anomalib.models.components import AnomalyModule
 
 
+def _snake_to_pascal_case(model_name: str) -> str:
+    """Convert model name from snake case to Pascal case.
+
+    Args:
+        model_name (str): Model name in snake case.
+
+    Returns:
+        str: Model name in Pascal case.
+    """
+    return "".join([split.capitalize() for split in model_name.split("_")])
+
+
 def get_model(config: Union[DictConfig, ListConfig]) -> AnomalyModule:
     """Load model from the configuration file.
 
     Works only when the convention for model naming is followed.
 
     The convention for writing model classes is
-    `anomalib.models.<model_name>.model.<Model_name>Lightning`
-    `anomalib.models.stfpm.model.StfpmLightning`
+    `anomalib.models.<model_name>.lightning_model.<ModelName>Lightning`
+    `anomalib.models.stfpm.lightning_model.StfpmLightning`
 
     Args:
         config (Union[DictConfig, ListConfig]): Config.yaml loaded using OmegaConf
@@ -42,12 +54,23 @@ def get_model(config: Union[DictConfig, ListConfig]) -> AnomalyModule:
     Returns:
         AnomalyModule: Anomaly Model
     """
-    model_list: List[str] = ["cflow", "dfkde", "dfm", "draem", "fastflow", "ganomaly", "padim", "patchcore", "stfpm"]
+    model_list: List[str] = [
+        "cflow",
+        "dfkde",
+        "dfm",
+        "draem",
+        "fastflow",
+        "ganomaly",
+        "padim",
+        "patchcore",
+        "reverse_distillation",
+        "stfpm",
+    ]
     model: AnomalyModule
 
     if config.model.name in model_list:
         module = import_module(f"anomalib.models.{config.model.name}")
-        model = getattr(module, f"{config.model.name.capitalize()}Lightning")(config)
+        model = getattr(module, f"{_snake_to_pascal_case(config.model.name)}Lightning")(config)
 
     else:
         raise ValueError(f"Unknown model {config.model.name}!")

anomalib/models/draem/config.yaml

@@ -45,8 +45,11 @@ metrics:
 project:
   seed: 42
   path: ./results
-  log_images_to: ["local"]
-  logger: false # options: [tensorboard, wandb, csv] or combinations.
+
+logging:
+  log_images_to: ["local"] # options: [wandb, tensorboard, local]. Make sure you also set logger with using wandb or tensorboard.
+  logger: [] # options: [tensorboard, wandb, csv] or combinations.
+  log_graph: false # Logs the model graph to respective logger.
 
 optimization:
   openvino:

anomalib/models/fastflow/config.yaml

@@ -49,8 +49,11 @@ metrics:
 project:
   seed: 0
   path: ./results
-  log_images_to: [local]
-  logger: false # options: [tensorboard, wandb, csv] or combinations.
+
+logging:
+  log_images_to: ["local"] # options: [wandb, tensorboard, local]. Make sure you also set logger with using wandb or tensorboard.
+  logger: [] # options: [tensorboard, wandb, csv] or combinations.
+  log_graph: false # Logs the model graph to respective logger.
 
 # PL Trainer Args. Don't add extra parameter here.
 trainer:

anomalib/models/reverse_distillation/LICENSE

@@ -0,0 +1,29 @@
+Copyright (c) 2022 Intel Corporation
+SPDX-License-Identifier: Apache-2.0
+
+Some files in this folder are based on the original Reverse Distillation implementation by hq-deng
+
+Original license
+----------------
+
+    MIT License
+
+    Copyright (c) 2022 hq-deng
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy
+    of this software and associated documentation files (the "Software"), to deal
+    in the Software without restriction, including without limitation the rights
+    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the Software is
+    furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all
+    copies or substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE.

anomalib/models/reverse_distillation/README.md

@@ -0,0 +1,98 @@
+# Anomaly Detection via Reverse Distillation from One-Class Embedding
+
+This is the implementation of the [Reverse Distillation](https://arxiv.org/pdf/2201.10703v2.pdf) paper.
+
+Model Type: Segmentation
+
+## Description
+
+Reverse Distillation model consists of three networks. The first is a pre-trained feature extractor (E). The next two are the one-class bottleneck embedding (OCBE) and the student decoder network (D). The backbone E is a ResNet model pre-trained on ImageNet dataset. During the forward pass, features from three ResNet block are extracted. These features are encoded by concatenating the three feature maps using the multi-scale feature fusion block of OCBE and passed to the decoder D. The decoder network is symmetrical to the feature extractor but reversed. During training, outputs from these symmetrical blocks are forced to be similar to the corresponding feature extractor layers by using cosine distance as the loss metric.
+
+During testing, a similar step is followed but this time the cosine distance between the feature maps is used to indicate the presence of anomalies. The distance maps from all the three layers are up-sampled to the image size and added (or multiplied) to produce the final feature map. Gaussian blur is applied to the output map to make it smoother. Finally, the anomaly map is generated by applying min-max normalization on the output map.
+
+## Architecture
+
+![Anomaly Detection via Reverse Distillation from One-Class Embedding Architecture](../../../docs/source/images/reverse_distillation/architecture.png "Reverse Distillation Architecture")
+
+## Usage
+
+`python tools/train.py --model reverse_distillation`
+
+## Benchmark
+
+All results gathered with seed `42`.
+
+Note: Early Stopping (with patience 3) was enabled during training.
+
+## [MVTec AD Dataset](https://www.mvtec.com/company/research/datasets/mvtec-ad)
+
+### Image-Level AUC
+
+|            | ResNet 18 | Wide ResNet 50 |
+| :--------- | --------: | -------------: |
+| Bottle     |         1 |          0.977 |
+| Cable      |     0.977 |          0.954 |
+| Capsule    |     0.947 |          0.738 |
+| Carpet     |      0.87 |          0.983 |
+| Grid       |     0.995 |          0.646 |
+| Hazelnut   |     0.915 |          0.994 |
+| Leather    |     0.994 |          0.999 |
+| Metal_nut  |         1 |           0.95 |
+| Pill       |     0.939 |          0.641 |
+| Screw      |       0.9 |          0.832 |
+| Tile       |     0.905 |           0.72 |
+| Toothbrush |       0.9 |          0.817 |
+| Transistor |     0.953 |          0.569 |
+| Wood       |     0.988 |          0.973 |
+| Zipper     |     0.888 |          0.952 |
+| Average    |     0.945 |           0.85 |
+
+### Pixel-Level AUC
+
+|            | ResNet 18 | Wide ResNet 50 |
+| :--------- | --------: | -------------: |
+| Bottle     |     0.982 |          0.985 |
+| Cable      |     0.979 |           0.97 |
+| Capsule    |     0.975 |          0.966 |
+| Carpet     |     0.945 |          0.974 |
+| Grid       |     0.973 |          0.678 |
+| Hazelnut   |     0.981 |          0.987 |
+| Leather    |     0.972 |          0.957 |
+| Metal_nut  |     0.975 |          0.989 |
+| Pill       |     0.981 |          0.984 |
+| Screw      |     0.989 |          0.988 |
+| Tile       |     0.805 |           0.76 |
+| Toothbrush |     0.987 |          0.988 |
+| Transistor |     0.912 |          0.889 |
+| Wood       |     0.814 |          0.956 |
+| Zipper     |     0.943 |          0.968 |
+| Average    |     0.948 |          0.936 |
+
+### Image F1 Score
+
+|            | ResNet 18 | Wide ResNet 50 |
+| :--------- | --------: | -------------: |
+| Bottle     |         1 |          0.977 |
+| Cable      |     0.977 |          0.954 |
+| Capsule    |     0.947 |          0.738 |
+| Carpet     |      0.87 |          0.983 |
+| Grid       |     0.995 |          0.646 |
+| Hazelnut   |     0.915 |          0.994 |
+| Leather    |     0.994 |          0.999 |
+| Metal_nut  |         1 |           0.95 |
+| Pill       |     0.939 |          0.641 |
+| Screw      |       0.9 |          0.832 |
+| Tile       |     0.905 |           0.72 |
+| Toothbrush |       0.9 |          0.817 |
+| Transistor |     0.953 |          0.569 |
+| Wood       |     0.988 |          0.973 |
+| Zipper     |     0.888 |          0.952 |
+| Average    |     0.945 |           0.85 |
+
+### Sample Results
+
+![Sample Result 1](../../../docs/source/images/reverse_distillation/results/0.png "Sample Result 1")
+
+![Sample Result 2](../../../docs/source/images/reverse_distillation/results/1.png "Sample Result 2")
+
+![Sample Result 3](../../../docs/source/images/reverse_distillation/results/2.png "Sample Result 3")

anomalib/models/reverse_distillation/__init__.py

@@ -0,0 +1,8 @@
+"""Reverse Distillation Model."""
+
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .lightning_model import ReverseDistillationLightning
+
+__all__ = ["ReverseDistillationLightning"]

anomalib/models/reverse_distillation/anomaly_map.py

@@ -0,0 +1,75 @@
+"""Compute Anomaly map."""
+
+# Original Code
+# Copyright (c) 2022 hq-deng
+# https://github.com/hq-deng/RD4AD
+# SPDX-License-Identifier: MIT
+#
+# Modified
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import List, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from kornia.filters import gaussian_blur2d
+from omegaconf import ListConfig
+from torch import Tensor
+
+
+class AnomalyMapGenerator:
+    """Generate Anomaly Heatmap.
+
+    Args:
+        image_size (Union[ListConfig, Tuple]): Size of original image used for upscaling the anomaly map.
+        sigma (int): Standard deviation of the gaussian kernel used to smooth anomaly map.
+        mode (str, optional): Operation used to generate anomaly map. Options are `add` and `multiply`.
+                Defaults to "multiply".
+
+    Raises:
+        ValueError: In case modes other than multiply and add are passed.
+    """
+
+    def __init__(self, image_size: Union[ListConfig, Tuple], sigma: int = 4, mode: str = "multiply"):
+        self.image_size = image_size if isinstance(image_size, tuple) else tuple(image_size)
+        self.sigma = sigma
+        self.kernel_size = 2 * int(4.0 * sigma + 0.5) + 1
+
+        if mode not in ("add", "multiply"):
+            raise ValueError(f"Found mode {mode}. Only multiply and add are supported.")
+        self.mode = mode
+
+    def __call__(self, student_features: List[Tensor], teacher_features: List[Tensor]) -> Tensor:
+        """Computes anomaly map given encoder and decoder features.
+
+        Args:
+            student_features (List[Tensor]): List of encoder features
+            teacher_features (List[Tensor]): List of decoder features
+
+        Returns:
+            Tensor: Anomaly maps of length batch.
+        """
+        if self.mode == "multiply":
+            anomaly_map = torch.ones(
+                [student_features[0].shape[0], 1, *self.image_size], device=student_features[0].device
+            )  # b c h w
+        elif self.mode == "add":
+            anomaly_map = torch.zeros(
+                [student_features[0].shape[0], 1, *self.image_size], device=student_features[0].device
+            )
+
+        for student_feature, teacher_feature in zip(student_features, teacher_features):
+            distance_map = 1 - F.cosine_similarity(student_feature, teacher_feature)
+            distance_map = torch.unsqueeze(distance_map, dim=1)
+            distance_map = F.interpolate(distance_map, size=self.image_size, mode="bilinear", align_corners=True)
+            if self.mode == "multiply":
+                anomaly_map *= distance_map
+            elif self.mode == "add":
+                anomaly_map += distance_map
+
+        anomaly_map = gaussian_blur2d(
+            anomaly_map, kernel_size=(self.kernel_size, self.kernel_size), sigma=(self.sigma, self.sigma)
+        )
+
+        return anomaly_map

anomalib/models/reverse_distillation/components/__init__.py

@@ -0,0 +1,20 @@
+"""PyTorch modules for Reverse Distillation."""
+
+# Copyright (C) 2022 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions
+# and limitations under the License.
+
+from .bottleneck import get_bottleneck_layer
+from .de_resnet import get_decoder
+
+__all__ = ["get_bottleneck_layer", "get_decoder"]