📋 [TASK] Implement UniNet - CVPR 2025

[samet-akcay]

UniNet - CVPR 2025

Implementation Overview

Integration task for UniNet, a CVPR 2025 unified anomaly detection framework, into anomalib. This adds multi-domain capabilities with contrastive learning and feature selection.

Paper Details

• Title: UniNet: A Contrastive Learning-guided Unified Framework with Feature Selection for Anomaly Detection
• Venue: CVPR 2025
• Repository: https://github.com/pangdatangtt/UniNet
• Authors: Wei, Shun and Jiang, Jielin and Xu, Xiaolong

Key Features & Advantages

1. Unified Multi-Domain Framework

• Industrial Domain: MVTec AD, VisA, BTAD, MVTec 3D-AD, VAD
• Medical Domain: OCT2017, APTOS, ISIC2018, Kvasir, CVC-ColonDB, CVC-ClinicDB
• Video Domain: Ped2

2. Flexible Learning Paradigms

• Unsupervised Anomaly Detection: Traditional one-class learning
• Supervised Anomaly Detection: Including medical polyp segmentation
• Multi-class Anomaly Detection: Handling multiple anomaly types

3. Technical Innovations

• Contrastive learning-guided approach
• Feature selection mechanism
• Unified architecture across domains
• Strong performance on standard benchmarks

Implementation Requirements

Implementation Tasks

1. Core Model Architecture

• UniNet Base Module (src/anomalib/models/image/uninet/)
  • Implement UniNetModel class inheriting from AnomalyModule
  • Add contrastive learning components with temperature scaling
  • Implement feature selection mechanism with attention weights
  • Support for ResNet18/34/50 backbones (following paper specs)

2. Domain-Specific Components

• Industrial Branch

  • MVTec AD, VisA, BTAD compatibility
  • One-class and multi-class detection modes
  • Patch-based feature extraction

• Medical Branch

  • Medical image preprocessing pipeline
  • Segmentation support for polyp detection
  • Binary classification for OCT/APTOS/ISIC datasets

• Video Branch

  • Temporal feature aggregation for Ped2
  • Frame-level anomaly scoring
  • Motion pattern analysis

3. Loss Functions & Training

• Contrastive Loss Implementation

  class ContrastiveLoss(nn.Module):
      def __init__(self, temperature=0.07):
          self.temperature = temperature
      
      def forward(self, features, labels):
          # Implement InfoNCE-style contrastive loss

• Feature Selection Loss

  • Attention-based feature weighting
  • Sparsity regularization
  • Domain adaptation components

4. Configuration & Hyperparameters

• Domain-specific configs

  configs/model/uninet/industrial.yaml
  configs/model/uninet/medical.yaml  
  configs/model/uninet/video.yaml
  

• Mode-specific parameters

  • Unsupervised: contrastive_weight: 0.1
  • Supervised: classification_weight: 1.0
  • Multi-class: num_classes: auto

5. Integration with Anomalib Framework

• Lightning Module Integration

  # src/anomalib/models/image/uninet/lightning_model.py
  class UniNet(AnomalyModule):
      def __init__(
          self,
          backbone: str = "resnet18",
          ...
      ):

• Metrics & Visualization

  • AUROC/AUPR for all domains
  • Pixel-level metrics for segmentation tasks
  • Attention visualization for feature selection
  • Domain-specific visualizations

• Callbacks & Logging

  • Early stopping based on validation AUROC
  • Learning rate scheduling
  • Feature attention weight logging
  • Domain-specific metric tracking

6. Testing & Validation

• Unit Tests

  • Model architecture tests
  • Loss function validation
  • Feature selection mechanism tests
  • Multi-domain compatibility tests

• Integration Tests

  • End-to-end training on MVTec AD
  • Multi-class validation on VisA
  • Medical domain testing on APTOS
  • Video anomaly detection on Ped2

• Benchmark Validation

  • Reproduce paper results (AUROC targets)
  • Performance comparison with existing models
  • Memory usage and training speed analysis

7. Code Structure

src/anomalib/models/image/uninet/
├── __init__.py
├── lightning_model.py      # Main UniNet Lightning module
├── torch_model.py          # Core PyTorch implementation
├── loss.py                 # Contrastive + feature selection losses
├── anomaly_map.py          # Anomaly score computation
└── utils.py                # Helper functions

tests/unit/models/image/uninet/
├── test_lightning_model.py
├── test_torch_model.py
└── test_loss.py

configs/model/
├── uninet.yaml             # Default config
├── uninet_industrial.yaml
├── uninet_medical.yaml
└── uninet_video.yaml

Technical Implementation Details

Model Architecture

class UniNetTorchModel(nn.Module):
    def __init__(self, backbone, domain, feature_dim=512):
        super().__init__()
        self.backbone = get_backbone(backbone)  # ResNet18/34/50
        ...
        
        # Contrastive learning head
        self.contrastive_head = nn.Sequential(
            nn.Linear(self.backbone.feature_dim, feature_dim),
            nn.ReLU(),
            nn.Linear(feature_dim, feature_dim)
        )
        
        # Feature selection attention
        self.attention = nn.MultiheadAttention(feature_dim, num_heads=8)
        
        ...

Training Strategy

1. Phase 1: Contrastive pre-training on normal samples
2. Phase 2: Feature selection optimization with attention
3. Phase 3: Domain-specific fine-tuning (supervised modes)

Key Hyperparameters (from paper)

• Learning rate: 1e-4 (Adam optimizer)
• Batch size: 32 for industrial, 16 for medical
• Training epochs: 100 for unsupervised, 50 for supervised
• Temperature: 0.07 for contrastive loss
• Feature dimension: 512

Performance Targets & Validation

Expected Results (from CVPR 2025 paper)

• MVTec AD: AUROC ~98.5% (image-level), ~97.8% (pixel-level)
• VisA: AUROC ~95.2% (multi-class setting)
• APTOS: AUROC ~92.1% (medical domain)
• Ped2: AUROC ~96.3% (video domain)

Integration Checklist

• Model passes anomalib's standard tests
• Configs follow anomalib naming conventions
• CLI integration works (anomalib train/test/predict)
• Benchmarking results match paper
• Memory usage within acceptable limits
• Training time competitive with existing models