🚀 feat(model): add UniNet

examples/configs/model/uninet.yaml

@@ -0,0 +1,15 @@
+model:
+  class_path: anomalib.models.UniNet
+  init_args:
+    student_backbone: wide_resnet50_2
+    teacher_backbone: wide_resnet50_2
+    temperature: 0.1
+
+trainer:
+  max_epochs: 100
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.EarlyStopping
+      init_args:
+        patience: 20
+        monitor: image_AUROC
+        mode: max

src/anomalib/models/__init__.py

@@ -72,6 +72,7 @@
     Stfpm,
     Supersimplenet,
     Uflow,
+    UniNet,
     VlmAd,
     WinClip,
 )
@@ -108,6 +109,7 @@ class UnknownModelError(ModuleNotFoundError):
     "Stfpm",
     "Supersimplenet",
     "Uflow",
+    "UniNet",
     "VlmAd",
     "WinClip",
     "AiVad",

src/anomalib/models/components/backbone/__init__.py

@@ -0,0 +1,17 @@
+"""Common backbone models.
+
+Example:
+    >>> from anomalib.models.components.backbone import get_decoder
+    >>> decoder = get_decoder()
+
+See Also:
+    - :func:`anomalib.models.components.backbone.de_resnet`:
+        Decoder network implementation
+"""
+
+# Copyright (C) 2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .de_resnet import get_decoder
+
+__all__ = ["get_decoder"]

src/anomalib/models/image/reverse_distillation/components/de_resnet.py → src/anomalib/models/components/backbone/de_resnet.py

@@ -1,6 +1,6 @@
-"""PyTorch model defining the decoder network for Reverse Distillation.
+"""PyTorch model defining the decoder network for Reverse Distillation and UniNet.
 
-This module implements the decoder network used in the Reverse Distillation model
+This module implements the decoder network used in the Reverse Distillation and UniNet models
 architecture. The decoder reconstructs features from the bottleneck representation
 back to the original feature space.
 
@@ -10,17 +10,15 @@
 - Full decoder network architecture
 
 Example:
-    >>> from anomalib.models.image.reverse_distillation.components.de_resnet import (
+    >>> from anomalib.models.components.backbone.de_resnet import (
     ...     get_decoder
     ... )
     >>> decoder = get_decoder()
     >>> features = torch.randn(32, 512, 28, 28)
     >>> reconstructed = decoder(features)
 
 See Also:
-    - :class:`anomalib.models.image.reverse_distillation.torch_model.ReverseDistillationModel`:
-        Main model implementation using this decoder
-    - :class:`anomalib.models.image.reverse_distillation.components.DecoderBasicBlock`:
+    - :class:`anomalib.models.components.backbone.de_resnet.DecoderBasicBlock`:
         Basic building block for the decoder network
 """
 

src/anomalib/models/image/__init__.py

@@ -61,6 +61,7 @@
 from .stfpm import Stfpm
 from .supersimplenet import Supersimplenet
 from .uflow import Uflow
+from .uninet import UniNet
 from .vlm_ad import VlmAd
 from .winclip import WinClip
 
@@ -82,6 +83,7 @@
     "Stfpm",
     "Supersimplenet",
     "Uflow",
+    "UniNet",
     "VlmAd",
     "WinClip",
 ]

src/anomalib/models/image/reverse_distillation/components/__init__.py

@@ -7,7 +7,6 @@
 through distillation and reconstruction:
 
 - Bottleneck layer: Compresses features into a lower dimensional space
-- Decoder network: Reconstructs features from the bottleneck representation
 
 Example:
     >>> from anomalib.models.image.reverse_distillation.components import (
@@ -20,14 +19,11 @@
 See Also:
     - :func:`anomalib.models.image.reverse_distillation.components.bottleneck`:
         Bottleneck layer implementation
-    - :func:`anomalib.models.image.reverse_distillation.components.de_resnet`:
-        Decoder network implementation
 """
 
 # Copyright (C) 2022-2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
 from .bottleneck import get_bottleneck_layer
-from .de_resnet import get_decoder
 
-__all__ = ["get_bottleneck_layer", "get_decoder"]
+__all__ = ["get_bottleneck_layer"]

src/anomalib/models/image/reverse_distillation/torch_model.py

@@ -39,9 +39,10 @@
 
 from anomalib.data import InferenceBatch
 from anomalib.models.components import TimmFeatureExtractor
+from anomalib.models.components.backbone import get_decoder
 
 from .anomaly_map import AnomalyMapGenerationMode, AnomalyMapGenerator
-from .components import get_bottleneck_layer, get_decoder
+from .components import get_bottleneck_layer
 
 if TYPE_CHECKING:
     from anomalib.data.utils.tiler import Tiler

src/anomalib/models/image/uninet/LICENSE

@@ -0,0 +1,29 @@
+Copyright (c) 2025 Intel Corporation
+SPDX-License-Identifier: Apache-2.0
+
+Some files in this folder are based on the original UniNet implementation by Shun Wei, Jielin Jiang, and Xiaolong Xu.
+
+Original license:
+-----------------
+
+    MIT License
+
+    Copyright (c) 2025 Shun Wei
+
+    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.

src/anomalib/models/image/uninet/__init__.py

@@ -0,0 +1,33 @@
+"""UniNet Model for anomaly detection.
+
+This module implements anomaly detection using the UniNet model. It is a model designed for diverse domains and is
+suited for both supervised and unsupervised anomaly detection. It also focuses on multi-class anomaly detection.
+
+Example:
+    >>> from anomalib.models import UniNet
+    >>> from anomalib.data import MVTecAD
+    >>> from anomalib.engine import Engine
+
+    >>> # Initialize model and data
+    >>> datamodule = MVTecAD()
+    >>> model = UniNet()
+
+    >>> # Train the model
+    >>> engine = Engine()
+    >>> engine.train(model=model, datamodule=datamodule)
+
+    >>> # Get predictions
+    >>> engine.predict(model=model, datamodule=datamodule)
+
+See Also:
+    - :class:`UniNet`: Main model class for UniNet-based anomaly detection
+    - :class:`UniNetModel`: PyTorch implementation of the UniNet model
+"""
+
+# Copyright (C) 2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .lightning_model import UniNet
+from .torch_model import UniNetModel
+
+__all__ = ["UniNet", "UniNetModel"]

src/anomalib/models/image/uninet/anomaly_map.py

@@ -0,0 +1,74 @@
+"""Utilities for computing anomaly maps."""
+
+# Original Code
+# Copyright (c) 2025 Shun Wei
+# https://github.com/pangdatangtt/UniNet
+# SPDX-License-Identifier: MIT
+#
+# Modified
+# Copyright (C) 2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import einops
+import torch
+from kornia.filters import gaussian_blur2d
+from torch.nn import functional as F  # noqa: N812
+
+
+def weighted_decision_mechanism(
+    batch_size: int,
+    output_list: list[list[torch.Tensor]],
+    alpha: float,
+    beta: float,
+    output_size: tuple[int, int] = (256, 256),
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute anomaly maps using weighted decision mechanism.
+
+    Args:
+        batch_size (int): Batch size.
+        output_list (list[list[torch.Tensor]]): List of output tensors.
+        alpha (float): Alpha parameter. Used for controlling the upper limit
+        beta (float): Beta parameter. Used for controlling the lower limit
+        output_size (tuple[int, int]): Output size.
+
+    Returns:
+        tuple[torch.Tensor, torch.Tensor]: Anomaly score and anomaly map.
+    """
+    total_weights_list = []
+    for i in range(batch_size):
+        low_similarity_list = [torch.max(output_list[j][i]) for j in range(len(output_list))]
+        probs = F.softmax(torch.tensor(low_similarity_list), dim=0)
+        weight_list = []  # set P consists of L high probability values, where L ranges from n-1 to n+1
+        for idx, prob in enumerate(probs):
+            weight_list.append(low_similarity_list[idx]) if prob > torch.mean(probs) else None
+        weight = torch.max(torch.tensor([torch.mean(torch.tensor(weight_list)) * alpha, beta]))
+        total_weights_list.append(weight)
+
+    assert len(total_weights_list) == batch_size, "The number of weights do not match the number of samples"
+
+    anomaly_maps: list[list[torch.Tensor]] = [[] for _ in output_list]
+    for idx, output in enumerate(output_list):
+        anomaly_map_ = torch.unsqueeze(output, dim=1)  # Bx1xhxw
+        # Bx256x256
+        anomaly_maps[idx] = F.interpolate(anomaly_map_, output_size, mode="bilinear", align_corners=True)[:, 0, :, :]
+
+    processed_anomaly_maps: torch.Tensor = sum(anomaly_maps)
+
+    anomaly_score = []
+    for idx in range(batch_size):
+        top_k = int(output_size[0] * output_size[1] * total_weights_list[idx])
+        assert top_k >= 1 / (output_size[0] * output_size[1]), "weight can not be smaller than 1 / (H * W)!"
+
+        single_anomaly_score_exp = processed_anomaly_maps[idx]
+        single_anomaly_score_exp = gaussian_blur2d(
+            einops.rearrange(single_anomaly_score_exp, "h w -> 1 1 h w"),  # kornia expects 4D tensor
+            kernel_size=(5, 5),
+            sigma=(4, 4),
+        ).squeeze()
+        assert single_anomaly_score_exp.reshape(1, -1).shape[-1] == output_size[0] * output_size[1], (
+            "something wrong with the last dimension of reshaped map!"
+        )
+        single_map = single_anomaly_score_exp.reshape(1, -1)
+        single_anomaly_score = single_map.topk(top_k).values[0][0]
+        anomaly_score.append(single_anomaly_score.detach())
+    return torch.vstack(anomaly_score), processed_anomaly_maps.detach()

src/anomalib/models/image/uninet/components/__init__.py

@@ -0,0 +1,13 @@
+"""PyTorch modules for the UniNet model implementation."""
+
+# Copyright (C) 2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .attention_bottleneck import AttentionBottleneck, BottleneckLayer
+from .dfs import DomainRelatedFeatureSelection
+
+__all__ = [
+    "DomainRelatedFeatureSelection",
+    "AttentionBottleneck",
+    "BottleneckLayer",
+]