✨ Add Reverse Distillation

anomalib/models/__init__.py

@@ -42,7 +42,7 @@ def get_model(config: Union[DictConfig, ListConfig]) -> AnomalyModule:
     Returns:
         AnomalyModule: Anomaly Model
     """
-    model_list: List[str] = ["cflow", "dfkde", "dfm", "ganomaly", "padim", "patchcore", "stfpm"]
+    model_list: List[str] = ["cflow", "dfkde", "dfm", "ganomaly", "padim", "patchcore", "reversedistillation", "stfpm"]
     model: AnomalyModule
 
     if config.model.name in model_list:

anomalib/models/reversedistillation/README.md

@@ -0,0 +1,32 @@
+# 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. The next two are the one-class embedding and the decoder networks.
+
+
+## Architecture
+
+![Anomaly Detection via Reverse Distillation from One-Class Embedding Architecture](../../../docs/source/images/reversedistillation/architecture.png "Reverse Distillation Architecture")
+
+## Usage
+
+`python tools/train.py --model reversedistillation`
+
+## Benchmark
+
+All results gathered with seed `42`.
+
+## [MVTec AD Dataset](https://www.mvtec.com/company/research/datasets/mvtec-ad)
+
+### Image-Level AUC
+
+### Pixel-Level AUC
+
+### Image F1 Score
+
+### Sample Results

anomalib/models/reversedistillation/__init__.py

@@ -0,0 +1,19 @@
+"""Reverse Distillation Model."""
+
+# 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 .lightning_model import ReversedistillationLightning
+
+__all__ = ["ReversedistillationLightning"]

anomalib/models/reversedistillation/anomaly_map.py

@@ -0,0 +1,77 @@
+"""compute Anomaly map."""
+
+# 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 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.
+    """
+
+    def __init__(self, image_size: Union[ListConfig, Tuple], sigma: int = 4):
+        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
+
+    def __call__(self, fs_list: List[Tensor], ft_list: List[Tensor], amap_mode: str = "mul") -> Tensor:
+        """Computes anomaly map given encoder and decoder features.
+
+        Args:
+            fs_list (List[Tensor]): List of encoder features
+            ft_list (List[Tensor]): List of decoder features
+            amap_mode (str, optional): Operation used to generate anomaly map. Options are `add` and `mul`.
+                Defaults to "mul".
+
+        Raises:
+            ValueError: _description_
+
+        Returns:
+            Tensor: _description_
+        """
+        if amap_mode == "mul":
+            anomaly_map = torch.ones([fs_list[0].shape[0], 1, *self.image_size], device=fs_list[0].device)  # b c h w
+        elif amap_mode == "add":
+            anomaly_map = torch.zeros([fs_list[0].shape[0], 1, *self.image_size], device=fs_list[0].device)
+        else:
+            raise ValueError(f"Found amap_mode {amap_mode}. Only mul and add are supported.")
+        for i in range(len(ft_list)):
+            fs = fs_list[i]
+            ft = ft_list[i]
+            a_map = 1 - F.cosine_similarity(fs, ft)
+            a_map = torch.unsqueeze(a_map, dim=1)
+            a_map = F.interpolate(a_map, size=self.image_size, mode="bilinear", align_corners=True)
+            if amap_mode == "mul":
+                anomaly_map *= a_map
+            elif amap_mode == "add":
+                anomaly_map += a_map
+            else:
+                raise ValueError(f"Operation {amap_mode} not supported. Only ``add`` and ``mul`` are supported")
+
+        anomaly_map = gaussian_blur2d(
+            anomaly_map, kernel_size=(self.kernel_size, self.kernel_size), sigma=(self.sigma, self.sigma)
+        )
+
+        return anomaly_map

anomalib/models/reversedistillation/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"]

anomalib/models/reversedistillation/components/bottleneck.py

@@ -0,0 +1,173 @@
+"""Torch model defining the bottleneck layer."""
+
+# 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.
+#
+# Code adapted from https://github.com/hq-deng/RD4AD
+
+
+from typing import Callable, List, Optional, Type
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torchvision.models.resnet import Bottleneck
+
+
+def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:
+    """3x3 convolution with padding."""
+    return nn.Conv2d(
+        in_planes,
+        out_planes,
+        kernel_size=3,
+        stride=stride,
+        padding=dilation,
+        groups=groups,
+        bias=False,
+        dilation=dilation,
+    )
+
+
+def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
+    """1x1 convolution."""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class OCBE(nn.Module):
+    """One-Class Bottleneck Embedding module.
+
+    Args:
+        block (Bottleneck): Expansion value is extracted from this block.
+        layers (int): Numbers of OCE layers to create after multiscale feature fusion.
+        groups (int, optional): Number of blocked connections from input channels to output channels.
+            Defaults to 1.
+        width_per_group (int, optional): Number of layers in each intermediate convolution layer. Defaults to 64.
+        norm_layer (Optional[Callable[..., nn.Module]], optional): Batch norm layer to use. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        block: Type[Bottleneck],
+        layers: int,
+        groups: int = 1,
+        width_per_group: int = 64,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ):
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+        self.groups = groups
+        self.base_width = width_per_group
+        self.inplanes = 256 * block.expansion
+        self.dilation = 1
+        self.bn_layer = self._make_layer(block, 512, layers, stride=2)
+
+        self.conv1 = conv3x3(64 * block.expansion, 128 * block.expansion, 2)
+        self.bn1 = norm_layer(128 * block.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(128 * block.expansion, 256 * block.expansion, 2)
+        self.bn2 = norm_layer(256 * block.expansion)
+        self.conv3 = conv3x3(128 * block.expansion, 256 * block.expansion, 2)
+        self.bn3 = norm_layer(256 * block.expansion)
+
+        self.conv4 = conv1x1(1024 * block.expansion, 512 * block.expansion, 1)
+        self.bn4 = norm_layer(512 * block.expansion)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def _make_layer(
+        self,
+        block: Type[Bottleneck],
+        planes: int,
+        blocks: int,
+        stride: int = 1,
+        dilate: bool = False,
+    ) -> nn.Sequential:
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes * 3, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(
+            block(
+                self.inplanes * 3,
+                planes,
+                stride,
+                downsample,
+                self.groups,
+                self.base_width,
+                previous_dilation,
+                norm_layer,
+            )
+        )
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(
+                block(
+                    self.inplanes,
+                    planes,
+                    groups=self.groups,
+                    base_width=self.base_width,
+                    dilation=self.dilation,
+                    norm_layer=norm_layer,
+                )
+            )
+
+        return nn.Sequential(*layers)
+
+    def forward(self, features: List[Tensor]) -> Tensor:
+        """Forward-pass of Bottleneck layer.
+
+        Args:
+            features (List[Tensor]): List of features extracted from the encoder.
+
+        Returns:
+            Tensor: Output of the bottleneck layer
+        """
+        # Always assumes that features has length of 3
+        l1 = self.relu(self.bn2(self.conv2(self.relu(self.bn1(self.conv1(features[0]))))))
+        l2 = self.relu(self.bn3(self.conv3(features[1])))
+        feature = torch.cat([l1, l2, features[2]], 1)
+        output = self.bn_layer(feature)
+
+        return output.contiguous()
+
+
+def get_bottleneck_layer(backbone: str, **kwargs) -> OCBE:
+    """Get appropriate bottleneck layer based on the name of the backbone.
+
+    Args:
+        backbone (str): Name of the backbone.
+
+    Returns:
+        Bottleneck_layer: One-Class Bottleneck Embedding module.
+    """
+    if backbone == "resnet18":
+        return OCBE(Bottleneck, 2, **kwargs)
+    else:
+        return OCBE(Bottleneck, 3, **kwargs)