🚀 feat(model): add UniNet (#2797)

* stash changes

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

* Add initial code

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

* Working model

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

* Working model

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

* Fix pre-commit

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

* Minor refactor

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

* Refactor components

Signed-off-by: Ashwin Vaidya <ashwin.vaidya@intel.com>

* single line

Signed-off-by: Ashwin Vaidya <ashwin.vaidya@intel.com>

* de_resnet->resnet_decoder

Signed-off-by: Ashwin Vaidya <ashwin.vaidya@intel.com>

* bug fixes

Signed-off-by: Ashwin Vaidya <ashwin.vaidya@intel.com>

* openvino export + pr comments

Signed-off-by: Ashwin Vaidya <ashwin.vaidya@intel.com>

* Add documentation

Signed-off-by: Ashwin Vaidya <ashwin.vaidya@intel.com>

* Remove else branch

Signed-off-by: Ashwin Vaidya <ashwin.vaidya@intel.com>

* Move copyright

Signed-off-by: Ashwin Vaidya <ashwin.vaidya@intel.com>

* PR comments

- Add comments in AttentionBlock
- Refacor torch module

Signed-off-by: Ashwin Vaidya <ashwin.vaidya@intel.com>

---------

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>
Signed-off-by: Ashwin Vaidya <ashwin.vaidya@intel.com>

Author: ashwinvaidya17

docs/source/markdown/guides/reference/models/image/uninet.md

@@ -0,0 +1,37 @@
+# UniNet
+
+```{eval-rst}
+.. automodule:: anomalib.models.image.uninet.lightning_model
+   :members:
+   :show-inheritance:
+```
+
+```{eval-rst}
+.. automodule:: anomalib.models.image.uninet.torch_model
+   :members:
+   :show-inheritance:
+```
+
+```{eval-rst}
+.. automodule:: anomalib.models.image.uninet.components.loss
+   :members:
+   :show-inheritance:
+```
+
+```{eval-rst}
+.. automodule:: anomalib.models.image.uninet.components.anomaly_map
+   :members:
+   :show-inheritance:
+```
+
+```{eval-rst}
+.. automodule:: anomalib.models.image.uninet.components.attention_bottleneck
+   :members:
+   :show-inheritance:
+```
+
+```{eval-rst}
+.. automodule:: anomalib.models.image.uninet.components.dfs
+   :members:
+   :show-inheritance:
+```

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

@@ -73,6 +73,7 @@
     Stfpm,
     Supersimplenet,
     Uflow,
+    UniNet,
     VlmAd,
     WinClip,
 )
@@ -110,6 +111,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 .resnet_decoder import get_decoder
+
+__all__ = ["get_decoder"]

src/anomalib/models/image/reverse_distillation/components/de_resnet.py renamed to src/anomalib/models/components/backbone/resnet_decoder.py

@@ -7,9 +7,9 @@
 # Copyright (C) 2022-2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-"""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.
 
@@ -19,17 +19,15 @@
 - Full decoder network architecture
 
 Example:
-    >>> from anomalib.models.image.reverse_distillation.components.de_resnet import (
+    >>> from anomalib.models.components.backbone.resnet_decoder 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.resnet_decoder.DecoderBasicBlock`:
         Basic building block for the decoder network
 """
 
@@ -157,8 +155,8 @@ class DecoderBottleneck(nn.Module):
     """Bottleneck block for the decoder network.
 
     This module implements a bottleneck block used in the decoder part of the Reverse
-    Distillation model. It performs upsampling and feature reconstruction through a series of
-    convolutional layers.
+    Distillation and UniNet models. It performs upsampling and feature reconstruction
+    through a series of convolutional layers.
 
     The block consists of three convolution layers:
     1. 1x1 conv to adjust channels
@@ -186,7 +184,7 @@ class DecoderBottleneck(nn.Module):
 
     Example:
         >>> import torch
-        >>> from anomalib.models.image.reverse_distillation.components.de_resnet import (
+        >>> from anomalib.models.components.backbone.resnet_decoder import (
         ...     DecoderBottleneck
         ... )
         >>> layer = DecoderBottleneck(256, 64)
@@ -269,7 +267,7 @@ def forward(self, batch: torch.Tensor) -> torch.Tensor:
         return self.relu(out)
 
 
-class ResNet(nn.Module):
+class ResNetDecoder(nn.Module):
     """Decoder ResNet model for feature reconstruction.
 
     This module implements a decoder version of the ResNet architecture, which
@@ -297,11 +295,11 @@ class ResNet(nn.Module):
             layer to use. If ``None``, uses ``BatchNorm2d``. Defaults to ``None``.
 
     Example:
-        >>> from anomalib.models.image.reverse_distillation.components import (
+        >>> from anomalib.models.components.backbone.resnet_decoder import (
         ...     DecoderBasicBlock,
-        ...     ResNet
+        ...     ResNetDecoder
         ... )
-        >>> model = ResNet(
+        >>> model = ResNetDecoder(
         ...     block=DecoderBasicBlock,
         ...     layers=[2, 2, 2, 2]
         ... )
@@ -437,63 +435,63 @@ def forward(self, batch: torch.Tensor) -> list[torch.Tensor]:
         return [feature_c, feature_b, feature_a]
 
 
-def _resnet(block: type[DecoderBasicBlock | DecoderBottleneck], layers: list[int], **kwargs) -> ResNet:
-    return ResNet(block, layers, **kwargs)
+def _resnet(block: type[DecoderBasicBlock | DecoderBottleneck], layers: list[int], **kwargs) -> ResNetDecoder:
+    return ResNetDecoder(block, layers, **kwargs)
 
 
-def de_resnet18() -> ResNet:
+def de_resnet18() -> ResNetDecoder:
     """ResNet-18 model."""
     return _resnet(DecoderBasicBlock, [2, 2, 2, 2])
 
 
-def de_resnet34() -> ResNet:
+def de_resnet34() -> ResNetDecoder:
     """ResNet-34 model."""
     return _resnet(DecoderBasicBlock, [3, 4, 6, 3])
 
 
-def de_resnet50() -> ResNet:
+def de_resnet50() -> ResNetDecoder:
     """ResNet-50 model."""
     return _resnet(DecoderBottleneck, [3, 4, 6, 3])
 
 
-def de_resnet101() -> ResNet:
+def de_resnet101() -> ResNetDecoder:
     """ResNet-101 model."""
     return _resnet(DecoderBottleneck, [3, 4, 23, 3])
 
 
-def de_resnet152() -> ResNet:
+def de_resnet152() -> ResNetDecoder:
     """ResNet-152 model."""
     return _resnet(DecoderBottleneck, [3, 8, 36, 3])
 
 
-def de_resnext50_32x4d() -> ResNet:
+def de_resnext50_32x4d() -> ResNetDecoder:
     """ResNeXt-50 32x4d model."""
     return _resnet(DecoderBottleneck, [3, 4, 6, 3], groups=32, width_per_group=4)
 
 
-def de_resnext101_32x8d() -> ResNet:
+def de_resnext101_32x8d() -> ResNetDecoder:
     """ResNeXt-101 32x8d model."""
     return _resnet(DecoderBottleneck, [3, 4, 23, 3], groups=32, width_per_group=8)
 
 
-def de_wide_resnet50_2() -> ResNet:
+def de_wide_resnet50_2() -> ResNetDecoder:
     """Wide ResNet-50-2 model."""
     return _resnet(DecoderBottleneck, [3, 4, 6, 3], width_per_group=128)
 
 
-def de_wide_resnet101_2() -> ResNet:
+def de_wide_resnet101_2() -> ResNetDecoder:
     """Wide ResNet-101-2 model."""
     return _resnet(DecoderBottleneck, [3, 4, 23, 3], width_per_group=128)
 
 
-def get_decoder(name: str) -> ResNet:
+def get_decoder(name: str) -> ResNetDecoder:
     """Get decoder model based on the name of the backbone.
 
     Args:
         name (str): Name of the backbone.
 
     Returns:
-        ResNet: Decoder ResNet architecture.
+        ResNetDecoder: Decoder ResNet architecture.
     """
     decoder_map = {
         "resnet18": de_resnet18,

src/anomalib/models/image/__init__.py

@@ -62,6 +62,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
 
@@ -83,6 +84,7 @@
     "Stfpm",
     "Supersimplenet",
     "Uflow",
+    "UniNet",
     "VlmAd",
     "WinClip",
     "Dinomaly",

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

@@ -10,7 +10,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 (
@@ -23,11 +22,8 @@
 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
 """
 
 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/README.md

@@ -0,0 +1,65 @@
+# UniNet: Unified Contrastive Learning Framework for Anomaly Detection
+
+This is the implementation of the UniNet model, a unified contrastive learning framework for anomaly detection presented in CVPR 2025.
+
+Model Type: Classification and Segmentation
+
+## Description
+
+UniNet is a contrastive learning-based anomaly detection model that uses a teacher-student architecture with attention bottleneck mechanisms. The model is designed for diverse domains and supports both supervised and unsupervised anomaly detection scenarios. It focuses on multi-class anomaly detection and leverages domain-related feature selection to improve performance across different categories.
+
+The model consists of:
+
+- **Teacher Networks**: Pre-trained backbone networks that provide reference features for normal samples
+- **Student Network**: A decoder network that learns to reconstruct normal patterns
+- **Attention Bottleneck**: Mechanisms that help focus on relevant features
+- **Domain-Related Feature Selection**: Adaptive feature selection for different domains
+- **Contrastive Loss**: Temperature-controlled similarity computation between student and teacher features
+
+During training, the student network learns to match teacher features for normal samples while being trained to distinguish anomalous patterns through contrastive learning. The model uses a weighted decision mechanism during inference to combine multi-scale features for final anomaly scoring.
+
+## Architecture
+
+The UniNet architecture leverages contrastive learning with teacher-student networks, incorporating attention bottlenecks and domain-specific feature selection for robust anomaly detection across diverse domains.
+
+## Usage
+
+`anomalib train --model UniNet --data MVTecAD --data.category <category>`
+
+## Benchmark
+
+All results gathered with seed `42`.
+
+## [MVTecAD Dataset](https://www.mvtec.com/company/research/datasets/mvtec-ad)
+
+### Image-Level AUC
+
+|  Avg  | Bottle | Cable | Capsule | Carpet | Grid  | Hazelnut | Leather | Metal Nut | Pill  | Screw | Tile  | Toothbrush | Transistor | Wood  | Zipper |
+| :---: | :----: | :---: | :-----: | :----: | :---: | :------: | :-----: | :-------: | :---: | :---: | :---: | :--------: | :--------: | :---: | :----: |
+| 0.956 | 0.999  | 0.982 |  0.939  | 0.896  | 0.996 |  0.999   |  1.000  |   1.000   | 0.816 | 0.919 | 0.970 |   1.000    |   0.984    | 0.993 | 0.945  |
+
+### Pixel-Level AUC
+
+|  Avg  | Bottle | Cable | Capsule | Carpet | Grid  | Hazelnut | Leather | Metal Nut | Pill  | Screw | Tile  | Toothbrush | Transistor | Wood  | Zipper |
+| :---: | :----: | :---: | :-----: | :----: | :---: | :------: | :-----: | :-------: | :---: | :---: | :---: | :--------: | :--------: | :---: | :----: |
+| 0.976 | 0.989  | 0.983 |  0.985  | 0.973  | 0.992 |  0.987   |  0.993  |   0.984   | 0.964 | 0.992 | 0.965 |   0.992    |   0.923    | 0.961 | 0.984  |
+
+### Image F1 Score
+
+|  Avg  | Bottle | Cable | Capsule | Carpet | Grid  | Hazelnut | Leather | Metal Nut | Pill  | Screw | Tile  | Toothbrush | Transistor | Wood  | Zipper |
+| :---: | :----: | :---: | :-----: | :----: | :---: | :------: | :-----: | :-------: | :---: | :---: | :---: | :--------: | :--------: | :---: | :----: |
+| 0.957 | 0.984  | 0.944 |  0.964  | 0.883  | 0.973 |  0.986   |  0.995  |   0.989   | 0.921 | 0.905 | 0.946 |   0.983    |   0.961    | 0.974 | 0.959  |
+
+## Model Features
+
+- **Contrastive Learning**: Uses temperature-controlled contrastive loss for effective anomaly detection
+- **Multi-Domain Support**: Designed to work across diverse domains with domain-related feature selection
+- **Flexible Training**: Supports both supervised and unsupervised training modes
+- **Attention Mechanisms**: Incorporates attention bottlenecks for focused feature learning
+- **Multi-Scale Features**: Leverages multi-scale feature matching for robust detection
+
+## Parameters
+
+- `student_backbone`: Backbone model for student network (default: "wide_resnet50_2")
+- `teacher_backbone`: Backbone model for teacher network (default: "wide_resnet50_2")
+- `temperature`: Temperature parameter for contrastive loss (default: 0.1)