🚀 feat(model): add UniNet

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/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,32 @@
+"""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,77 @@
+"""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 numpy as np
+import torch
+from scipy.ndimage import gaussian_filter
+from torch.nn import functional
+
+
+def weighted_decision_mechanism(
+    batch_size: int,
+    output_list: list[list[torch.Tensor]],
+    alpha: float,
+    beta: float,
+    out_size: int = 25,
+) -> 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
+        out_size (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 = functional.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].cpu().numpy()) if prob > torch.mean(probs) else None
+        # TODO(ashwinvaidya17): see if all the numpy operations can be replaced with torch operations
+        weight = np.max([np.mean(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_map_lists: list[list[torch.Tensor]] = [[] for _ in output_list]
+    for idx, output in enumerate(output_list):
+        cat_output = torch.cat(output, dim=0)
+        anomaly_map = torch.unsqueeze(cat_output, dim=1)  # Bx1xhxw
+        # Bx256x256
+        anomaly_map_lists[idx] = functional.interpolate(anomaly_map, out_size, mode="bilinear", align_corners=True)[
+            :,
+            0,
+            :,
+            :,
+        ]
+
+    anomaly_map = sum(anomaly_map_lists)
+
+    anomaly_score = []
+    for idx in range(batch_size):
+        top_k = int(out_size * out_size * total_weights_list[idx])
+        assert top_k >= 1 / (out_size * out_size), "weight can not be smaller than 1 / (H * W)!"
+
+        single_anomaly_score_exp = anomaly_map[idx]
+        single_anomaly_score_exp = torch.tensor(gaussian_filter(single_anomaly_score_exp.cpu().numpy(), sigma=4))
+        assert (
+            single_anomaly_score_exp.reshape(1, -1).shape[-1] == out_size * out_size
+        ), "something wrong with the last dimension of reshaped map!"
+        single_map = single_anomaly_score_exp.reshape(1, -1)
+        single_anomaly_score = np.sort(single_map.topk(top_k, dim=-1)[0].detach().cpu().numpy(), axis=1)
+        anomaly_score.append(single_anomaly_score)
+    return anomaly_score, anomaly_map.detach().cpu().numpy()

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

@@ -0,0 +1,14 @@
+"""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, domain_related_feature_selection
+
+__all__ = [
+    "domain_related_feature_selection",
+    "DomainRelatedFeatureSelection",
+    "AttentionBottleneck",
+    "BottleneckLayer",
+]