🤖 UniNet: A Contrastive Learning-guided Unified Framework with Feature Selection for Anomaly Detection

July 9, 2025 · View on GitHub

🤖 UniNet: A Contrastive Learning-guided Unified Framework with Feature Selection for Anomaly Detection

Shun Wei,  Jielin Jiang,  Xiaolong Xu

-- CVPR 2025 --

Last commit  Stars  Forks  Visits  PRs  📄 PDF

PWC PWC PWC PWC PWC PWC PWC PWC PWC

🔔 News

  • 07-09-2025: This project and code will be further improved after I complete the exam and paper writing, probably in early August. The citation format has been updated! If you find this method useful and want to compare it with yours, feel free to cite!
  • 05-15-2025: Fix some minor errors.
  • 04-27-2025: The checkpoints have been uploaded.
  • 04-26-2025: The code has been released 😊!
  • 04-21-2025: The code will be released in recent days!
  • 02-27-2025: Accepted by CVPR2025 💥!

🔗 Citation

Please cite our paper if the method and code help and inspire your project:

@InProceedings{Wei_2025_CVPR,
    author    = {Wei, Shun and Jiang, Jielin and Xu, Xiaolong},
    title     = {UniNet: A Contrastive Learning-guided Unified Framework with Feature Selection for Anomaly Detection},
    booktitle = {Proceedings of the Computer Vision and Pattern Recognition Conference (CVPR)},
    month     = {June},
    year      = {2025},
    pages     = {9994-10003}
}

List of TODOs

  • 📖 Introduction
  • 🔧 Environments
  • 📊 Data Preparation
  • 🚀 Run Experiments
  • 📂 Results
  • 🙏 Acknowledgements
  • 📜 License
  • 💬 Contact

📖 Introduction

This repository will contain source code for UniNet implemented with PyTorch.

UniNet is a unified framework designed for diverse domains, such as industrial, medical, and video, by addressing the limitations in general ability of existing methods. Besides, UniNet can be suited for both unsupervised and supervised settings simultaneously. As a unified framework, UniNet also focuses on multi-class anomaly detection.

🔧 Environments

Create a new conda environment and install required packages.

conda create -n uninet_env python=3.9.7
conda activate uninet_env
pip install -r requirements.txt

Experiments are conducted on NVIDIA GeForce RTX 3090. Same GPU and package version are recommended.

📊 Data Preparation

The public datasets employed in the paper are listed below.

1. MVTec AD 
|-- mvtec
    |-- bottle
        |-- ground_truth
        |-- test
        |-- train
    |-- cable
        |-- ground_truth
        |-- test
        |-- train
    |-- ...
2. VisA

Unzip the file to ../VisA/. Preprocess the dataset to ../VisA_pytorch/ in 1-class mode by their official splitting code. Alternately, you can also preprocess the dataset using this code from ReContrast.

|-- VisA_pytorch
    |-- 1cls
        |-- candle
            |-- ground_truth
            |-- test
                    |-- good
                    |-- bad
            |-- train
                    |-- good
        |-- capsules
        |-- ....
3. APTOS

Creat a new directory ../APTOS. Unzip the file to ../APTOS/original/. Now, the directory would be like:

|-- APTOS
    |-- original
        |-- test_images
        |-- train_images
        |-- test.csv
        |-- train.csv

Run the following command to preprocess the data to ../APTOS/.

python ./prepare_data/prepare_aptos.py --data-folder ../APTOS/original --save-folder ../APTOS

Then, it woule be like:

|-- APTOS
    |-- test
        |-- NORMAL
        |-- ABNORMAL
    |-- train
        |-- NORMAL
    |-- original
4. OCT2017

Creat a new directory ../OCT2017. Unzip the file, and move everything in ZhangLabData/CellData/OCT to ../OCT2017/. The directory should be like:

|-- OCT2017
    |-- test
        |-- CNV
        |-- DME
        |-- DRUSEN
        |-- NORMAL
    |-- train
        |-- CNV
        |-- DME
        |-- DRUSEN
        |-- NORMAL
5. ISIC2018

Creat a new directory ../ISIC2018. After downloading "Training Data","Training Ground Truth", "Validation Data", and "Validation Ground Truth" of Task 3, please unzip them to ../ISIC2018/original/. Now, the directory would be like:

|-- ISIC2018
    |-- original
        |-- ISIC2018_Task3_Training_GroundTruth
        |-- ISIC2018_Task3_Training_Input
        |-- ISIC2018_Task3_Validation_GroundTruth
        |-- ISIC2018_Task3_Validation_Input

Run the following command to preprocess the data to ../ISIC2018/.

python ./prepare_data/prepare_isic2018.py --data-folder ../ISIC2018/original --save-folder ../ISIC2018

Then, it would be like:

|-- ISIC2018
    |-- test
        |-- NORMAL
        |-- ABNORMAL
    |-- train
        |-- NORMAL
    |-- original

🚀 Run Experiments

1. Unsupervised AD

Run the following command for industrial domain, such as MVTec AD dataset:

python main.py --setting oc --dataset "MVTec AD"

Run the following command for medical domain, such as APTOS dataset:

python main.py --dataset APTOS

Alternatively, you can train and test all medical datasets at once using the following command:

python main.py --dataset APTOS --train_and_test_all

Run the following command for video domain:

python main.py --dataset Ped2

Run the following command for multiclass anomaly detection, such as VisA dataset:

python main.py --setting mc --dataset VisA
2. Supervised AD

Run the following command for medical polyp segmentation, sucha as Kvasir-SEG dataset:

python main.py --dataset "Kvasir-SEG"

Similarly, you can also train and test all medical datasets at once using the following command:

python main.py --dataset "Kvasir-SEG" --train_and_test_all

Run the following command for industrial domain:

python main.py --dataset VAD
3. Testing

With training on your own or saving the weight file we upload, you can test the model using the following command:

python main.py --setting oc --dataset "MVTec AD" --load_ckpts

python main.py --setting mc --dataset VisA --load_ckpts

python main.py --dataset APTOS --load_ckpts

For unsupervised industrial AD, the ''--setting'' must be considered.

📂 Results

1. Unsupervised anomaly detection

2. Supervised anomaly detection

3. Multi-class setting

4. Visualization results

🙏 Acknowledgements

We sincerely appreciate ReContrast for its clear and excellent approach.

📜 License

The code in this repository is licensed under the MIT license.

💬 Contact

For the questions or collaborations:

⭐Please consider giving me a star if you find it useful!!!